diff --git a/cmd/picoclaw/internal/config/command.go b/cmd/picoclaw/internal/config/command.go new file mode 100644 index 00000000..3329d491 --- /dev/null +++ b/cmd/picoclaw/internal/config/command.go @@ -0,0 +1,57 @@ +package config + +import ( + "fmt" + "strings" + + "github.com/spf13/cobra" + + "github.com/sipeed/picoclaw/cmd/picoclaw/internal" + "github.com/sipeed/picoclaw/pkg/config" +) + +func NewConfigCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "config", + Short: "Manage configuration", + } + + cmd.AddCommand(newResetCommand()) + return cmd +} + +func newResetCommand() *cobra.Command { + var force bool + + cmd := &cobra.Command{ + Use: "reset", + Short: "Reset configuration to factory defaults", + Args: cobra.NoArgs, + Example: ` picoclaw config reset + picoclaw config reset --force`, + RunE: func(_ *cobra.Command, _ []string) error { + if !force { + fmt.Print("Reset config to factory defaults? API keys will be preserved. (y/n): ") + var response string + fmt.Scanln(&response) + if strings.ToLower(strings.TrimSpace(response)) != "y" { + fmt.Println("Aborted.") + return nil + } + } + + configPath := internal.GetConfigPath() + if err := config.ResetToDefaults(configPath); err != nil { + return fmt.Errorf("reset failed: %w", err) + } + fmt.Println("Configuration has been reset to factory defaults.") + fmt.Println("A backup of the previous config was created in the same directory.") + return nil + }, + } + + cmd.Flags().BoolVarP(&force, "force", "f", false, + "Skip confirmation prompt") + + return cmd +} diff --git a/cmd/picoclaw/main.go b/cmd/picoclaw/main.go index abcf03a3..1b5d0974 100644 --- a/cmd/picoclaw/main.go +++ b/cmd/picoclaw/main.go @@ -17,6 +17,7 @@ import ( "github.com/sipeed/picoclaw/cmd/picoclaw/internal/agent" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/auth" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/cliui" + configcmd "github.com/sipeed/picoclaw/cmd/picoclaw/internal/config" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/cron" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/gateway" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/mcp" @@ -82,6 +83,7 @@ picoclaw --no-color status`, }) cmd.AddCommand( + configcmd.NewConfigCommand(), onboard.NewOnboardCommand(), agent.NewAgentCommand(), auth.NewAuthCommand(), diff --git a/cmd/picoclaw/main_test.go b/cmd/picoclaw/main_test.go index 037c7c2e..e4f9de5e 100644 --- a/cmd/picoclaw/main_test.go +++ b/cmd/picoclaw/main_test.go @@ -39,6 +39,7 @@ func TestNewPicoclawCommand(t *testing.T) { allowedCommands := []string{ "agent", "auth", + "config", "cron", "gateway", "mcp", diff --git a/docs/guides/configuration.fr.md b/docs/guides/configuration.fr.md index 786a0c28..dcacc187 100644 --- a/docs/guides/configuration.fr.md +++ b/docs/guides/configuration.fr.md @@ -66,6 +66,50 @@ PicoClaw stocke les données dans votre workspace configuré (par défaut : `~/. > **Remarque :** Les modifications apportées à `AGENT.md`, `SOUL.md`, `USER.md` et `memory/MEMORY.md` sont détectées automatiquement au moment de l'exécution via le suivi de la date de modification (mtime). Il n'est **pas nécessaire de redémarrer le gateway** après avoir modifié ces fichiers — l'agent charge le nouveau contenu à la prochaine requête. +### Politique de contexte de requête + +`turn_profile` est une politique facultative sous `agents.defaults.turn_profile` pour contrôler le contexte chargé par chaque nouveau tour : historique, prompt système, prompts de skills et outils autorisés. Sans cette configuration, ou avec `"enabled": false`, PicoClaw garde son comportement normal. Avec `"enabled": true`, la politique ci-dessous s'applique à chaque nouveau tour. + +Tous les blocs utilisent les mêmes valeurs de `mode` : + +| Mode | Signification | +| --- | --- | +| `default` | Garde le comportement normal de PicoClaw. Un bloc absent ou sans `mode` vaut `default`. | +| `off` | Désactive ce bloc pour le tour. | +| `custom` | Utilise une liste d'autorisation. Dans cette version, `custom` est pris en charge seulement pour `skills` et `tools`; l'utiliser pour `history` ou `system_prompt` produit une erreur de validation. | + +Blocs disponibles : + +| Bloc | Ce qu'il contrôle | +| --- | --- | +| `history` | Lecture de l'historique et du résumé, écriture des messages utilisateur/assistant/outil, ingestion de contexte, compression et résumé. | +| `system_prompt` | Injection de l'identité PicoClaw, des instructions de l'espace de travail, de la mémoire, du contexte d'exécution et du résumé. Les prompts système externes restent autorisés quand ce bloc est `off`. | +| `skills` | Chargement du catalogue de skills et du contenu des skills actifs. `custom.allow` ne garde que les noms listés. | +| `tools` | Outils exposés au modèle et autorisés à l'exécution. `custom.allow` ne garde que les outils enregistrés et listés. | + +Quand `system_prompt.mode` vaut `off`, que des outils restent visibles et qu'aucun prompt système externe n'est fourni, PicoClaw réutilise sa règle existante d'utilisation des outils comme prompt minimal de secours. Si `tools.mode` vaut `off`, ce prompt de secours n'est pas ajouté. + +Exemple de contexte propre avec outils web : + +```json +{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": { "mode": "off" }, + "system_prompt": { "mode": "off" }, + "skills": { "mode": "off" }, + "tools": { + "mode": "custom", + "allow": ["web_search", "web_fetch"] + } + } + } + } +} +``` + ### Sources de Compétences Par défaut, les compétences sont chargées depuis : @@ -364,6 +408,55 @@ Configurez plusieurs endpoints pour le même nom de modèle — PicoClaw effectu L'ancienne configuration `providers` est **dépréciée** et a été supprimée dans V2. Les configs V0/V1 existantes sont auto-migrées. Voir [docs/migration/model-list-migration.md](../migration/model-list-migration.md). +#### Configuration du Streaming + +Le streaming provider utilise un double opt-in et est désactivé par défaut. L'agent ne tente le streaming que lorsque le channel courant a `settings.streaming.enabled: true`, que l'entrée de modèle active a `streaming.enabled: true`, et que le provider comme le channel prennent en charge le streaming. Si une condition manque, PicoClaw utilise le chemin de requête non-streaming normal. + +Pico WebUI est le premier channel entièrement câblé. Pico crée le premier message assistant avec le message wire existant `message.create`, puis met à jour ce même message avec `message.update`; aucun nouveau type de message Pico n'est introduit. + +Laissez `streaming` absent si vous ne voulez pas de streaming. Un bloc `streaming` omis signifie désactivé; il n'est pas nécessaire d'écrire `"streaming": {"enabled": false}`. + +Exemple d'activation : + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| Champ | Type | Défaut | Description | +| ----- | ---- | ------ | ----------- | +| `channel_list..settings.streaming.enabled` | bool | `false` | Autorise ce channel à afficher la sortie streaming du provider | +| `channel_list..settings.streaming.throttle_seconds` | int | Défaut Pico après activation : `0` | Intervalle minimal entre les mises à jour intermédiaires; le contenu final est toujours envoyé | +| `channel_list..settings.streaming.min_growth_chars` | int | Défaut Pico après activation : `1` | Croissance minimale du texte avant une mise à jour intermédiaire; le contenu final est toujours envoyé | +| `model_list[].streaming.enabled` | bool | `false` | Autorise cette entrée de modèle à tenter des requêtes provider en streaming | + +Les anciennes variables d'environnement Telegram restent compatibles : `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`, `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS` et `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS`. Elles s'appliquent uniquement aux settings Telegram et n'activent ni ne modifient `settings.streaming` de Pico. + +Le comportement d'échec est volontairement conservateur : si le streaming échoue avant l'envoi d'un chunk visible, PicoClaw réessaie une fois via le chemin `Chat()` normal. Si un chunk a déjà été affiché à l'utilisateur, PicoClaw n'envoie pas une deuxième réponse non-streaming, afin d'éviter une sortie dupliquée. + ### Architecture des Providers PicoClaw route les providers par famille de protocole : diff --git a/docs/guides/configuration.it.md b/docs/guides/configuration.it.md index d7de4689..213ca107 100644 --- a/docs/guides/configuration.it.md +++ b/docs/guides/configuration.it.md @@ -31,6 +31,55 @@ PICOCLAW_HOME=/opt/picoclaw picoclaw agent PICOCLAW_HOME=/srv/picoclaw PICOCLAW_CONFIG=/srv/picoclaw/main.json picoclaw gateway ``` +### Configurazione Streaming + +Lo streaming del provider usa un double opt-in ed è disattivato per impostazione predefinita. L'agent prova lo streaming solo quando il canale corrente ha `settings.streaming.enabled: true`, l'entry del modello attivo ha `streaming.enabled: true`, e sia il provider sia il canale supportano lo streaming. Se manca una qualsiasi condizione, PicoClaw usa il normale percorso di richiesta non streaming. + +Pico WebUI è il primo canale completamente collegato. Pico crea il primo messaggio assistant con il wire message esistente `message.create`, poi aggiorna lo stesso messaggio con `message.update`; non viene introdotto alcun nuovo tipo di wire message Pico. + +Lascia `streaming` assente quando non vuoi usare lo streaming. Un blocco `streaming` omesso significa disattivato; non è necessario scrivere `"streaming": {"enabled": false}`. + +Esempio di attivazione: + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| Campo | Tipo | Predefinito | Descrizione | +| ----- | ---- | ----------- | ----------- | +| `channel_list..settings.streaming.enabled` | bool | `false` | Permette a questo canale di mostrare l'output streaming del provider | +| `channel_list..settings.streaming.throttle_seconds` | int | Predefinito Pico dopo l'attivazione: `0` | Intervallo minimo tra aggiornamenti intermedi; il contenuto finale viene sempre inviato | +| `channel_list..settings.streaming.min_growth_chars` | int | Predefinito Pico dopo l'attivazione: `1` | Crescita minima del testo prima di inviare un aggiornamento intermedio; il contenuto finale viene sempre inviato | +| `model_list[].streaming.enabled` | bool | `false` | Permette a questa entry di modello di provare richieste provider streaming | + +Le variabili d'ambiente legacy di Telegram restano compatibili: `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`, `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS` e `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS`. Si applicano solo alle settings Telegram e non attivano né modificano `settings.streaming` di Pico. + +Il comportamento in caso di errore è intenzionalmente conservativo: se lo streaming fallisce prima che venga inviato un chunk visibile, PicoClaw riprova una volta tramite il normale percorso `Chat()`. Se un chunk è già stato mostrato all'utente, PicoClaw non invia una seconda risposta non streaming, evitando output duplicato. + ### Struttura del Workspace PicoClaw salva i dati nel workspace configurato (predefinito: `~/.picoclaw/workspace`): diff --git a/docs/guides/configuration.ja.md b/docs/guides/configuration.ja.md index 0234edbd..7dd67644 100644 --- a/docs/guides/configuration.ja.md +++ b/docs/guides/configuration.ja.md @@ -67,6 +67,50 @@ PicoClaw は設定されたワークスペース（デフォルト: `~/.picoclaw > **注意：** `AGENT.md`、`SOUL.md`、`USER.md` および `memory/MEMORY.md` への変更は、ファイル更新時刻（mtime）の追跡により実行時に自動検出されます。これらのファイルを編集した後に **gateway を再起動する必要はありません** — Agent は次のリクエスト時に最新の内容を自動的に読み込みます。 +### リクエストコンテキストポリシー + +`turn_profile` は `agents.defaults.turn_profile` に置く任意のリクエストコンテキストポリシーです。各ターンに履歴、system prompt、skill prompt、許可ツールを含めるかどうかを制御します。未設定、または `"enabled": false` の場合、PicoClaw は通常動作のままです。`"enabled": true` にすると、このポリシーが各新規ターンに適用されます。 + +各ブロックは同じ `mode` を使います。 + +| Mode | 意味 | +| --- | --- | +| `default` | PicoClaw の通常動作を維持します。ブロックまたは `mode` が省略された場合も `default` です。 | +| `off` | そのブロックを無効にします。 | +| `custom` | 許可リストを使います。このバージョンでは `skills` と `tools` のみ対応し、`history` や `system_prompt` で使うと検証エラーになります。 | + +ブロックの意味： + +| ブロック | 制御する内容 | +| --- | --- | +| `history` | 履歴と要約の読み込み、ユーザー/アシスタント/ツールメッセージの保存、コンテキスト取り込み、圧縮と要約。 | +| `system_prompt` | PicoClaw の既定の identity、ワークスペース指示、メモリ、実行時コンテキスト、要約の注入。`off` でも外部 system prompt は利用できます。 | +| `skills` | Skill カタログと active skill のプロンプト内容。`custom.allow` は列挙した skill 名だけを残します。 | +| `tools` | モデルに見せ、実行を許可するツール。`custom.allow` は登録済みで列挙されたツール名だけを残します。 | + +`system_prompt.mode` が `off` で、ツールが表示され、外部 system prompt がない場合、PicoClaw は既存のツール使用ルールを最小のフォールバックプロンプトとして再利用します。`tools.mode` が `off` の場合、このフォールバックは追加されません。 + +Web ツールだけを残すクリーンなコンテキスト例： + +```json +{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": { "mode": "off" }, + "system_prompt": { "mode": "off" }, + "skills": { "mode": "off" }, + "tools": { + "mode": "custom", + "allow": ["web_search", "web_fetch"] + } + } + } + } +} +``` + ### スキルソース デフォルトでは、スキルは以下の順序で読み込まれます： @@ -365,6 +409,55 @@ HEARTBEAT_OK を返信 ユーザーが直接結果を受信 旧 `providers` 設定は**非推奨**となり、V2 で削除されました。既存の V0/V1 設定は自動的に移行されます。[docs/migration/model-list-migration.md](../migration/model-list-migration.md) を参照してください。 +#### ストリーミング設定 + +Provider ストリーミングは二重の opt-in 方式で、デフォルトでは無効です。現在の channel に `settings.streaming.enabled: true` があり、アクティブなモデルエントリに `streaming.enabled: true` があり、さらに provider と channel の両方がストリーミングをサポートしている場合にのみ、agent はストリーミングリクエストを試行します。いずれかの条件が欠ける場合、PicoClaw は通常の非ストリーミングリクエスト経路を使います。 + +Pico WebUI が最初に完全対応した channel です。Pico は既存の `message.create` wire message で最初の assistant メッセージを作成し、その後 `message.update` で同じメッセージを更新します。新しい Pico wire message type は追加されません。 + +ストリーミングを使わない場合は `streaming` を省略してください。`streaming` ブロックの省略は無効を意味するため、`"streaming": {"enabled": false}` を書く必要はありません。 + +有効化例： + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| フィールド | 型 | デフォルト | 説明 | +| ---------- | -- | ---------- | ---- | +| `channel_list..settings.streaming.enabled` | bool | `false` | この channel で provider のストリーミング出力を表示できるようにします | +| `channel_list..settings.streaming.throttle_seconds` | int | Pico で有効化後のデフォルト：`0` | 中間更新の最小間隔。最終内容は常に flush されます | +| `channel_list..settings.streaming.min_growth_chars` | int | Pico で有効化後のデフォルト：`1` | 次の中間更新を送るために必要な最小文字増加数。最終内容は常に flush されます | +| `model_list[].streaming.enabled` | bool | `false` | このモデルエントリで provider ストリーミングリクエストを試行できるようにします | + +既存の Telegram 環境変数 `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`、`PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS`、`PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS` は互換性のため引き続き使えます。これらは Telegram settings にのみ適用され、Pico の `settings.streaming` を有効化または変更しません。 + +失敗時の動作は保守的です。可視 chunk が送信される前にストリーミングが失敗した場合、PicoClaw は通常の `Chat()` 経路で一度だけ再試行します。すでに chunk がユーザーに表示されている場合は、表示済み出力の重複を避けるため、二つ目の非ストリーミング回答は送信しません。 + ### Provider アーキテクチャ PicoClaw はプロトコルファミリーで Provider をルーティングします： diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 64a61f24..ef4802a7 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -99,6 +99,50 @@ The `evolution` block controls PicoClaw's self-evolution runtime. When enabled, Use `observe` first if you want to inspect learning records without generating skill changes. Use `draft` when you want PicoClaw to prepare reviewable improvements. Use `apply` only when you are comfortable letting accepted drafts update workspace skills. +### Request Context Policy + +`turn_profile` is an optional request context policy under `agents.defaults.turn_profile`. Leave it unset or set `"enabled": false` to keep PicoClaw's normal behavior. When `"enabled": true`, the same policy applies to every new turn. + +Each block uses the same `mode` values: + +| Mode | Meaning | +| --- | --- | +| `default` | Keep PicoClaw's normal behavior for that block. Missing blocks and missing `mode` fields are treated as `default`. | +| `off` | Disable that block for the turn. | +| `custom` | Use an allow list. In this version, `custom` is supported only for `skills` and `tools`; using it for `history` or `system_prompt` is a validation error. | + +Profile blocks: + +| Block | What it controls | +| --- | --- | +| `history` | Whether the turn reads prior session history and summary, writes user/assistant/tool messages, ingests context, and runs compaction or summarization. | +| `system_prompt` | Whether PicoClaw injects its default identity, workspace instructions, memory, runtime context, and summary. External request system prompts are still allowed when this is `off`. | +| `skills` | Whether the skill catalog and active skill prompt content are loaded. `custom.allow` keeps only the listed skill names in prompt context. | +| `tools` | Which callable tools are exposed to the model and allowed at execution time. `custom.allow` keeps only listed registered tool names. | + +When `system_prompt.mode` is `off`, tools are still visible, and no external system prompt is supplied, PicoClaw uses its existing tool-use rule as the minimal fallback prompt. If `tools.mode` is `off`, no fallback prompt is added. + +Example clean web policy: + +```json +{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": { "mode": "off" }, + "system_prompt": { "mode": "off" }, + "skills": { "mode": "off" }, + "tools": { + "mode": "custom", + "allow": ["web_search", "web_fetch"] + } + } + } + } +} +``` + ### Web launcher dashboard **picoclaw-launcher** serves a browser UI that requires password sign-in first. On first run, open `/launcher-setup` to create the dashboard password. Later manual sign-ins use `/launcher-login`. @@ -744,6 +788,55 @@ Resolution rules: - If `provider` is omitted, PicoClaw treats the first `/` segment in `model` as the provider and everything after that first `/` as the runtime model ID. - This means `"model": "openrouter/openai/gpt-5.4"` still works as a compatibility form and sends `openai/gpt-5.4` to OpenRouter. +#### Streaming Configuration + +Provider streaming uses a double opt-in and is disabled by default. The agent only tries streaming when the current channel has `settings.streaming.enabled: true`, the active model entry has `streaming.enabled: true`, and both the provider and channel support streaming. If any condition is missing, PicoClaw uses the normal non-streaming request path. + +Pico WebUI is the first fully wired channel. Pico creates the first assistant message with the existing `message.create` wire message, then updates that same message with `message.update`; no new Pico wire message type is introduced. + +Leave `streaming` unset when you do not want streaming. An omitted `streaming` block means disabled; you do not need to write `"streaming": {"enabled": false}`. + +Opt-in example: + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| Field | Type | Default | Description | +| ----- | ---- | ------- | ----------- | +| `channel_list..settings.streaming.enabled` | bool | `false` | Allows this channel to display provider streaming output | +| `channel_list..settings.streaming.throttle_seconds` | int | Pico default after enabling: `0` | Minimum interval for intermediate updates; final content is always flushed | +| `channel_list..settings.streaming.min_growth_chars` | int | Pico default after enabling: `1` | Minimum character growth before sending an intermediate update; final content is always flushed | +| `model_list[].streaming.enabled` | bool | `false` | Allows this model entry to try provider streaming requests | + +Legacy Telegram environment variables remain compatible: `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`, `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS`, and `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS`. They only apply to Telegram settings and do not enable or modify Pico `settings.streaming`. + +Failure behavior is intentionally conservative: if streaming fails before any visible chunk is sent, PicoClaw retries once through the normal `Chat()` path. If a chunk has already been shown to the user, PicoClaw does not send a second non-streaming answer, because that would duplicate visible output. + #### Vendor-Specific Examples > **Tip**: You can omit `api_key` fields and store them in `.security.yml` for better security. See [Security Configuration](#-security-configuration-recommended). diff --git a/docs/guides/configuration.ms.md b/docs/guides/configuration.ms.md index bcd17afa..4448829a 100644 --- a/docs/guides/configuration.ms.md +++ b/docs/guides/configuration.ms.md @@ -31,6 +31,55 @@ PICOCLAW_HOME=/opt/picoclaw picoclaw agent PICOCLAW_HOME=/srv/picoclaw PICOCLAW_CONFIG=/srv/picoclaw/main.json picoclaw gateway ``` +### Konfigurasi Streaming + +Provider streaming menggunakan double opt-in dan dimatikan secara lalai. Agent hanya mencuba streaming apabila saluran semasa mempunyai `settings.streaming.enabled: true`, entry model aktif mempunyai `streaming.enabled: true`, dan kedua-dua provider serta saluran menyokong streaming. Jika mana-mana syarat tiada, PicoClaw menggunakan laluan permintaan bukan streaming biasa. + +Pico WebUI ialah saluran pertama yang disambungkan sepenuhnya. Pico mencipta mesej assistant pertama dengan wire message sedia ada `message.create`, kemudian mengemas kini mesej yang sama dengan `message.update`; tiada jenis wire message Pico baharu ditambah. + +Biarkan `streaming` tidak ditetapkan jika anda tidak mahu streaming. Blok `streaming` yang tiada bermaksud dimatikan; anda tidak perlu menulis `"streaming": {"enabled": false}`. + +Contoh mengaktifkan streaming: + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| Kunci | Jenis | Lalai | Penerangan | +| ----- | ----- | ----- | ---------- | +| `channel_list..settings.streaming.enabled` | bool | `false` | Membenarkan saluran ini memaparkan output streaming provider | +| `channel_list..settings.streaming.throttle_seconds` | int | Lalai Pico selepas diaktifkan: `0` | Jarak masa minimum antara kemas kini pertengahan; kandungan akhir sentiasa dihantar | +| `channel_list..settings.streaming.min_growth_chars` | int | Lalai Pico selepas diaktifkan: `1` | Pertambahan aksara minimum sebelum menghantar kemas kini pertengahan; kandungan akhir sentiasa dihantar | +| `model_list[].streaming.enabled` | bool | `false` | Membenarkan entry model ini mencuba permintaan provider streaming | + +Pemboleh ubah persekitaran Telegram lama masih serasi: `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`, `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS`, dan `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS`. Ia hanya digunakan untuk settings Telegram dan tidak mengaktifkan atau mengubah `settings.streaming` Pico. + +Tingkah laku kegagalan adalah konservatif: jika streaming gagal sebelum mana-mana chunk kelihatan dihantar, PicoClaw mencuba semula sekali melalui laluan `Chat()` biasa. Jika chunk sudah dipaparkan kepada pengguna, PicoClaw tidak menghantar jawapan bukan streaming kedua untuk mengelakkan output berganda. + ### Susun Atur Workspace PicoClaw menyimpan data dalam workspace yang dikonfigurasikan (lalai: `~/.picoclaw/workspace`): diff --git a/docs/guides/configuration.pt-br.md b/docs/guides/configuration.pt-br.md index e5d904e2..bb5d1413 100644 --- a/docs/guides/configuration.pt-br.md +++ b/docs/guides/configuration.pt-br.md @@ -67,6 +67,50 @@ O PicoClaw armazena dados no seu workspace configurado (padrão: `~/.picoclaw/wo > **Nota:** Alterações em `AGENT.md`, `SOUL.md`, `USER.md` e `memory/MEMORY.md` são detectadas automaticamente em tempo de execução via rastreamento de data de modificação (mtime). **Não é necessário reiniciar o gateway** após editar esses arquivos — o agente carrega o novo conteúdo na próxima requisição. +### Política de contexto da requisição + +`turn_profile` é uma política opcional em `agents.defaults.turn_profile` para controlar qual contexto cada novo turno carrega: histórico, prompt de sistema, prompts de skills e ferramentas permitidas. Sem essa configuração, ou com `"enabled": false`, o PicoClaw mantém o comportamento normal. Com `"enabled": true`, a política abaixo se aplica a cada novo turno. + +Todos os blocos usam os mesmos valores de `mode`: + +| Mode | Significado | +| --- | --- | +| `default` | Mantém o comportamento normal do PicoClaw. Blocos ausentes ou sem `mode` são tratados como `default`. | +| `off` | Desativa esse bloco para o turno. | +| `custom` | Usa uma lista de permissão. Nesta versão, `custom` só é suportado para `skills` e `tools`; usá-lo em `history` ou `system_prompt` gera erro de validação. | + +Blocos disponíveis: + +| Bloco | O que controla | +| --- | --- | +| `history` | Leitura de histórico e resumo, gravação de mensagens de usuário/assistente/ferramenta, ingestão de contexto, compactação e resumo. | +| `system_prompt` | Injeção da identidade padrão do PicoClaw, instruções do workspace, memória, contexto de execução e resumo. Prompts de sistema externos ainda são permitidos quando este bloco está `off`. | +| `skills` | Catálogo de skills e conteúdo de skills ativas no prompt. `custom.allow` mantém apenas os nomes listados. | +| `tools` | Ferramentas visíveis ao modelo e permitidas na execução. `custom.allow` mantém apenas ferramentas registradas e listadas. | + +Quando `system_prompt.mode` é `off`, ferramentas continuam visíveis e nenhum prompt de sistema externo é fornecido, o PicoClaw reutiliza sua regra existente de uso de ferramentas como prompt mínimo de fallback. Se `tools.mode` é `off`, esse fallback não é adicionado. + +Exemplo de contexto limpo com ferramentas web: + +```json +{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": { "mode": "off" }, + "system_prompt": { "mode": "off" }, + "skills": { "mode": "off" }, + "tools": { + "mode": "custom", + "allow": ["web_search", "web_fetch"] + } + } + } + } +} +``` + ### Fontes de Skills Por padrão, as skills são carregadas de: @@ -365,6 +409,55 @@ Configure múltiplos endpoints para o mesmo nome de modelo — PicoClaw fará ro A configuração antiga `providers` está **depreciada** e foi removida no V2. Configs V0/V1 existentes são auto-migradas. Veja [docs/migration/model-list-migration.md](../migration/model-list-migration.md). +#### Configuração de Streaming + +O streaming do provider usa double opt-in e fica desativado por padrão. O agent só tenta streaming quando o canal atual tem `settings.streaming.enabled: true`, a entrada de modelo ativa tem `streaming.enabled: true`, e tanto o provider quanto o canal suportam streaming. Se qualquer condição estiver ausente, o PicoClaw usa o caminho normal de requisição sem streaming. + +O Pico WebUI é o primeiro canal totalmente integrado. O Pico cria a primeira mensagem assistant com o wire message existente `message.create` e depois atualiza a mesma mensagem com `message.update`; nenhum novo tipo de wire message do Pico é introduzido. + +Deixe `streaming` ausente quando não quiser streaming. Um bloco `streaming` omitido significa desativado; você não precisa escrever `"streaming": {"enabled": false}`. + +Exemplo de ativação: + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| Campo | Tipo | Padrão | Descrição | +| ----- | ---- | ------ | --------- | +| `channel_list..settings.streaming.enabled` | bool | `false` | Permite que este canal exiba output streaming do provider | +| `channel_list..settings.streaming.throttle_seconds` | int | Padrão do Pico após ativar: `0` | Intervalo mínimo entre atualizações intermediárias; o conteúdo final sempre é enviado | +| `channel_list..settings.streaming.min_growth_chars` | int | Padrão do Pico após ativar: `1` | Crescimento mínimo de texto antes de enviar outra atualização intermediária; o conteúdo final sempre é enviado | +| `model_list[].streaming.enabled` | bool | `false` | Permite que esta entrada de modelo tente requisições de provider streaming | + +As variáveis de ambiente legadas do Telegram continuam compatíveis: `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`, `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS` e `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS`. Elas se aplicam apenas às settings do Telegram e não ativam nem modificam `settings.streaming` do Pico. + +O comportamento de falha é intencionalmente conservador: se o streaming falhar antes de qualquer chunk visível ser enviado, o PicoClaw tenta novamente uma vez pelo caminho normal `Chat()`. Se um chunk já foi mostrado ao usuário, o PicoClaw não envia uma segunda resposta sem streaming, evitando output duplicado. + ### Arquitetura de Providers PicoClaw roteia providers por família de protocolo: diff --git a/docs/guides/configuration.vi.md b/docs/guides/configuration.vi.md index d905b6d2..516ebe9f 100644 --- a/docs/guides/configuration.vi.md +++ b/docs/guides/configuration.vi.md @@ -67,6 +67,50 @@ PicoClaw lưu trữ dữ liệu trong workspace đã cấu hình (mặc định: > **Lưu ý:** Các thay đổi đối với `AGENT.md`, `SOUL.md`, `USER.md` và `memory/MEMORY.md` được tự động phát hiện trong thời gian chạy thông qua theo dõi thời gian sửa đổi file (mtime). **Không cần khởi động lại gateway** sau khi chỉnh sửa các file này — agent sẽ tải nội dung mới vào yêu cầu tiếp theo. +### Chính sách ngữ cảnh request + +`turn_profile` là chính sách tùy chọn trong `agents.defaults.turn_profile` để kiểm soát ngữ cảnh mỗi turn mới mang theo: lịch sử, system prompt, prompt skills và các tool được phép gọi. Nếu không cấu hình, hoặc đặt `"enabled": false`, PicoClaw giữ nguyên hành vi mặc định. Khi đặt `"enabled": true`, chính sách bên dưới áp dụng cho mỗi turn mới. + +Mỗi block dùng chung các giá trị `mode`: + +| Mode | Ý nghĩa | +| --- | --- | +| `default` | Giữ hành vi bình thường của PicoClaw. Block bị thiếu hoặc thiếu `mode` đều được xem là `default`. | +| `off` | Tắt block đó cho turn. | +| `custom` | Dùng danh sách cho phép. Phiên bản này chỉ hỗ trợ `custom` cho `skills` và `tools`; dùng cho `history` hoặc `system_prompt` sẽ lỗi validate. | + +Các block: + +| Block | Nội dung kiểm soát | +| --- | --- | +| `history` | Đọc lịch sử và tóm tắt, ghi tin nhắn user/assistant/tool, nạp context, compact và summarize. | +| `system_prompt` | Chèn identity mặc định của PicoClaw, chỉ dẫn workspace, memory, runtime context và summary. System prompt từ request bên ngoài vẫn được dùng khi block này `off`. | +| `skills` | Catalog skills và nội dung active skill trong prompt. `custom.allow` chỉ giữ các tên skill được liệt kê. | +| `tools` | Công cụ hiển thị cho model và được phép thực thi. `custom.allow` chỉ giữ các tool đã đăng ký và được liệt kê. | + +Khi `system_prompt.mode` là `off`, tools vẫn hiển thị và không có system prompt bên ngoài, PicoClaw dùng lại quy tắc dùng tool hiện có làm prompt fallback tối thiểu. Nếu `tools.mode` là `off`, fallback này không được thêm. + +Ví dụ ngữ cảnh sạch chỉ giữ tool web: + +```json +{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": { "mode": "off" }, + "system_prompt": { "mode": "off" }, + "skills": { "mode": "off" }, + "tools": { + "mode": "custom", + "allow": ["web_search", "web_fetch"] + } + } + } + } +} +``` + ### Nguồn Skill Mặc định, skill được tải từ: @@ -365,6 +409,55 @@ Cấu hình nhiều endpoint cho cùng tên mô hình — PicoClaw sẽ tự đ Cấu hình `providers` cũ đã **bị deprecated** và đã được loại bỏ trong V2. Các cấu hình V0/V1 hiện có sẽ được tự động migrate. Xem [docs/migration/model-list-migration.md](../migration/model-list-migration.md). +#### Cấu Hình Streaming + +Provider streaming dùng cơ chế double opt-in và bị tắt theo mặc định. Agent chỉ thử streaming khi channel hiện tại có `settings.streaming.enabled: true`, entry model đang dùng có `streaming.enabled: true`, và cả provider lẫn channel đều hỗ trợ streaming. Nếu thiếu bất kỳ điều kiện nào, PicoClaw dùng đường dẫn yêu cầu không streaming thông thường. + +Pico WebUI là channel đầu tiên được nối đầy đủ. Pico tạo message assistant đầu tiên bằng wire message hiện có `message.create`, sau đó cập nhật chính message đó bằng `message.update`; không thêm loại wire message Pico mới. + +Hãy để trống `streaming` khi bạn không muốn dùng streaming. Bỏ qua block `streaming` nghĩa là đã tắt; bạn không cần viết `"streaming": {"enabled": false}`. + +Ví dụ bật streaming: + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| Trường | Kiểu | Mặc định | Mô tả | +| ------ | ---- | -------- | ----- | +| `channel_list..settings.streaming.enabled` | bool | `false` | Cho phép channel này hiển thị output streaming từ provider | +| `channel_list..settings.streaming.throttle_seconds` | int | Mặc định Pico sau khi bật: `0` | Khoảng cách tối thiểu giữa các cập nhật trung gian; nội dung cuối luôn được flush | +| `channel_list..settings.streaming.min_growth_chars` | int | Mặc định Pico sau khi bật: `1` | Số ký tự tăng tối thiểu trước khi gửi cập nhật trung gian; nội dung cuối luôn được flush | +| `model_list[].streaming.enabled` | bool | `false` | Cho phép entry model này thử yêu cầu provider streaming | + +Các biến môi trường Telegram cũ vẫn tương thích: `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`, `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS`, và `PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS`. Chúng chỉ áp dụng cho Telegram settings và không bật hoặc thay đổi `settings.streaming` của Pico. + +Hành vi lỗi được giữ thận trọng: nếu streaming lỗi trước khi gửi bất kỳ chunk hiển thị nào, PicoClaw thử lại một lần qua đường dẫn `Chat()` thông thường. Nếu đã có chunk hiển thị cho người dùng, PicoClaw không gửi thêm một câu trả lời non-streaming thứ hai để tránh lặp output. + ### Kiến Trúc Provider PicoClaw định tuyến provider theo họ giao thức: diff --git a/docs/guides/configuration.zh.md b/docs/guides/configuration.zh.md index a6e4d42b..1d04a0d4 100644 --- a/docs/guides/configuration.zh.md +++ b/docs/guides/configuration.zh.md @@ -97,6 +97,50 @@ PicoClaw 将数据存储在您配置的工作区中（默认：`~/.picoclaw/work 如果你只想先检查学习记录，建议从 `observe` 开始。需要生成可审查改进时使用 `draft`。只有在你接受让已通过的草稿更新工作区技能时，才使用 `apply`。 +### 请求上下文策略 + +`turn_profile` 是 `agents.defaults.turn_profile` 下的可选请求上下文策略，用来控制每个新回合是否带入历史、系统提示、技能提示，以及允许调用哪些工具。不写该配置或设置 `"enabled": false` 时，PicoClaw 完全保持原逻辑；设置 `"enabled": true` 后，下面的策略会应用到每个新回合。 + +所有块都使用同一组 `mode`： + +| Mode | 含义 | +| --- | --- | +| `default` | 保持 PicoClaw 原逻辑。块缺失或 `mode` 缺失都按 `default` 处理。 | +| `off` | 关闭该块。 | +| `custom` | 使用允许列表。本版本仅支持 `skills` 和 `tools`，在 `history` 或 `system_prompt` 中使用会触发配置校验错误。 | + +各块含义： + +| 块 | 控制内容 | +| --- | --- | +| `history` | 是否读取历史和摘要、写入用户/助手/工具消息、写入 context manager，以及是否压缩或总结本轮会话。 | +| `system_prompt` | 是否注入 PicoClaw 默认身份、工作区指令、记忆、运行时上下文和摘要。关闭后仍可使用外部传入的 system prompt。 | +| `skills` | 是否加载技能目录和 active skill 提示。`custom.allow` 只保留列出的技能名。 | +| `tools` | 暴露给模型并允许执行的工具。`custom.allow` 只保留已注册且列出的工具名。 | + +当 `system_prompt.mode` 为 `off`、工具仍可见且没有外部 system prompt 时，PicoClaw 会复用现有的工具使用规则作为最小兜底提示。如果 `tools.mode` 为 `off`，则不会添加兜底提示。 + +只保留 Web 工具的干净上下文示例： + +```json +{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": { "mode": "off" }, + "system_prompt": { "mode": "off" }, + "skills": { "mode": "off" }, + "tools": { + "mode": "custom", + "allow": ["web_search", "web_fetch"] + } + } + } + } +} +``` + ### Web 启动器控制台 用 **picoclaw-launcher** 打开浏览器控制台前需要先使用密码登录。首次启动时打开 `/launcher-setup` 创建 dashboard 登录密码；后续手动登录使用 `/launcher-login`。 @@ -293,7 +337,7 @@ PicoClaw 默认在沙箱环境中运行。Agent 只能访问配置的工作区 | `tools.exec.custom_deny_patterns` | string[] | `[]` | 自定义阻止的正则表达式模式 | | `tools.exec.custom_allow_patterns` | string[] | `[]` | 自定义允许的正则表达式模式 | -> **安全提示：** Symlink 保护默认启用——所有文件路径在白名单匹配前都会通过 `filepath.EvalSymlinks` 解析，防止符号链接逃逸攻击。 +> **安全提示：** Symlink 保护默认启用——所有文件路径在允许列表匹配前都会通过 `filepath.EvalSymlinks` 解析，防止符号链接逃逸攻击。 #### 已知限制：构建工具的子进程 @@ -537,6 +581,55 @@ Agent 读取 HEARTBEAT.md - 如果未设置 `provider`，PicoClaw 会把 `model` 第一个 `/` 之前的字段当作 provider，并把第一个 `/` 之后的全部内容当作最终模型 ID。 - 这意味着 `"model": "openrouter/openai/gpt-5.4"` 这样的兼容写法仍然可用，并会把 `openai/gpt-5.4` 发送给 OpenRouter。 +#### 流式输出配置 + +Provider 流式输出采用双开关，默认关闭。只有当前 channel 的 `settings.streaming.enabled` 和当前模型条目的 `streaming.enabled` 都为 `true`，并且 provider 与 channel 都支持流式能力时，Agent 才会尝试流式请求；任一条件不满足时仍使用普通非流式请求。 + +当前完整落地的是 Pico WebUI。Pico 使用已有的 `message.create` 创建第一条 assistant 消息，随后用 `message.update` 更新同一条消息，不新增协议消息类型。 + +不需要流式时请省略 `streaming` 配置块。省略表示关闭，不需要写 `"streaming": {"enabled": false}`。 + +开启示例： + +```json +{ + "model_list": [ + { + "model_name": "gpt-5.4", + "provider": "openai", + "model": "gpt-5.4", + "api_keys": ["sk-your-openai-key"], + "streaming": { + "enabled": true + } + } + ], + "channel_list": { + "pico": { + "enabled": true, + "type": "pico", + "settings": { + "token": "YOUR_PICO_TOKEN", + "streaming": { + "enabled": true + } + } + } + } +} +``` + +| 字段 | 类型 | 默认值 | 说明 | +|------|------|--------|------| +| `channel_list..settings.streaming.enabled` | bool | `false` | 是否允许该 channel 尝试展示 provider 流式输出 | +| `channel_list..settings.streaming.throttle_seconds` | int | Pico 开启后默认 `0` | 中间更新的最小时间间隔，最终内容不受此限制 | +| `channel_list..settings.streaming.min_growth_chars` | int | Pico 开启后默认 `1` | 中间更新相比上次发送至少增长的字符数，最终内容不受此限制 | +| `model_list[].streaming.enabled` | bool | `false` | 是否允许该模型条目尝试 provider 流式请求 | + +Telegram 旧环境变量仍兼容：`PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED`、`PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS`、`PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS`。这些环境变量只作用于 Telegram settings，不会开启或修改 Pico 的 `settings.streaming`。 + +失败处理保持保守：如果还没有任何可见 chunk 就失败，PicoClaw 会回退到普通 `Chat()` 路径重试一次；如果已经有 chunk 展示给用户，则不会再发送一条非流式最终答案，避免界面重复输出。 + #### 各厂商配置示例
diff --git a/docs/guides/providers.fr.md b/docs/guides/providers.fr.md index aff60035..a16d7c80 100644 --- a/docs/guides/providers.fr.md +++ b/docs/guides/providers.fr.md @@ -113,10 +113,13 @@ Cette conception permet également le **support multi-agents** avec une sélecti | `max_tokens_field` | string | Non | Remplace le nom du champ max tokens dans le corps de la requête (ex : `max_completion_tokens` pour les modèles o1) | | `thinking_level` | string | Non | Niveau de pensée étendue : `off`, `low`, `medium`, `high`, `xhigh` ou `adaptive` | | `extra_body` | object | Non | Champs supplémentaires à injecter dans chaque corps de requête | +| `streaming.enabled` | bool | Non | Opt-in pour le streaming provider sur cette entrée de modèle. Par défaut `false`, et le channel actif doit aussi avoir `settings.streaming.enabled` à `true` | | `rpm` | int | Non | Limite de requêtes par minute | | `fallbacks` | string[] | Non | Noms des modèles de secours pour le basculement automatique | | `enabled` | bool | Non | Activer ou désactiver cette entrée de modèle (par défaut : `true`) | +Lorsque le streaming est désactivé, omettez le bloc `streaming`. Écrire `"streaming": {"enabled": false}` est optionnel et n'est pas nécessaire. + #### Exemples par Vendor **OpenAI** diff --git a/docs/guides/providers.ja.md b/docs/guides/providers.ja.md index fecc7451..75ee28a5 100644 --- a/docs/guides/providers.ja.md +++ b/docs/guides/providers.ja.md @@ -114,10 +114,13 @@ | `max_tokens_field` | string | いいえ | リクエストボディの max tokens フィールド名を上書き（例：o1 モデルでは `max_completion_tokens`） | | `thinking_level` | string | いいえ | 拡張思考レベル：`off`、`low`、`medium`、`high`、`xhigh`、`adaptive` | | `extra_body` | object | いいえ | 各リクエストボディに注入する追加フィールド | +| `streaming.enabled` | bool | いいえ | このモデルエントリで provider ストリーミングを試行するための opt-in。デフォルトは `false` で、アクティブな channel の `settings.streaming.enabled` も `true` である必要があります | | `rpm` | int | いいえ | 1 分あたりのリクエストレート制限 | | `fallbacks` | string[] | いいえ | 自動フェイルオーバーのフォールバックモデル名 | | `enabled` | bool | いいえ | このモデルエントリを有効にするかどうか（デフォルト：`true`） | +ストリーミングを無効にする場合は `streaming` ブロックを省略してください。`"streaming": {"enabled": false}` を書くことは任意であり、必須ではありません。 + #### ベンダー別設定例 **OpenAI** diff --git a/docs/guides/providers.md b/docs/guides/providers.md index 7b078373..40edb87a 100644 --- a/docs/guides/providers.md +++ b/docs/guides/providers.md @@ -131,10 +131,13 @@ This design also enables **multi-agent support** with flexible provider selectio | `tool_schema_transform` | string | No | Optional compatibility transform for tool parameter schemas. Default: disabled. Supported values: `simple`. | | `extra_body` | object | No | Additional fields to inject into every request body | | `custom_headers` | object | No | Additional HTTP headers to inject into every request (e.g., `{"X-Source":"coding-plan"}`). If a key matches a built-in header, the custom value overrides the built-in one (e.g., `Authorization`, `User-Agent`, `Content-Type`, `Accept`). | +| `streaming.enabled` | bool | No | Opt-in for provider streaming on this model entry. Defaults to `false` and also requires the active channel's `settings.streaming.enabled` to be `true`. | | `rpm` | int | No | Per-minute request rate limit | | `fallbacks` | string[] | No | Fallback model names for automatic failover | | `enabled` | bool | No | Whether this model entry is active (default: `true`) | +When streaming is disabled, omit the `streaming` block. Writing `"streaming": {"enabled": false}` is optional and not needed in generated or hand-written config. + #### Tool Schema Compatibility By default, PicoClaw now forwards tool JSON Schemas unchanged. diff --git a/docs/guides/providers.pt-br.md b/docs/guides/providers.pt-br.md index 0d45dc30..ed99a27d 100644 --- a/docs/guides/providers.pt-br.md +++ b/docs/guides/providers.pt-br.md @@ -113,10 +113,13 @@ Este design também permite **suporte multi-agente** com seleção flexível de | `max_tokens_field` | string | Não | Substitui o nome do campo max tokens no corpo da requisição (ex: `max_completion_tokens` para modelos o1) | | `thinking_level` | string | Não | Nível de pensamento estendido: `off`, `low`, `medium`, `high`, `xhigh` ou `adaptive` | | `extra_body` | object | Não | Campos adicionais para injetar em cada corpo de requisição | +| `streaming.enabled` | bool | Não | Opt-in para provider streaming nesta entrada de modelo. O padrão é `false` e o canal ativo também precisa de `settings.streaming.enabled` como `true` | | `rpm` | int | Não | Limite de requisições por minuto | | `fallbacks` | string[] | Não | Nomes dos modelos de fallback para failover automático | | `enabled` | bool | Não | Ativar ou desativar esta entrada de modelo (padrão: `true`) | +Quando streaming estiver desativado, omita o bloco `streaming`. Escrever `"streaming": {"enabled": false}` é opcional e não é necessário. + #### Exemplos por Vendor **OpenAI** diff --git a/docs/guides/providers.vi.md b/docs/guides/providers.vi.md index c354461c..402e09fb 100644 --- a/docs/guides/providers.vi.md +++ b/docs/guides/providers.vi.md @@ -113,10 +113,13 @@ Thiết kế này cũng cho phép **hỗ trợ đa agent** với lựa chọn pr | `max_tokens_field` | string | Không | Ghi đè tên trường max tokens trong request body (ví dụ: `max_completion_tokens` cho model o1) | | `thinking_level` | string | Không | Mức độ tư duy mở rộng: `off`, `low`, `medium`, `high`, `xhigh` hoặc `adaptive` | | `extra_body` | object | Không | Các trường bổ sung để chèn vào mỗi request body | +| `streaming.enabled` | bool | Không | Opt-in cho provider streaming trên entry model này. Mặc định là `false` và channel đang hoạt động cũng cần `settings.streaming.enabled` là `true` | | `rpm` | int | Không | Giới hạn tốc độ yêu cầu mỗi phút | | `fallbacks` | string[] | Không | Tên model dự phòng cho failover tự động | | `enabled` | bool | Không | Kích hoạt hay vô hiệu hóa entry model này (mặc định: `true`) | +Khi không dùng streaming, hãy bỏ qua block `streaming`. Viết `"streaming": {"enabled": false}` là tùy chọn và không cần thiết. + #### Ví Dụ Theo Vendor **OpenAI** diff --git a/docs/guides/providers.zh.md b/docs/guides/providers.zh.md index 4bab65f6..c6392714 100644 --- a/docs/guides/providers.zh.md +++ b/docs/guides/providers.zh.md @@ -123,6 +123,7 @@ | `proxy` | string | 否 | 此模型条目的 HTTP 代理 URL | | `user_agent` | string | 否 | 自定义 `User-Agent` 请求头（支持 OpenAI 兼容、Gemini、Anthropic 和 Azure provider） | | `request_timeout` | int | 否 | 请求超时时间（秒），默认值因 provider 而异 | +| `streaming.enabled` | bool | 否 | 是否允许此模型条目尝试 provider 流式请求，默认 `false`。它只表达模型/端点能力 opt-in，实际还需要当前 channel 的 `settings.streaming.enabled` 同时开启 | | `max_tokens_field` | string | 否 | 覆盖请求体中 max tokens 的字段名（如 o1 模型使用 `max_completion_tokens`） | | `thinking_level` | string | 否 | 扩展思考级别：`off`、`low`、`medium`、`high`、`xhigh` 或 `adaptive` | | `extra_body` | object | 否 | 注入到每个请求体中的额外字段 | @@ -131,6 +132,8 @@ | `fallbacks` | string[] | 否 | 自动故障转移的备用模型名称 | | `enabled` | bool | 否 | 是否启用此模型条目（默认：`true`） | +不需要流式时请省略 `streaming` 配置块。写 `"streaming": {"enabled": false}` 是可选的，手写或生成配置时都不需要。 + #### `provider` / `model` 解析规则 PicoClaw 按下面的规则解析 `provider` 和最终发给上游的模型 ID： diff --git a/go.mod b/go.mod index f4e494d3..c3c05f96 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.25.10 require ( fyne.io/systray v1.12.1 github.com/SevereCloud/vksdk/v3 v3.3.1 - github.com/adhocore/gronx v1.19.7 + github.com/adhocore/gronx v1.20.0 github.com/anthropics/anthropic-sdk-go v1.26.0 github.com/atc0005/go-teams-notify/v2 v2.14.0 github.com/aws/aws-sdk-go-v2 v1.41.7 @@ -23,7 +23,7 @@ require ( github.com/gorilla/websocket v1.5.3 github.com/h2non/filetype v1.1.3 github.com/larksuite/oapi-sdk-go/v3 v3.7.5 - github.com/line/line-bot-sdk-go/v8 v8.19.0 + github.com/line/line-bot-sdk-go/v8 v8.20.0 github.com/mdp/qrterminal/v3 v3.2.1 github.com/minio/selfupdate v0.6.0 github.com/modelcontextprotocol/go-sdk v1.5.0 @@ -135,9 +135,9 @@ require ( github.com/yosida95/uritemplate/v3 v3.0.2 // indirect golang.org/x/arch v0.24.0 // indirect golang.org/x/crypto v0.51.0 - golang.org/x/net v0.54.0 + golang.org/x/net v0.55.0 golang.org/x/sync v0.20.0 - golang.org/x/sys v0.44.0 + golang.org/x/sys v0.45.0 ) replace github.com/bwmarrin/discordgo => github.com/yeongaori/discordgo-fork v0.0.0-20260319072544-e8e546f5d532 diff --git a/go.sum b/go.sum index a4137887..27fbdab1 100644 --- a/go.sum +++ b/go.sum @@ -9,8 +9,8 @@ github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7Oputl github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU= github.com/SevereCloud/vksdk/v3 v3.3.1 h1:O86zsp5LQnHE+O5acvuXM/s6S1LyxzVTkF6+Lup0Jyg= github.com/SevereCloud/vksdk/v3 v3.3.1/go.mod h1:c6WaA5aocUYsXfkcUbg2qy45V9M1VDcqHHmHIN14NAw= -github.com/adhocore/gronx v1.19.7 h1:7hhFwChgDw9eHC3+TQ+OKKBqJnP44oWkDCnnW9nrsuA= -github.com/adhocore/gronx v1.19.7/go.mod h1:7oUY1WAU8rEJWmAxXR2DN0JaO4gi9khSgKjiRypqteg= +github.com/adhocore/gronx v1.20.0 h1:PD13Mo0wekkZ7ZZR9yb1TqeqTfybs7/K3ez9DmjQwEs= +github.com/adhocore/gronx v1.20.0/go.mod h1:7oUY1WAU8rEJWmAxXR2DN0JaO4gi9khSgKjiRypqteg= github.com/agnivade/levenshtein v1.2.1 h1:EHBY3UOn1gwdy/VbFwgo4cxecRznFk7fKWN1KOX7eoM= github.com/agnivade/levenshtein v1.2.1/go.mod h1:QVVI16kDrtSuwcpd0p1+xMC6Z/VfhtCyDIjcwga4/DU= github.com/andreyvit/diff v0.0.0-20170406064948-c7f18ee00883 h1:bvNMNQO63//z+xNgfBlViaCIJKLlCJ6/fmUseuG0wVQ= @@ -181,8 +181,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/larksuite/oapi-sdk-go/v3 v3.7.5 h1:dimv+ZAGia01f4xCDGvCiBHKWMf4K1AB7fGsM+lv5Jw= github.com/larksuite/oapi-sdk-go/v3 v3.7.5/go.mod h1:ZEplY+kwuIrj/nqw5uSCINNATcH3KdxSN7y+UxYY5fI= -github.com/line/line-bot-sdk-go/v8 v8.19.0 h1:5FD/1SprRZ8Y0FiUI6syYiBewOs0ak2tuUBMYN0wzE4= -github.com/line/line-bot-sdk-go/v8 v8.19.0/go.mod h1:AeSRUuu7WGgveGDJb6DyKyFUOst2UB2aF6LO2cQeuXs= +github.com/line/line-bot-sdk-go/v8 v8.20.0 h1:Jv22DV3JuQ5qZvniqUbg504bJrVzffXs2CMpyoiuIZU= +github.com/line/line-bot-sdk-go/v8 v8.20.0/go.mod h1:QMXJwPka2ysSeVQKWXkBp8DzBFs+CFAXFNo75KJtWho= github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE= @@ -356,6 +356,8 @@ golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U= golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w= golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= golang.org/x/oauth2 v0.23.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= @@ -390,6 +392,8 @@ golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ= golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= diff --git a/pkg/agent/adapters/messagebus.go b/pkg/agent/adapters/messagebus.go index ccae7e8b..e7897337 100644 --- a/pkg/agent/adapters/messagebus.go +++ b/pkg/agent/adapters/messagebus.go @@ -31,6 +31,10 @@ func (a *messageBusAdapter) PublishOutboundMedia(ctx context.Context, msg bus.Ou return a.inner.PublishOutboundMedia(ctx, msg) } +func (a *messageBusAdapter) GetStreamer(ctx context.Context, channel, chatID, sessionKey string) (bus.Streamer, bool) { + return a.inner.GetStreamer(ctx, channel, chatID, sessionKey) +} + func (a *messageBusAdapter) InboundChan() <-chan bus.InboundMessage { return a.inner.InboundChan() } diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 5749149c..90745e5d 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -79,17 +79,18 @@ type AgentLoop struct { // processOptions configures how a message is processed type processOptions struct { - Dispatch DispatchRequest // Normalized routed request boundary for this turn - SessionKey string // Session identifier for history/context - SessionAliases []string // Compatibility aliases for the session key - Channel string // Target channel for tool execution - ChatID string // Target chat ID for tool execution - MessageID string // Current inbound platform message ID - ReplyToMessageID string // Current inbound reply target message ID - SenderID string // Current sender ID for dynamic context - SenderDisplayName string // Current sender display name for dynamic context - UserMessage string // User message content (may include prefix) - ForcedSkills []string // Skills explicitly requested for this message + Dispatch DispatchRequest // Normalized routed request boundary for this turn + SessionKey string // Session identifier for history/context + SessionAliases []string // Compatibility aliases for the session key + Channel string // Target channel for tool execution + ChatID string // Target chat ID for tool execution + MessageID string // Current inbound platform message ID + ReplyToMessageID string // Current inbound reply target message ID + SenderID string // Current sender ID for dynamic context + SenderDisplayName string // Current sender display name for dynamic context + UserMessage string // User message content (may include prefix) + ForcedSkills []string // Skills explicitly requested for this message + TurnProfile config.EffectiveTurnProfile SystemPromptOverride string // Override the default system prompt (Used by SubTurns) Media []string // media:// refs from inbound message InitialSteeringMessages []providers.Message // Steering messages from refactor/agent @@ -119,9 +120,11 @@ const ( pendingTurnPrefix = "pending-" metadataKeyMessageKind = "message_kind" metadataKeyToolCalls = "tool_calls" + metadataKeyOutboundKind = "outbound_kind" messageKindThought = "thought" messageKindToolFeedback = "tool_feedback" messageKindToolCalls = "tool_calls" + outboundKindFinal = "final" metadataKeyAccountID = "account_id" metadataKeyGuildID = "guild_id" metadataKeyTeamID = "team_id" @@ -532,17 +535,22 @@ func (al *AgentLoop) runAgentLoop( opts processOptions, ) (string, error) { opts = normalizeProcessOptions(opts) + var err error + opts, err = resolveTurnProfileOptions(al.GetConfig(), opts) + if err != nil { + return "", err + } // Record last channel for heartbeat notifications (skip internal channels and cli) if opts.Dispatch.Channel() != "" && opts.Dispatch.ChatID() != "" && !constants.IsInternalChannel(opts.Dispatch.Channel()) { channelKey := fmt.Sprintf("%s:%s", opts.Dispatch.Channel(), opts.Dispatch.ChatID()) - if err := al.RecordLastChannel(channelKey); err != nil { + if recordErr := al.RecordLastChannel(channelKey); recordErr != nil { logger.WarnCF( "agent", "Failed to record last channel", - map[string]any{"error": err.Error()}, + map[string]any{"error": recordErr.Error()}, ) } } @@ -585,7 +593,7 @@ func (al *AgentLoop) runAgentLoop( opts.Dispatch.SessionKey, opts.Dispatch.SessionScope, ) - al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + msg := bus.OutboundMessage{ Context: outboundContextFromInbound( opts.Dispatch.InboundContext, opts.Dispatch.Channel(), @@ -597,7 +605,15 @@ func (al *AgentLoop) runAgentLoop( Scope: scope, Content: result.finalContent, ContextUsage: computeContextUsage(agent, opts.Dispatch.SessionKey), - }) + } + if modelName := strings.TrimSpace(result.modelName); modelName != "" { + if msg.Context.Raw == nil { + msg.Context.Raw = make(map[string]string, 1) + } + msg.Context.Raw["model_name"] = modelName + } + markFinalOutbound(&msg) + al.bus.PublishOutbound(ctx, msg) } if result.finalContent != "" { diff --git a/pkg/agent/agent_command.go b/pkg/agent/agent_command.go index ae0293d7..944c58f3 100644 --- a/pkg/agent/agent_command.go +++ b/pkg/agent/agent_command.go @@ -319,6 +319,7 @@ func (al *AgentLoop) buildCommandsRuntime( agent.Provider = nextProvider agent.Candidates = nextCandidates agent.ThinkingLevel = parseThinkingLevel(modelCfg.ThinkingLevel) + agent.ThinkingLevelConfigured = isConfiguredThinkingLevel(modelCfg.ThinkingLevel) if oldProvider != nil && oldProvider != nextProvider { if stateful, ok := oldProvider.(providers.StatefulProvider); ok { diff --git a/pkg/agent/agent_message.go b/pkg/agent/agent_message.go index 4d2886a8..419664ed 100644 --- a/pkg/agent/agent_message.go +++ b/pkg/agent/agent_message.go @@ -193,6 +193,11 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) SendResponse: false, AllowInterimPicoPublish: true, } + var err error + opts, err = resolveTurnProfileOptions(al.GetConfig(), opts) + if err != nil { + return "", err + } // context-dependent commands check their own Runtime fields and report // "unavailable" when the required capability is nil. diff --git a/pkg/agent/agent_outbound.go b/pkg/agent/agent_outbound.go index f4a01adf..f5de6fa4 100644 --- a/pkg/agent/agent_outbound.go +++ b/pkg/agent/agent_outbound.go @@ -75,12 +75,14 @@ func (al *AgentLoop) PublishResponseIfNeeded(ctx context.Context, channel, chatI } msg := bus.OutboundMessage{ - Context: bus.NewOutboundContext(channel, chatID, ""), - Content: response, + Context: bus.NewOutboundContext(channel, chatID, ""), + SessionKey: sessionKey, + Content: response, } if sessionKey != "" { msg.ContextUsage = computeContextUsage(al.agentForSession(sessionKey), sessionKey) } + markFinalOutbound(&msg) al.bus.PublishOutbound(ctx, msg) logger.InfoCF("agent", "Published outbound response", map[string]any{ @@ -100,7 +102,10 @@ func (al *AgentLoop) targetReasoningChannelID(channelName string) (chatID string return "" } -func (al *AgentLoop) publishPicoReasoning(ctx context.Context, reasoningContent, chatID string) { +func (al *AgentLoop) publishPicoReasoning( + ctx context.Context, + reasoningContent, chatID, sessionKey, modelName string, +) { if reasoningContent == "" || chatID == "" { return } @@ -112,15 +117,19 @@ func (al *AgentLoop) publishPicoReasoning(ctx context.Context, reasoningContent, pubCtx, pubCancel := context.WithTimeout(ctx, 5*time.Second) defer pubCancel() + raw := map[string]string{metadataKeyMessageKind: messageKindThought} + if trimmedModelName := strings.TrimSpace(modelName); trimmedModelName != "" { + raw["model_name"] = trimmedModelName + } + if err := al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{ Context: bus.InboundContext{ Channel: "pico", ChatID: chatID, - Raw: map[string]string{ - metadataKeyMessageKind: messageKindThought, - }, + Raw: raw, }, - Content: reasoningContent, + SessionKey: sessionKey, + Content: reasoningContent, }); err != nil { if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || errors.Is(err, bus.ErrBusClosed) { @@ -140,6 +149,7 @@ func (al *AgentLoop) publishPicoReasoning(ctx context.Context, reasoningContent, func (al *AgentLoop) publishPicoToolCallInterim( ctx context.Context, ts *turnState, + modelName string, reasoningContent string, content string, toolCalls []providers.ToolCall, @@ -152,7 +162,14 @@ func (al *AgentLoop) publishPicoToolCallInterim( pubCtx, pubCancel := context.WithTimeout(ctx, 3*time.Second) err := al.bus.PublishOutbound( pubCtx, - outboundMessageForTurnWithKind(ts, reasoningContent, messageKindThought), + outboundMessageForTurnWithOptions( + ts, + reasoningContent, + outboundTurnMessageOptions{ + kind: messageKindThought, + modelName: modelName, + }, + ), ) pubCancel() if err != nil && !errors.Is(err, context.DeadlineExceeded) && @@ -179,7 +196,12 @@ func (al *AgentLoop) publishPicoToolCallInterim( if strings.TrimSpace(content) != "" && !duplicateToolCallContent { pubCtx, pubCancel := context.WithTimeout(ctx, 3*time.Second) - err := al.bus.PublishOutbound(pubCtx, outboundMessageForTurn(ts, content)) + err := al.bus.PublishOutbound( + pubCtx, + outboundMessageForTurnWithOptions(ts, content, outboundTurnMessageOptions{ + modelName: modelName, + }), + ) pubCancel() if err != nil && !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) && @@ -206,11 +228,13 @@ func (al *AgentLoop) publishPicoToolCallInterim( return } - msg := outboundMessageForTurnWithKind(ts, "", messageKindToolCalls) - if msg.Context.Raw == nil { - msg.Context.Raw = map[string]string{} - } - msg.Context.Raw[metadataKeyToolCalls] = string(rawToolCalls) + msg := outboundMessageForTurnWithOptions(ts, "", outboundTurnMessageOptions{ + kind: messageKindToolCalls, + modelName: modelName, + raw: map[string]string{ + metadataKeyToolCalls: string(rawToolCalls), + }, + }) pubCtx, pubCancel := context.WithTimeout(ctx, 3*time.Second) err = al.bus.PublishOutbound(pubCtx, msg) diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index 7a869ec9..08cf71d7 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "maps" "net/http" "net/http/httptest" "os" @@ -138,6 +139,121 @@ func (r *recordingProvider) GetDefaultModel() string { return "mock-model" } +type thinkingRecordingProvider struct { + lastOptions map[string]any +} + +func (r *thinkingRecordingProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + r.lastOptions = maps.Clone(opts) + return &providers.LLMResponse{ + Content: "Mock response", + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (r *thinkingRecordingProvider) GetDefaultModel() string { + return "mock-model" +} + +func (r *thinkingRecordingProvider) SupportsThinking() bool { + return true +} + +type thinkingOptionRecordingProvider struct { + lastOptions map[string]any +} + +func (r *thinkingOptionRecordingProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + r.lastOptions = maps.Clone(opts) + return &providers.LLMResponse{ + Content: "Mock response", + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (r *thinkingOptionRecordingProvider) GetDefaultModel() string { + return "mock-model" +} + +type reasoningOptionRecordingProvider struct { + lastOptions map[string]any +} + +func (r *reasoningOptionRecordingProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + r.lastOptions = maps.Clone(opts) + return &providers.LLMResponse{ + Content: "final answer", + ReasoningContent: "thinking trace", + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (r *reasoningOptionRecordingProvider) GetDefaultModel() string { + return "mock-model" +} + +type reasoningResponseProvider struct{} + +func (p *reasoningResponseProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + return &providers.LLMResponse{ + Content: "Mock response", + ReasoningContent: "thinking trace", + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (p *reasoningResponseProvider) GetDefaultModel() string { + return "mock-model" +} + +type sideQuestionFallbackTestProvider struct { + model string +} + +func (p *sideQuestionFallbackTestProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + if p.model == "test-model" { + return nil, context.DeadlineExceeded + } + return &providers.LLMResponse{ + ReasoningContent: "thinking trace", + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (p *sideQuestionFallbackTestProvider) GetDefaultModel() string { + return p.model +} + type modelRewriteHook struct { model string } @@ -284,6 +400,55 @@ func TestPublishResponseIfNeeded_DismissesToolFeedbackWhenMessageToolAlreadySent } } +func TestPublishResponseIfNeeded_MarksFinalOutbound(t *testing.T) { + al, _, msgBus, provider, cleanup := newTestAgentLoop(t) + defer cleanup() + _ = provider + + al.PublishResponseIfNeeded(context.Background(), "pico", "pico:session-1", "session-1", "final reply") + + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "final reply" { + t.Fatalf("outbound content = %q, want final reply", outbound.Content) + } + if outbound.Context.Raw[metadataKeyOutboundKind] != outboundKindFinal { + t.Fatalf("outbound kind = %q, want %q", outbound.Context.Raw[metadataKeyOutboundKind], outboundKindFinal) + } + if outbound.SessionKey != "session-1" { + t.Fatalf("outbound session key = %q, want session-1", outbound.SessionKey) + } + case <-time.After(time.Second): + t.Fatal("expected final outbound") + } +} + +func TestPublishPicoReasoningIncludesSessionKey(t *testing.T) { + al, _, msgBus, provider, cleanup := newTestAgentLoop(t) + defer cleanup() + _ = provider + + al.publishPicoReasoning(context.Background(), "reasoning", "pico-chat", "session-1", "") + + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Channel != "pico" || outbound.ChatID != "pico-chat" { + t.Fatalf("unexpected outbound target: %+v", outbound) + } + if outbound.Content != "reasoning" { + t.Fatalf("outbound content = %q, want reasoning", outbound.Content) + } + if outbound.SessionKey != "session-1" { + t.Fatalf("outbound session key = %q, want session-1", outbound.SessionKey) + } + if outbound.Context.Raw[metadataKeyMessageKind] != messageKindThought { + t.Fatalf("message kind = %q, want %q", outbound.Context.Raw[metadataKeyMessageKind], messageKindThought) + } + case <-time.After(time.Second): + t.Fatal("expected pico reasoning outbound") + } +} + func TestProcessMessage_IncludesCurrentSenderInDynamicContext(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { @@ -337,6 +502,500 @@ func TestProcessMessage_IncludesCurrentSenderInDynamicContext(t *testing.T) { } } +func TestProcessMessage_DoesNotPassImplicitThinkingOffToCapableProvider(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + provider := &thinkingRecordingProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + + response, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "pico", + ChatID: "chat-1", + Content: "hello", + })) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "Mock response" { + t.Fatalf("processMessage() response = %q, want %q", response, "Mock response") + } + if _, ok := provider.lastOptions["thinking_level"]; ok { + t.Fatalf("thinking_level option should be omitted when unset, got %#v", provider.lastOptions["thinking_level"]) + } +} + +func TestProcessMessage_PassesExplicitThinkingOffToCapableProvider(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "test-model", + ThinkingLevel: "off", + }}, + } + + provider := &thinkingRecordingProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + + response, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "pico", + ChatID: "chat-1", + Content: "hello", + })) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "Mock response" { + t.Fatalf("processMessage() response = %q, want %q", response, "Mock response") + } + if got := provider.lastOptions["thinking_level"]; got != "off" { + t.Fatalf("thinking_level option = %#v, want %q", got, "off") + } +} + +func TestProcessMessage_PassesExplicitThinkingOffToProviderWithoutThinkingCapability(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "test-model", + ThinkingLevel: "off", + }}, + } + + provider := &thinkingOptionRecordingProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + + response, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "pico", + ChatID: "chat-1", + Content: "hello", + })) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "Mock response" { + t.Fatalf("processMessage() response = %q, want %q", response, "Mock response") + } + if got := provider.lastOptions["thinking_level"]; got != "off" { + t.Fatalf("thinking_level option = %#v, want %q", got, "off") + } +} + +func TestProcessMessage_PassesDeepSeekThinkingLevelToThinkingCapableProvider(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "deepseek-v4-flash", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "deepseek-v4-flash", + Provider: "deepseek", + Model: "deepseek-v4-flash", + ThinkingLevel: "xhigh", + }}, + } + + provider := &thinkingRecordingProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + + response, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "pico", + ChatID: "chat-1", + Content: "hello", + })) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "Mock response" { + t.Fatalf("processMessage() response = %q, want %q", response, "Mock response") + } + if got := provider.lastOptions["thinking_level"]; got != "xhigh" { + t.Fatalf("thinking_level option = %#v, want %q", got, "xhigh") + } +} + +func TestProcessMessage_SuppressesReasoningWhenThinkingOff(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "test-model", + ThinkingLevel: "off", + }}, + } + + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, &reasoningResponseProvider{}) + + response, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:main:pico:chat-1", + Channel: "pico", + ChatID: "chat-1", + UserMessage: "hello", + SendResponse: false, + DefaultResponse: defaultResponse, + NoHistory: true, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if response != "Mock response" { + t.Fatalf("response = %q, want %q", response, "Mock response") + } + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("expected no reasoning outbound when thinking is off, got %+v", outbound) + case <-time.After(50 * time.Millisecond): + } +} + +func TestProcessMessage_BeforeLLMModelRewriteReevaluatesThinkingOff(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "plain-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "plain-model", + Model: "openai/plain-model", + }, + { + ModelName: "off-model", + Model: "openai/off-model", + ThinkingLevel: "off", + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &reasoningOptionRecordingProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + if err := al.MountHook(NamedHook("rewrite-model", modelRewriteHook{model: "off-model"})); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "pico", + SenderID: "user1", + ChatID: "pico:test-session", + Content: "hello", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "final answer" { + t.Fatalf("processMessage() response = %q, want final answer", response) + } + if got := provider.lastOptions["thinking_level"]; got != "off" { + t.Fatalf("thinking_level option = %#v, want off after hook model rewrite", got) + } + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("expected no reasoning outbound after hook rewrote to off model, got %+v", outbound) + case <-time.After(50 * time.Millisecond): + } +} + +func TestProcessMessage_BeforeLLMModelRewriteDoesNotLeakThinkingOff(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "off-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "off-model", + Model: "openai/off-model", + ThinkingLevel: "off", + }, + { + ModelName: "plain-model", + Model: "openai/plain-model", + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &reasoningOptionRecordingProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + if err := al.MountHook(NamedHook("rewrite-model", modelRewriteHook{model: "plain-model"})); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "pico", + SenderID: "user1", + ChatID: "pico:test-session", + Content: "hello", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "final answer" { + t.Fatalf("processMessage() response = %q, want final answer", response) + } + if _, ok := provider.lastOptions["thinking_level"]; ok { + t.Fatalf( + "thinking_level option should be cleared after hook rewrote away from off model, got %#v", + provider.lastOptions["thinking_level"], + ) + } + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "thinking trace" { + t.Fatalf("reasoning outbound content = %q, want thinking trace", outbound.Content) + } + case <-time.After(3 * time.Second): + t.Fatal("expected reasoning outbound after hook rewrote away from off model") + } +} + +func TestProcessMessage_BtwCommandSuppressesReasoningWhenThinkingOff(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "openai/test-model", + ThinkingLevel: "off", + }}, + } + + al := NewAgentLoop(cfg, bus.NewMessageBus(), &sideQuestionFallbackTestProvider{model: "test-model"}) + al.providerFactory = func(mc *config.ModelConfig) (providers.LLMProvider, string, error) { + model := "" + if mc != nil { + _, model = providers.ExtractProtocol(mc) + } + if model == "" { + model = "test-model" + } + return &sideQuestionFallbackTestProvider{model: model}, model, nil + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "/btw explain privately", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if strings.Contains(response, "thinking trace") { + t.Fatalf("processMessage() response = %q, should not expose reasoning with thinking off", response) + } +} + +func TestProcessMessage_BtwHookModelRewriteReevaluatesThinkingOff(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "plain-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "plain-model", + Model: "openai/plain-model", + }, + { + ModelName: "off-model", + Model: "openai/off-model", + ThinkingLevel: "off", + }, + }, + } + + al := NewAgentLoop(cfg, bus.NewMessageBus(), &sideQuestionFallbackTestProvider{model: "plain-model"}) + al.providerFactory = func(mc *config.ModelConfig) (providers.LLMProvider, string, error) { + model := "" + if mc != nil { + _, model = providers.ExtractProtocol(mc) + } + if model == "" { + model = "plain-model" + } + return &sideQuestionFallbackTestProvider{model: model}, model, nil + } + if err := al.MountHook(NamedHook("rewrite-model", modelRewriteHook{model: "off-model"})); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "/btw explain privately", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if strings.Contains(response, "thinking trace") { + t.Fatalf( + "processMessage() response = %q, should not expose reasoning after hook rewrote to off model", + response, + ) + } +} + +func TestProcessMessage_BtwHookModelRewriteDoesNotLeakThinkingOff(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "off-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "off-model", + Model: "openai/off-model", + ThinkingLevel: "off", + }, + { + ModelName: "plain-model", + Model: "openai/plain-model", + }, + }, + } + + al := NewAgentLoop(cfg, bus.NewMessageBus(), &sideQuestionFallbackTestProvider{model: "off-model"}) + al.providerFactory = func(mc *config.ModelConfig) (providers.LLMProvider, string, error) { + model := "" + if mc != nil { + _, model = providers.ExtractProtocol(mc) + } + if model == "" { + model = "off-model" + } + return &sideQuestionFallbackTestProvider{model: model}, model, nil + } + if err := al.MountHook(NamedHook("rewrite-model", modelRewriteHook{model: "plain-model"})); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "/btw explain privately", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "thinking trace" { + t.Fatalf("processMessage() response = %q, want reasoning after hook rewrote away from off model", response) + } +} + +func TestProcessMessage_BtwFallbackDoesNotInheritPrimaryThinkingOff(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + ModelFallbacks: []string{"openai/fallback-model"}, + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "openai/test-model", + ThinkingLevel: "off", + }}, + } + + al := NewAgentLoop(cfg, bus.NewMessageBus(), &sideQuestionFallbackTestProvider{model: "test-model"}) + al.providerFactory = func(mc *config.ModelConfig) (providers.LLMProvider, string, error) { + model := "" + if mc != nil { + _, model = providers.ExtractProtocol(mc) + } + if model == "" { + model = "test-model" + } + return &sideQuestionFallbackTestProvider{model: model}, model, nil + } + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "/btw explain fallback reasoning", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "thinking trace" { + t.Fatalf("processMessage() response = %q, want fallback reasoning when fallback has no off", response) + } +} + func TestProcessMessage_UseCommandLoadsRequestedSkill(t *testing.T) { tmpDir := t.TempDir() skillDir := filepath.Join(tmpDir, "skills", "shell") @@ -440,6 +1099,8 @@ func TestProcessMessage_BtwCommandRunsWithoutPersistingHistory(t *testing.T) { defaultAgent.Sessions.SetHistory(sessionKey, initialHistory) defaultAgent.Sessions.SetSummary(sessionKey, "The team decided to keep state request-scoped.") + initialHistory = defaultAgent.Sessions.GetHistory(sessionKey) + response, err := al.processMessage(context.Background(), msg) if err != nil { t.Fatalf("processMessage() error = %v", err) @@ -558,6 +1219,8 @@ func TestProcessMessage_BtwCommandUsesIsolatedProvider(t *testing.T) { } defaultAgent.Sessions.SetHistory(mainSessionKey, initialHistory) + initialHistory = defaultAgent.Sessions.GetHistory(mainSessionKey) + // Process a /btw command response, err := al.processMessage(context.Background(), bus.InboundMessage{ Channel: "telegram", @@ -3129,6 +3792,300 @@ func TestProcessMessage_FallbackUsesPerCandidateProvider(t *testing.T) { } } +func TestProcessMessage_FallbackReceivesExplicitThinkingOff(t *testing.T) { + workspace := t.TempDir() + + primaryServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _ = json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "rate limit exceeded", + "type": "rate_limit_error", + }, + }) + })) + defer primaryServer.Close() + + fallbackCalls := 0 + fallbackServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fallbackCalls++ + if r.URL.Path != "/chat/completions" { + t.Fatalf("fallback server path = %q, want /chat/completions", r.URL.Path) + } + defer r.Body.Close() + + var req map[string]any + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Fatalf("decode fallback request: %v", err) + } + if got := req["model"]; got != "doubao-seed-1-6-flash-250828" { + t.Fatalf("fallback request model = %#v, want doubao-seed-1-6-flash-250828", got) + } + thinking, ok := req["thinking"].(map[string]any) + if !ok { + t.Fatalf("fallback request thinking = %#v, want map", req["thinking"]) + } + if got := thinking["type"]; got != "disabled" { + t.Fatalf("fallback request thinking.type = %#v, want disabled", got) + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "fallback reply"}, + "finish_reason": "stop", + }, + }, + }); err != nil { + t.Fatalf("encode fallback response: %v", err) + } + })) + defer fallbackServer.Close() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "primary-model", + ModelFallbacks: []string{"doubao-fallback"}, + MaxTokens: 4096, + MaxToolIterations: 3, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "primary-model", + Model: "openrouter/primary-model", + APIBase: primaryServer.URL, + APIKeys: config.SimpleSecureStrings("primary-key"), + Workspace: workspace, + }, + { + ModelName: "doubao-fallback", + Model: "openai/doubao-seed-1-6-flash-250828", + APIBase: fallbackServer.URL, + APIKeys: config.SimpleSecureStrings("fallback-key"), + ThinkingLevel: "off", + Workspace: workspace, + }, + }, + } + + provider, _, err := providers.CreateProvider(cfg) + if err != nil { + t.Fatalf("CreateProvider() error = %v", err) + } + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + helper := testHelper{al: al} + + resp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "hi", + }) + + if resp != "fallback reply" { + t.Fatalf("response = %q, want fallback reply", resp) + } + if fallbackCalls != 1 { + t.Fatalf("fallback server calls = %d, want 1", fallbackCalls) + } +} + +func TestProcessMessage_PrimaryThinkingOffDoesNotLeakToFallback(t *testing.T) { + workspace := t.TempDir() + + primaryServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _ = json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "rate limit exceeded", + "type": "rate_limit_error", + }, + }) + })) + defer primaryServer.Close() + + fallbackServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/chat/completions" { + t.Fatalf("fallback server path = %q, want /chat/completions", r.URL.Path) + } + defer r.Body.Close() + + var req map[string]any + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Fatalf("decode fallback request: %v", err) + } + if _, ok := req["thinking"]; ok { + t.Fatalf("fallback request should not inherit primary thinking off, got thinking=%#v", req["thinking"]) + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "fallback reply"}, + "finish_reason": "stop", + }, + }, + }); err != nil { + t.Fatalf("encode fallback response: %v", err) + } + })) + defer fallbackServer.Close() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "primary-model", + ModelFallbacks: []string{"doubao-fallback"}, + MaxTokens: 4096, + MaxToolIterations: 3, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "primary-model", + Model: "openrouter/primary-model", + APIBase: primaryServer.URL, + APIKeys: config.SimpleSecureStrings("primary-key"), + ThinkingLevel: "off", + Workspace: workspace, + }, + { + ModelName: "doubao-fallback", + Model: "openai/doubao-seed-1-6-flash-250828", + APIBase: fallbackServer.URL, + APIKeys: config.SimpleSecureStrings("fallback-key"), + Workspace: workspace, + }, + }, + } + + provider, _, err := providers.CreateProvider(cfg) + if err != nil { + t.Fatalf("CreateProvider() error = %v", err) + } + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + helper := testHelper{al: al} + + resp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "hi", + }) + if resp != "fallback reply" { + t.Fatalf("response = %q, want fallback reply", resp) + } +} + +func TestProcessMessage_FallbackThinkingOffUsesCandidateIdentity(t *testing.T) { + workspace := t.TempDir() + + primaryServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _ = json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "rate limit exceeded", + "type": "rate_limit_error", + }, + }) + })) + defer primaryServer.Close() + + fallbackServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/chat/completions" { + t.Fatalf("fallback server path = %q, want /chat/completions", r.URL.Path) + } + defer r.Body.Close() + + var req map[string]any + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Fatalf("decode fallback request: %v", err) + } + thinking, ok := req["thinking"].(map[string]any) + if !ok { + t.Fatalf("fallback request thinking = %#v, want map", req["thinking"]) + } + if got := thinking["type"]; got != "disabled" { + t.Fatalf("fallback request thinking.type = %#v, want disabled", got) + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "fallback reply"}, + "finish_reason": "stop", + }, + }, + }); err != nil { + t.Fatalf("encode fallback response: %v", err) + } + })) + defer fallbackServer.Close() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "primary-model", + ModelFallbacks: []string{"doubao-off"}, + MaxTokens: 4096, + MaxToolIterations: 3, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "primary-model", + Model: "openrouter/primary-model", + APIBase: primaryServer.URL, + APIKeys: config.SimpleSecureStrings("primary-key"), + Workspace: workspace, + }, + { + ModelName: "doubao-default", + Model: "openai/doubao-seed-1-6-flash-250828", + APIBase: fallbackServer.URL, + APIKeys: config.SimpleSecureStrings("fallback-key"), + Workspace: workspace, + }, + { + ModelName: "doubao-off", + Model: "openai/doubao-seed-1-6-flash-250828", + APIBase: fallbackServer.URL, + APIKeys: config.SimpleSecureStrings("fallback-key"), + ThinkingLevel: "off", + Workspace: workspace, + }, + }, + } + + provider, _, err := providers.CreateProvider(cfg) + if err != nil { + t.Fatalf("CreateProvider() error = %v", err) + } + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + helper := testHelper{al: al} + + resp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "hi", + }) + if resp != "fallback reply" { + t.Fatalf("response = %q, want fallback reply", resp) + } +} + // TestProcessMessage_FallbackUsesActiveProviderWhenCandidateNotRegistered verifies // that when a candidate has no model_list entry it is absent from CandidateProviders // and the fallback closure falls back to activeProvider instead of panicking. diff --git a/pkg/agent/agent_utils.go b/pkg/agent/agent_utils.go index 9228b6d5..effa2b19 100644 --- a/pkg/agent/agent_utils.go +++ b/pkg/agent/agent_utils.go @@ -86,15 +86,56 @@ func outboundMessageForTurn(ts *turnState, content string) bus.OutboundMessage { } } -func outboundMessageForTurnWithKind(ts *turnState, content, kind string) bus.OutboundMessage { - msg := outboundMessageForTurn(ts, content) - if strings.TrimSpace(kind) == "" { - return msg +func markFinalOutbound(msg *bus.OutboundMessage) { + if msg == nil { + return } if msg.Context.Raw == nil { msg.Context.Raw = make(map[string]string, 1) } - msg.Context.Raw[metadataKeyMessageKind] = kind + msg.Context.Raw[metadataKeyOutboundKind] = outboundKindFinal +} + +type outboundTurnMessageOptions struct { + kind string + modelName string + raw map[string]string +} + +func outboundMessageForTurnWithOptions( + ts *turnState, + content string, + opts outboundTurnMessageOptions, +) bus.OutboundMessage { + msg := outboundMessageForTurn(ts, content) + trimmedKind := strings.TrimSpace(opts.kind) + trimmedModelName := strings.TrimSpace(opts.modelName) + rawCount := len(opts.raw) + if trimmedKind != "" { + rawCount++ + } + if trimmedModelName != "" { + rawCount++ + } + if rawCount == 0 { + return msg + } + + if msg.Context.Raw == nil { + msg.Context.Raw = make(map[string]string, rawCount) + } + if trimmedKind != "" { + msg.Context.Raw[metadataKeyMessageKind] = trimmedKind + } + if trimmedModelName != "" { + msg.Context.Raw["model_name"] = trimmedModelName + } + for key, value := range opts.raw { + if strings.TrimSpace(key) == "" { + continue + } + msg.Context.Raw[key] = value + } return msg } @@ -440,6 +481,9 @@ func activeSkillNames(agent *AgentInstance, opts processOptions) []string { if agent == nil { return nil } + if turnProfileSkillsOff(opts.TurnProfile) { + return nil + } combined := make([]string, 0, len(agent.SkillsFilter)+len(opts.ForcedSkills)) combined = append(combined, agent.SkillsFilter...) @@ -468,6 +512,9 @@ func activeSkillNames(agent *AgentInstance, opts processOptions) []string { resolved = append(resolved, name) } + if turnProfileCustomSkills(opts.TurnProfile) { + return filterNamesByTurnProfile(resolved, opts.TurnProfile.AllowedSkills) + } return resolved } @@ -511,13 +558,15 @@ func hasMediaRefs(messages []providers.Message) bool { func sideQuestionModelName(agent *AgentInstance, usedLight bool) string { if usedLight && len(agent.LightCandidates) > 0 { - // Use the first light candidate's model - return agent.LightCandidates[0].Model + if name := resolvedCandidateModelName(agent.LightCandidates, ""); name != "" { + return name + } } return agent.Model } func modelNameFromIdentityKey(identityKey string) string { + identityKey = strings.TrimSpace(identityKey) if identityKey == "" { return "" } @@ -528,6 +577,14 @@ func modelNameFromIdentityKey(identityKey string) string { return identityKey } +func modelAliasFromCandidateIdentityKey(identityKey string) string { + const prefix = "model_name:" + if !strings.HasPrefix(identityKey, prefix) { + return "" + } + return strings.TrimSpace(strings.TrimPrefix(identityKey, prefix)) +} + func closeProviderIfStateful(provider providers.LLMProvider) { if stateful, ok := provider.(providers.StatefulProvider); ok { stateful.Close() diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 87bdd6b4..01e9ec0b 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -54,9 +54,13 @@ func (cb *ContextBuilder) WithToolDiscovery(useBM25, useRegex bool) *ContextBuil useBM25: useBM25, useRegex: useRegex, }); err != nil { - logger.WarnCF("agent", "Failed to register tool discovery prompt contributor", map[string]any{ - "error": err.Error(), - }) + logger.WarnCF( + "agent", + "Failed to register tool discovery prompt contributor", + map[string]any{ + "error": err.Error(), + }, + ) } } return cb @@ -77,9 +81,13 @@ func (cb *ContextBuilder) WithAgentDiscovery( agentID: agentID, discover: discover, }); err != nil { - logger.WarnCF("agent", "Failed to register agent discovery prompt contributor", map[string]any{ - "error": err.Error(), - }) + logger.WarnCF( + "agent", + "Failed to register agent discovery prompt contributor", + map[string]any{ + "error": err.Error(), + }, + ) } } return cb @@ -130,9 +138,34 @@ func (cb *ContextBuilder) promptRegistryOrDefault() *PromptRegistry { return cb.promptRegistry } -func (cb *ContextBuilder) getIdentity() string { +func (cb *ContextBuilder) getIdentity(includeToolUseRule bool) string { workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace)) version := config.FormatVersion() + rules := []string{} + if includeToolUseRule { + rules = append(rules, toolUseSystemPromptRule()) + } + accuracyRule := "**Be helpful and accurate** - Briefly explain what you're doing." + if includeToolUseRule { + accuracyRule = "**Be helpful and accurate** - When using tools, briefly explain what you're doing." + } + rules = append( + rules, + accuracyRule, + "**Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content.", + ) + if includeToolUseRule { + rules = append( + rules, + fmt.Sprintf( + "**Memory** - When interacting with me if something seems memorable, update %s/memory/MEMORY.md", + workspacePath, + ), + ) + } + for i, rule := range rules { + rules[i] = fmt.Sprintf("%d. %s", i+1, rule) + } return fmt.Sprintf( `# picoclaw 🦞 (%s) @@ -147,14 +180,15 @@ Your workspace is at: %s ## Important Rules -1. **ALWAYS use tools** - When you need to perform an action (schedule reminders, send messages, execute commands, etc.), you MUST call the appropriate tool. Do NOT just say you'll do it or pretend to do it. - -2. **Be helpful and accurate** - When using tools, briefly explain what you're doing. - -3. **Memory** - When interacting with me if something seems memorable, update %s/memory/MEMORY.md - -4. **Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content.`, - version, workspacePath, workspacePath, workspacePath, workspacePath, workspacePath) +%s +`, + version, + workspacePath, + workspacePath, + workspacePath, + workspacePath, + strings.Join(rules, "\n\n"), + ) } func formatToolDiscoveryRule(useBM25, useRegex bool) string { @@ -181,6 +215,20 @@ func (cb *ContextBuilder) BuildSystemPrompt() string { } func (cb *ContextBuilder) BuildSystemPromptParts() []PromptPart { + return cb.buildSystemPromptParts(systemPromptBuildOptions{ + IncludeSkillCatalog: true, + IncludeToolUseRule: true, + }) +} + +type systemPromptBuildOptions struct { + IncludeSkillCatalog bool + IncludeToolUseRule bool + AllowedSkills []string + AllowedTools []string +} + +func (cb *ContextBuilder) buildSystemPromptParts(opts systemPromptBuildOptions) []PromptPart { stack := NewPromptStack(cb.promptRegistryOrDefault()) add := func(part PromptPart) { if err := stack.Add(part); err != nil { @@ -201,7 +249,7 @@ func (cb *ContextBuilder) BuildSystemPromptParts() []PromptPart { Slot: PromptSlotIdentity, Source: PromptSource{ID: PromptSourceKernel, Name: "identity"}, Title: "picoclaw identity", - Content: cb.getIdentity(), + Content: cb.getIdentity(opts.IncludeToolUseRule), Stable: true, Cache: PromptCacheEphemeral, }) @@ -222,8 +270,19 @@ func (cb *ContextBuilder) BuildSystemPromptParts() []PromptPart { } // Skills - show summary, AI can read full content with read_file tool - skillsSummary := cb.skillsLoader.BuildSkillsSummary() + skillsSummary := "" + if opts.IncludeSkillCatalog { + skillsSummary = cb.buildSkillsSummary(opts.AllowedSkills) + } if skillsSummary != "" { + skillIntro := "The following skills extend your capabilities." + readFileAllowed := promptAllowsTool( + PromptBuildRequest{AllowedTools: opts.AllowedTools}, + "read_file", + ) + if opts.IncludeToolUseRule && readFileAllowed { + skillIntro += " To use a skill, read its SKILL.md file using the read_file tool." + } add(PromptPart{ ID: "capability.skill_catalog", Layer: PromptLayerCapability, @@ -232,9 +291,9 @@ func (cb *ContextBuilder) BuildSystemPromptParts() []PromptPart { Title: "skill catalog", Content: fmt.Sprintf(`# Skills -The following skills extend your capabilities. To use a skill, read its SKILL.md file using the read_file tool. +%s -%s`, skillsSummary), +%s`, skillIntro, skillsSummary), Stable: true, Cache: PromptCacheEphemeral, }) @@ -319,6 +378,96 @@ func (cb *ContextBuilder) BuildSystemPromptWithCache() string { return prompt } +func (cb *ContextBuilder) buildSystemPromptForRequest( + req PromptBuildRequest, +) (string, []providers.ContentBlock) { + if req.SuppressDefaultSystemPrompt { + return "", nil + } + + useDefaultCache := !req.SuppressSkillContext && + !req.SuppressToolUseRule && + len(req.AllowedSkills) == 0 && + len(req.AllowedTools) == 0 + if useDefaultCache { + staticPrompt := cb.BuildSystemPromptWithCache() + return staticPrompt, []providers.ContentBlock{ + promptContentBlock(PromptPart{ + ID: "kernel.static", + Layer: PromptLayerKernel, + Slot: PromptSlotIdentity, + Source: PromptSource{ID: PromptSourceKernel, Name: "static"}, + Content: staticPrompt, + }, &providers.CacheControl{Type: "ephemeral"}), + } + } + + parts := cb.buildSystemPromptParts(systemPromptBuildOptions{ + IncludeSkillCatalog: !req.SuppressSkillContext, + IncludeToolUseRule: !req.SuppressToolUseRule, + AllowedSkills: req.AllowedSkills, + AllowedTools: req.AllowedTools, + }) + staticPrompt := renderPromptPartsLegacy(parts) + blocks := make([]providers.ContentBlock, 0, len(parts)) + for _, part := range parts { + if strings.TrimSpace(part.Content) == "" { + continue + } + blocks = append(blocks, promptContentBlock(part, cacheControlForPromptPart(part))) + } + return staticPrompt, blocks +} + +func (cb *ContextBuilder) buildSkillsSummary(allowed []string) string { + if cb.skillsLoader == nil { + return "" + } + if len(allowed) == 0 { + return cb.skillsLoader.BuildSkillsSummary() + } + allowedSet := cleanAllowedSet(allowed) + if len(allowedSet) == 0 { + return "" + } + + var lines []string + lines = append(lines, "") + for _, s := range cb.skillsLoader.ListSkills() { + if _, ok := allowedSet[strings.ToLower(strings.TrimSpace(s.Name))]; !ok { + continue + } + lines = append(lines, " ") + lines = append(lines, fmt.Sprintf(" %s", xmlEscapeForPrompt(s.Name))) + lines = append( + lines, + fmt.Sprintf(" %s", xmlEscapeForPrompt(s.Description)), + ) + lines = append( + lines, + fmt.Sprintf(" %s", xmlEscapeForPrompt(s.Path)), + ) + lines = append(lines, fmt.Sprintf(" %s", xmlEscapeForPrompt(s.Source))) + lines = append(lines, " ") + } + if len(lines) == 1 { + return "" + } + lines = append(lines, "") + return strings.Join(lines, "\n") +} + +func xmlEscapeForPrompt(s string) string { + replacer := strings.NewReplacer( + "&", "&", + "<", "<", + ">", ">", + "\"", """, + "'", "'", + ) + return replacer.Replace(s) +} + // EstimateSystemTokens estimates the token count of the full system message // that would be sent to the LLM, mirroring the composition logic in BuildMessages. // It includes: static prompt, dynamic context, active skills, and summary with @@ -687,19 +836,18 @@ func (cb *ContextBuilder) BuildMessages( func (cb *ContextBuilder) BuildMessagesFromPrompt(req PromptBuildRequest) []providers.Message { messages := []providers.Message{} - // The static part (identity, bootstrap, skills, memory) is cached locally to - // avoid repeated file I/O and string building on every call (fixes issue #607). - // Dynamic parts (time, session, summary) are appended per request. + // The default static part (identity, bootstrap, skills, memory) is cached + // locally to avoid repeated file I/O and string building on every call + // (fixes issue #607). Profile-customized static prompts are built on demand. + // Dynamic parts (time, session, summary) are appended per request unless the + // profile suppresses PicoClaw system context. // Everything is sent as a single system message for provider compatibility: // - Anthropic adapter extracts messages[0] (Role=="system") and maps its content // to the top-level "system" parameter in the Messages API request. A single // contiguous system block makes this extraction straightforward. // - Codex maps only the first system message to its instructions field. // - OpenAI-compat passes messages through as-is. - staticPrompt := cb.BuildSystemPromptWithCache() - - // Build short dynamic context (time, runtime, session) — changes per request - dynamicCtx := cb.buildDynamicContext(req.Channel, req.ChatID, req.SenderID, req.SenderDisplayName) + staticPrompt, contentBlocks := cb.buildSystemPromptForRequest(req) // Compose a single system message: static (cached) + dynamic + optional summary. // Keeping all system content in one message ensures every provider adapter can @@ -710,26 +858,27 @@ func (cb *ContextBuilder) BuildMessagesFromPrompt(req PromptBuildRequest) []prov // cache-aware adapters (Anthropic) can set per-block cache_control. // The static block is marked "ephemeral" — its prefix hash is stable // across requests, enabling LLM-side KV cache reuse. - stringParts := []string{staticPrompt} - - contentBlocks := []providers.ContentBlock{ - promptContentBlock(PromptPart{ - ID: "kernel.static", - Layer: PromptLayerKernel, - Slot: PromptSlotIdentity, - Source: PromptSource{ID: PromptSourceKernel, Name: "static"}, - Content: staticPrompt, - }, &providers.CacheControl{Type: "ephemeral"}), + var stringParts []string + if strings.TrimSpace(staticPrompt) != "" { + stringParts = append(stringParts, staticPrompt) } promptParts := append([]PromptPart(nil), req.Overlays...) - promptParts = append(promptParts, cb.buildActiveSkillsPromptParts(req.ActiveSkills)...) - if contributedParts, err := cb.promptRegistryOrDefault().Collect(context.Background(), req); err != nil { - logger.WarnCF("agent", "Prompt contributor collection failed", map[string]any{ - "error": err.Error(), - }) - } else { - promptParts = append(promptParts, contributedParts...) + if !req.SuppressDefaultSystemPrompt && !req.SuppressSkillContext { + activeSkills := append([]string(nil), req.ActiveSkills...) + if len(req.AllowedSkills) > 0 { + activeSkills = filterNamesByTurnProfile(activeSkills, req.AllowedSkills) + } + promptParts = append(promptParts, cb.buildActiveSkillsPromptParts(activeSkills)...) + } + if !req.SuppressDefaultSystemPrompt { + if contributedParts, err := cb.promptRegistryOrDefault().Collect(context.Background(), req); err != nil { + logger.WarnCF("agent", "Prompt contributor collection failed", map[string]any{ + "error": err.Error(), + }) + } else { + promptParts = append(promptParts, contributedParts...) + } } if len(promptParts) > 0 { @@ -752,35 +901,62 @@ func (cb *ContextBuilder) BuildMessagesFromPrompt(req PromptBuildRequest) []prov } } - runtimePart := PromptPart{ - ID: "context.runtime", - Layer: PromptLayerContext, - Slot: PromptSlotRuntime, - Source: PromptSource{ID: PromptSourceRuntime, Name: "runtime"}, - Title: "runtime context", - Content: dynamicCtx, - Stable: false, - Cache: PromptCacheNone, - } - stringParts = append(stringParts, dynamicCtx) - contentBlocks = append(contentBlocks, promptContentBlock(runtimePart, nil)) - - if req.Summary != "" { - summaryPart := PromptPart{ - ID: "context.summary", - Layer: PromptLayerContext, - Slot: PromptSlotSummary, - Source: PromptSource{ID: PromptSourceSummary, Name: "context.summary"}, - Title: "context summary", - Content: fmt.Sprintf( - "CONTEXT_SUMMARY: The following is an approximate summary of prior conversation "+ - "for reference only. It may be incomplete or outdated — always defer to explicit instructions.\n\n%s", - req.Summary), - Stable: false, - Cache: PromptCacheNone, + dynamicChars := 0 + if !req.SuppressDefaultSystemPrompt { + // Build short dynamic context (time, runtime, session) — changes per request + dynamicCtx := cb.buildDynamicContext( + req.Channel, + req.ChatID, + req.SenderID, + req.SenderDisplayName, + ) + dynamicChars = len(dynamicCtx) + runtimePart := PromptPart{ + ID: "context.runtime", + Layer: PromptLayerContext, + Slot: PromptSlotRuntime, + Source: PromptSource{ID: PromptSourceRuntime, Name: "runtime"}, + Title: "runtime context", + Content: dynamicCtx, + Stable: false, + Cache: PromptCacheNone, } - stringParts = append(stringParts, summaryPart.Content) - contentBlocks = append(contentBlocks, promptContentBlock(summaryPart, nil)) + stringParts = append(stringParts, dynamicCtx) + contentBlocks = append(contentBlocks, promptContentBlock(runtimePart, nil)) + + if req.Summary != "" { + summaryPart := PromptPart{ + ID: "context.summary", + Layer: PromptLayerContext, + Slot: PromptSlotSummary, + Source: PromptSource{ID: PromptSourceSummary, Name: "context.summary"}, + Title: "context summary", + Content: fmt.Sprintf( + "CONTEXT_SUMMARY: The following is an approximate summary of prior conversation "+ + "for reference only. It may be incomplete or outdated — always defer to explicit instructions.\n\n%s", + req.Summary, + ), + Stable: false, + Cache: PromptCacheNone, + } + stringParts = append(stringParts, summaryPart.Content) + contentBlocks = append(contentBlocks, promptContentBlock(summaryPart, nil)) + } + } + + if len(stringParts) == 0 && req.ToolUseFallback { + fallbackPart := PromptPart{ + ID: "kernel.tool_use_fallback", + Layer: PromptLayerKernel, + Slot: PromptSlotIdentity, + Source: PromptSource{ID: PromptSourceKernel, Name: "tool_use_fallback"}, + Title: "tool use fallback", + Content: toolUseSystemPromptRule(), + Stable: true, + Cache: PromptCacheEphemeral, + } + stringParts = append(stringParts, fallbackPart.Content) + contentBlocks = append(contentBlocks, promptContentBlock(fallbackPart, nil)) } fullSystemPrompt := strings.Join(stringParts, "\n\n---\n\n") @@ -795,7 +971,7 @@ func (cb *ContextBuilder) BuildMessagesFromPrompt(req PromptBuildRequest) []prov logger.DebugCF("agent", "System prompt built", map[string]any{ "static_chars": len(staticPrompt), - "dynamic_chars": len(dynamicCtx), + "dynamic_chars": dynamicChars, "total_chars": len(fullSystemPrompt), "has_summary": req.Summary != "", "overlays": len(req.Overlays), @@ -814,11 +990,13 @@ func (cb *ContextBuilder) BuildMessagesFromPrompt(req PromptBuildRequest) []prov // Single system message containing all context — compatible with all providers. // SystemParts enables cache-aware adapters to set per-block cache_control; // Content is the concatenated fallback for adapters that don't read SystemParts. - messages = append(messages, providers.Message{ - Role: "system", - Content: fullSystemPrompt, - SystemParts: contentBlocks, - }) + if strings.TrimSpace(fullSystemPrompt) != "" { + messages = append(messages, providers.Message{ + Role: "system", + Content: fullSystemPrompt, + SystemParts: contentBlocks, + }) + } // Add conversation history messages = append(messages, history...) @@ -829,6 +1007,9 @@ func (cb *ContextBuilder) BuildMessagesFromPrompt(req PromptBuildRequest) []prov if strings.TrimSpace(req.CurrentMessage) != "" || len(req.Media) > 0 { messages = append(messages, userPromptMessage(req.CurrentMessage, req.Media)) } + if len(messages) == 0 { + messages = append(messages, userPromptMessage("", nil)) + } return messages } @@ -933,7 +1114,11 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message break } if next.ToolCallID == "" { - logger.DebugCF("agent", "Dropping tool result without tool_call_id", map[string]any{}) + logger.DebugCF( + "agent", + "Dropping tool result without tool_call_id", + map[string]any{}, + ) continue } if _, ok := expected[next.ToolCallID]; !ok { @@ -943,9 +1128,13 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message continue } if seenInBlock[next.ToolCallID] { - logger.DebugCF("agent", "Dropping duplicate tool result in tool block", map[string]any{ - "tool_call_id": next.ToolCallID, - }) + logger.DebugCF( + "agent", + "Dropping duplicate tool result in tool block", + map[string]any{ + "tool_call_id": next.ToolCallID, + }, + ) continue } seenInBlock[next.ToolCallID] = true @@ -955,7 +1144,11 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message allFound := !invalidToolCallID if invalidToolCallID { - logger.DebugCF("agent", "Dropping assistant message with empty tool_call_id", map[string]any{}) + logger.DebugCF( + "agent", + "Dropping assistant message with empty tool_call_id", + map[string]any{}, + ) } for toolCallID, found := range expected { if !found { @@ -985,9 +1178,13 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message } if msg.Role == "tool" { - logger.DebugCF("agent", "Dropping orphaned tool message after validation", map[string]any{ - "tool_call_id": msg.ToolCallID, - }) + logger.DebugCF( + "agent", + "Dropping orphaned tool message after validation", + map[string]any{ + "tool_call_id": msg.ToolCallID, + }, + ) continue } diff --git a/pkg/agent/context_seahorse.go b/pkg/agent/context_seahorse.go index c6e5b30a..2a10d245 100644 --- a/pkg/agent/context_seahorse.go +++ b/pkg/agent/context_seahorse.go @@ -197,6 +197,7 @@ func providerToSeahorseMessage(msg protocoltypes.Message) seahorse.Message { result := seahorse.Message{ Role: msg.Role, Content: msg.Content, + ModelName: msg.ModelName, ReasoningContent: msg.ReasoningContent, TokenCount: tokenizer.EstimateMessageTokens(msg), } @@ -243,6 +244,7 @@ func seahorseToProviderMessages(result *seahorse.AssembleResult) []protocoltypes pm := protocoltypes.Message{ Role: msg.Role, Content: msg.Content, + ModelName: msg.ModelName, ReasoningContent: msg.ReasoningContent, } diff --git a/pkg/agent/context_seahorse_test.go b/pkg/agent/context_seahorse_test.go index 05c83183..101f72ee 100644 --- a/pkg/agent/context_seahorse_test.go +++ b/pkg/agent/context_seahorse_test.go @@ -174,6 +174,7 @@ func TestProviderToSeahorseMessageWithReasoning(t *testing.T) { msg := protocoltypes.Message{ Role: "assistant", Content: "response text", + ModelName: "gpt-5.4-mini", ReasoningContent: "I thought about this carefully", } @@ -181,6 +182,9 @@ func TestProviderToSeahorseMessageWithReasoning(t *testing.T) { if result.ReasoningContent != "I thought about this carefully" { t.Errorf("ReasoningContent = %q, want 'I thought about this carefully'", result.ReasoningContent) } + if result.ModelName != "gpt-5.4-mini" { + t.Errorf("ModelName = %q, want %q", result.ModelName, "gpt-5.4-mini") + } } func TestSeahorseToProviderMessagesWithReasoning(t *testing.T) { @@ -189,6 +193,7 @@ func TestSeahorseToProviderMessagesWithReasoning(t *testing.T) { { Role: "assistant", Content: "response", + ModelName: "gpt-5.4", ReasoningContent: "thinking process", }, }, @@ -201,6 +206,9 @@ func TestSeahorseToProviderMessagesWithReasoning(t *testing.T) { if messages[0].ReasoningContent != "thinking process" { t.Errorf("ReasoningContent = %q, want 'thinking process'", messages[0].ReasoningContent) } + if messages[0].ModelName != "gpt-5.4" { + t.Errorf("ModelName = %q, want %q", messages[0].ModelName, "gpt-5.4") + } } func TestSeahorseToProviderMessages(t *testing.T) { diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 4ed71303..68150ce7 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -31,6 +31,7 @@ type AgentInstance struct { MaxTokens int Temperature float64 ThinkingLevel ThinkingLevel + ThinkingLevelConfigured bool ContextWindow int SummarizeMessageThreshold int SummarizeTokenPercent int @@ -184,6 +185,7 @@ func NewAgentInstance( thinkingLevelStr = mc.ThinkingLevel } thinkingLevel := parseThinkingLevel(thinkingLevelStr) + thinkingLevelConfigured := isConfiguredThinkingLevel(thinkingLevelStr) summarizeMessageThreshold := defaults.SummarizeMessageThreshold if summarizeMessageThreshold == 0 { @@ -251,6 +253,7 @@ func NewAgentInstance( MaxTokens: maxTokens, Temperature: temperature, ThinkingLevel: thinkingLevel, + ThinkingLevelConfigured: thinkingLevelConfigured, ContextWindow: contextWindow, SummarizeMessageThreshold: summarizeMessageThreshold, SummarizeTokenPercent: summarizeTokenPercent, diff --git a/pkg/agent/interfaces/interfaces.go b/pkg/agent/interfaces/interfaces.go index 2efec05e..87428261 100644 --- a/pkg/agent/interfaces/interfaces.go +++ b/pkg/agent/interfaces/interfaces.go @@ -21,6 +21,9 @@ type MessageBus interface { // PublishOutboundMedia sends an outbound media message. PublishOutboundMedia(ctx context.Context, msg bus.OutboundMediaMessage) error + // GetStreamer returns a channel streamer when the active channel supports streaming. + GetStreamer(ctx context.Context, channel, chatID, sessionKey string) (bus.Streamer, bool) + // InboundChan returns the channel for receiving inbound messages. InboundChan() <-chan bus.InboundMessage } diff --git a/pkg/agent/model_resolution.go b/pkg/agent/model_resolution.go index 6065f640..487a3072 100644 --- a/pkg/agent/model_resolution.go +++ b/pkg/agent/model_resolution.go @@ -29,6 +29,36 @@ func modelConfigIdentityKey(mc *config.ModelConfig) string { return "" } +func effectiveDefaultProvider(defaultProvider string) string { + defaultProvider = strings.TrimSpace(defaultProvider) + if defaultProvider == "" { + return "openai" + } + return providers.NormalizeProvider(defaultProvider) +} + +func modelProviderAndIDForResolution(defaultProvider string, mc *config.ModelConfig) (provider string, modelID string) { + if mc == nil { + return "", "" + } + return providers.ExtractProtocol(mc) +} + +func cloneModelConfigForResolution( + defaultProvider string, + mc *config.ModelConfig, + workspace string, +) *config.ModelConfig { + if mc == nil { + return nil + } + clone := *mc + if clone.Workspace == "" { + clone.Workspace = workspace + } + return &clone +} + func candidateFromModelConfig( defaultProvider string, mc *config.ModelConfig, @@ -37,7 +67,7 @@ func candidateFromModelConfig( return providers.FallbackCandidate{}, false } - protocol, modelID := providers.ExtractProtocol(mc) + protocol, modelID := modelProviderAndIDForResolution(defaultProvider, mc) if strings.TrimSpace(modelID) == "" { return providers.FallbackCandidate{}, false } @@ -45,12 +75,13 @@ func candidateFromModelConfig( return providers.FallbackCandidate{ Provider: protocol, Model: modelID, + DisplayName: strings.TrimSpace(mc.ModelName), RPM: mc.RPM, IdentityKey: modelConfigIdentityKey(mc), }, true } -func lookupModelConfigByRef(cfg *config.Config, raw string) *config.ModelConfig { +func lookupModelConfigByRef(cfg *config.Config, raw string, defaultProvider ...string) *config.ModelConfig { raw = strings.TrimSpace(raw) if raw == "" || cfg == nil { return nil @@ -66,6 +97,10 @@ func lookupModelConfigByRef(cfg *config.Config, raw string) *config.ModelConfig rawKey = providers.ModelKey(rawRef.Provider, rawRef.Model) } + fallbackProvider := "" + if len(defaultProvider) > 0 { + fallbackProvider = effectiveDefaultProvider(defaultProvider[0]) + } for i := range cfg.ModelList { mc := cfg.ModelList[i] if mc == nil { @@ -75,12 +110,14 @@ func lookupModelConfigByRef(cfg *config.Config, raw string) *config.ModelConfig if fullModel == "" { continue } + protocol, modelID := modelProviderAndIDForResolution(fallbackProvider, mc) if fullModel == raw { return mc } - protocol, modelID := providers.ExtractProtocol(mc) if modelID == raw { - return mc + if fallbackProvider == "" || providers.NormalizeProvider(protocol) == fallbackProvider { + return mc + } } if rawKey != "" && providers.ModelKey(protocol, modelID) == rawKey { return mc @@ -99,8 +136,9 @@ func resolveModelCandidate( if raw == "" { return providers.FallbackCandidate{}, false } + defaultProvider = effectiveDefaultProvider(defaultProvider) - if mc := lookupModelConfigByRef(cfg, raw); mc != nil { + if mc := lookupModelConfigByRef(cfg, raw, defaultProvider); mc != nil { return candidateFromModelConfig(defaultProvider, mc) } @@ -110,8 +148,9 @@ func resolveModelCandidate( } return providers.FallbackCandidate{ - Provider: ref.Provider, - Model: ref.Model, + Provider: ref.Provider, + Model: ref.Model, + DisplayName: raw, }, true } @@ -160,6 +199,18 @@ func resolvedCandidateProvider(candidates []providers.FallbackCandidate, fallbac return fallback } +func resolvedCandidateModelName(candidates []providers.FallbackCandidate, fallback string) string { + if len(candidates) > 0 { + if name := modelAliasFromCandidateIdentityKey(candidates[0].IdentityKey); strings.TrimSpace(name) != "" { + return name + } + if displayName := strings.TrimSpace(candidates[0].DisplayName); displayName != "" { + return displayName + } + } + return strings.TrimSpace(fallback) +} + func resolvedModelConfig(cfg *config.Config, modelName, workspace string) (*config.ModelConfig, error) { if cfg == nil { return nil, fmt.Errorf("config is nil") @@ -177,3 +228,48 @@ func resolvedModelConfig(cfg *config.Config, modelName, workspace string) (*conf return &clone, nil } + +func resolveActiveModelConfig( + cfg *config.Config, + workspace string, + candidates []providers.FallbackCandidate, + activeModel string, + defaultProvider string, +) *config.ModelConfig { + if cfg == nil { + return nil + } + defaultProvider = effectiveDefaultProvider(defaultProvider) + + if len(candidates) > 0 { + candidate := candidates[0] + identityKey := strings.TrimSpace(candidate.IdentityKey) + if identityKey != "" { + for _, mc := range cfg.ModelList { + if mc == nil || modelConfigIdentityKey(mc) != identityKey { + continue + } + protocol, modelID := modelProviderAndIDForResolution(defaultProvider, mc) + if providers.ModelKey(protocol, modelID) == providers.ModelKey(candidate.Provider, candidate.Model) { + return cloneModelConfigForResolution(defaultProvider, mc, workspace) + } + } + } + for _, mc := range cfg.ModelList { + if mc == nil { + continue + } + protocol, modelID := modelProviderAndIDForResolution(defaultProvider, mc) + if providers.ModelKey(protocol, modelID) == providers.ModelKey(candidate.Provider, candidate.Model) { + return cloneModelConfigForResolution(defaultProvider, mc, workspace) + } + } + return nil + } + + if mc := lookupModelConfigByRef(cfg, activeModel, defaultProvider); mc != nil { + return cloneModelConfigForResolution(defaultProvider, mc, workspace) + } + + return nil +} diff --git a/pkg/agent/model_resolution_test.go b/pkg/agent/model_resolution_test.go new file mode 100644 index 00000000..ea23e5da --- /dev/null +++ b/pkg/agent/model_resolution_test.go @@ -0,0 +1,165 @@ +package agent + +import ( + "testing" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" +) + +func TestModelNameFromIdentityKey_LegacyProviderModel(t *testing.T) { + if got := modelNameFromIdentityKey("openai/gpt-5.4"); got != "gpt-5.4" { + t.Fatalf("modelNameFromIdentityKey() = %q, want %q", got, "gpt-5.4") + } +} + +func TestModelNameFromIdentityKey_PreservesNonLegacyIdentity(t *testing.T) { + if got := modelNameFromIdentityKey("model_name:primary"); got != "model_name:primary" { + t.Fatalf("modelNameFromIdentityKey() = %q, want %q", got, "model_name:primary") + } +} + +func TestModelAliasFromCandidateIdentityKey(t *testing.T) { + if got := modelAliasFromCandidateIdentityKey("model_name:primary"); got != "primary" { + t.Fatalf("modelAliasFromCandidateIdentityKey() = %q, want %q", got, "primary") + } + if got := modelAliasFromCandidateIdentityKey("openai/gpt-5.4"); got != "" { + t.Fatalf("modelAliasFromCandidateIdentityKey() = %q, want empty", got) + } +} + +func TestResolvedCandidateModelName_PrefersIdentityAlias(t *testing.T) { + got := resolvedCandidateModelName([]providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4", IdentityKey: "model_name:primary"}, + }, "fallback-model") + if got != "primary" { + t.Fatalf("resolvedCandidateModelName() = %q, want %q", got, "primary") + } +} + +func TestResolvedCandidateModelName_DoesNotScanFallbackAliases(t *testing.T) { + got := resolvedCandidateModelName([]providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4"}, + {Provider: "openai", Model: "gpt-5.4-mini", IdentityKey: "model_name:fallback"}, + }, "primary-model") + if got != "primary-model" { + t.Fatalf("resolvedCandidateModelName() = %q, want %q", got, "primary-model") + } +} + +func TestResolvedCandidateModelName_UsesCandidateDisplayName(t *testing.T) { + got := resolvedCandidateModelName([]providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4", DisplayName: "gpt-5.4-display"}, + }, "fallback-model") + if got != "gpt-5.4-display" { + t.Fatalf("resolvedCandidateModelName() = %q, want %q", got, "gpt-5.4-display") + } +} + +func TestResolveActiveModelConfig_PrefersCandidateIdentityKey(t *testing.T) { + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + { + ModelName: "glm-4.7", + Provider: "zhipu", + Model: "glm-4.7", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + { + ModelName: "suanneng-glm-4.7", + Provider: "zhipu", + Model: "glm-4.7", + Streaming: config.ModelStreamingConfig{Enabled: true}, + }, + }, + } + + got := resolveActiveModelConfig( + cfg, + "/workspace", + []providers.FallbackCandidate{{ + Provider: "zhipu", + Model: "glm-4.7", + IdentityKey: "model_name:suanneng-glm-4.7", + }}, + "glm-4.7", + "openai", + ) + + if got == nil { + t.Fatal("resolveActiveModelConfig() = nil, want model config") + } + if got.ModelName != "suanneng-glm-4.7" { + t.Fatalf("model_name = %q, want %q", got.ModelName, "suanneng-glm-4.7") + } + if !got.Streaming.Enabled { + t.Fatal("streaming.enabled = false, want true from identity-matched model config") + } +} + +func TestResolveActiveModelConfig_LoadBalancedAliasUsesSelectedCandidate(t *testing.T) { + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + { + ModelName: "lb-model", + Model: "openai/primary", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + { + ModelName: "lb-model", + Model: "openai/secondary", + Streaming: config.ModelStreamingConfig{Enabled: true}, + }, + }, + } + + got := resolveActiveModelConfig( + cfg, + "/workspace", + []providers.FallbackCandidate{{ + Provider: "openai", + Model: "secondary", + IdentityKey: "model_name:lb-model", + }}, + "lb-model", + "openai", + ) + + if got == nil { + t.Fatal("resolveActiveModelConfig() = nil, want model config") + } + if got.Model != "openai/secondary" { + t.Fatalf("model = %q, want openai/secondary", got.Model) + } + if !got.Streaming.Enabled { + t.Fatal("streaming.enabled = false, want true from selected load-balanced entry") + } +} + +func TestResolveActiveModelConfig_DoesNotFallbackToOpenAIForDefaultProviderCandidate(t *testing.T) { + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + { + ModelName: "openai-gpt", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: true}, + }, + }, + } + + got := resolveActiveModelConfig( + cfg, + "/workspace", + []providers.FallbackCandidate{{ + Provider: "nvidia", + Model: "gpt-4o", + }}, + "gpt-4o", + "nvidia", + ) + + if got != nil { + t.Fatalf("resolveActiveModelConfig() = %#v, want nil for non-active provider config", got) + } +} diff --git a/pkg/agent/pipeline_execute.go b/pkg/agent/pipeline_execute.go index 567e56d1..81fc06ec 100644 --- a/pkg/agent/pipeline_execute.go +++ b/pkg/agent/pipeline_execute.go @@ -131,6 +131,36 @@ toolLoop: toolName := tc.Name toolArgs := cloneStringAnyMap(tc.Arguments) + denyByTurnProfile := func() bool { + if turnProfileToolAllowed(ts.profile, toolName) { + return false + } + exec.allResponsesHandled = false + denyContent := fmt.Sprintf("Tool %q is not allowed by the active turn profile.", toolName) + al.emitEvent( + runtimeevents.KindAgentToolExecSkipped, + ts.eventMeta("runTurn", "turn.tool.skipped"), + ToolExecSkippedPayload{ + Tool: toolName, + Reason: denyContent, + }, + ) + deniedMsg := providers.Message{ + Role: "tool", + Content: denyContent, + ToolCallID: tc.ID, + } + messages = append(messages, deniedMsg) + if !ts.opts.NoHistory { + ts.agent.Sessions.AddFullMessage(ts.sessionKey, deniedMsg) + ts.recordPersistedMessage(deniedMsg) + } + return true + } + + if denyByTurnProfile() { + continue + } if al.hooks != nil { toolReq, decision := al.hooks.BeforeTool(turnCtx, &ToolCallHookRequest{ @@ -180,7 +210,11 @@ toolLoop: toolFeedbackArgsPreview(toolArgs, toolFeedbackMaxLen), ) fbCtx, fbCancel := context.WithTimeout(turnCtx, 3*time.Second) - _ = al.bus.PublishOutbound(fbCtx, outboundMessageForTurnWithKind(ts, feedbackMsg, messageKindToolFeedback)) + _ = al.bus.PublishOutbound(fbCtx, outboundMessageForTurnWithOptions( + ts, + feedbackMsg, + outboundTurnMessageOptions{kind: messageKindToolFeedback}, + )) fbCancel() } @@ -359,7 +393,9 @@ toolLoop: content := al.cfg.FilterSensitiveData(result.ForLLM) msg := subTurnResultPromptMessage(content) messages = append(messages, msg) - ts.agent.Sessions.AddFullMessage(ts.sessionKey, msg) + if !ts.opts.NoHistory { + ts.agent.Sessions.AddFullMessage(ts.sessionKey, msg) + } } default: } @@ -437,6 +473,10 @@ toolLoop: } } + if denyByTurnProfile() { + continue + } + argsJSON, _ := json.Marshal(toolArgs) argsPreview := utils.Truncate(string(argsJSON), 200) logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", toolName, argsPreview), @@ -467,7 +507,11 @@ toolLoop: toolFeedbackArgsPreview(toolArgs, toolFeedbackMaxLen), ) fbCtx, fbCancel := context.WithTimeout(turnCtx, 3*time.Second) - _ = al.bus.PublishOutbound(fbCtx, outboundMessageForTurnWithKind(ts, feedbackMsg, messageKindToolFeedback)) + _ = al.bus.PublishOutbound(fbCtx, outboundMessageForTurnWithOptions( + ts, + feedbackMsg, + outboundTurnMessageOptions{kind: messageKindToolFeedback}, + )) fbCancel() } @@ -740,7 +784,9 @@ toolLoop: content := al.cfg.FilterSensitiveData(result.ForLLM) msg := subTurnResultPromptMessage(content) messages = append(messages, msg) - ts.agent.Sessions.AddFullMessage(ts.sessionKey, msg) + if !ts.opts.NoHistory { + ts.agent.Sessions.AddFullMessage(ts.sessionKey, msg) + } } default: } @@ -794,7 +840,7 @@ toolLoop: }) } } - if ts.opts.EnableSummary { + if !ts.opts.NoHistory && ts.opts.EnableSummary { al.contextManager.Compact(turnCtx, &CompactRequest{ SessionKey: ts.sessionKey, Reason: ContextCompressReasonSummarize, diff --git a/pkg/agent/pipeline_finalize.go b/pkg/agent/pipeline_finalize.go index 1f407825..d42f1091 100644 --- a/pkg/agent/pipeline_finalize.go +++ b/pkg/agent/pipeline_finalize.go @@ -33,6 +33,7 @@ func (p *Pipeline) Finalize( ts.setPhase(TurnPhaseCompleted) return turnResult{ finalContent: finalContent, + modelName: exec.llmModelName, status: turnStatus, followUps: append([]bus.InboundMessage(nil), ts.followUps...), }, nil @@ -44,6 +45,7 @@ func (p *Pipeline) Finalize( finalMsg := providers.Message{ Role: "assistant", Content: finalContent, + ModelName: exec.llmModelName, ReasoningContent: responseReasoningContent(exec.response), } ts.agent.Sessions.AddFullMessage(ts.sessionKey, finalMsg) @@ -58,11 +60,12 @@ func (p *Pipeline) Finalize( Message: err.Error(), }, ) + cancelConfiguredStreamingLLM(turnCtx, exec) return turnResult{status: TurnEndStatusError}, err } } - if ts.opts.EnableSummary { + if !ts.opts.NoHistory && ts.opts.EnableSummary { al.contextManager.Compact( turnCtx, &CompactRequest{ @@ -73,9 +76,31 @@ func (p *Pipeline) Finalize( ) } + contextUsage := computeContextUsage(ts.agent, ts.sessionKey) + streamErr := finalizeConfiguredStreamingLLM(turnCtx, ts, exec, finalContent, contextUsage) + // If streaming never became visible, keep the legacy Pico interim publish path + // so the final answer is still delivered outside normal SendResponse. + if ((streamErr != nil && !isConfiguredStreamingVisibleError(streamErr)) || exec.streamingFallback) && + !ts.opts.SendResponse && ts.opts.AllowInterimPicoPublish && finalContent != "" { + msg := outboundMessageForTurnWithOptions(ts, finalContent, outboundTurnMessageOptions{ + modelName: exec.llmModelName, + }) + msg.ContextUsage = contextUsage + markFinalOutbound(&msg) + _ = al.bus.PublishOutbound(turnCtx, msg) + } + if streamErr != nil && isConfiguredStreamingVisibleError(streamErr) { + ts.setPhase(TurnPhaseCompleted) + return turnResult{ + finalContent: finalContent, + status: TurnEndStatusError, + followUps: append([]bus.InboundMessage(nil), ts.followUps...), + }, streamErr + } ts.setPhase(TurnPhaseCompleted) return turnResult{ finalContent: finalContent, + modelName: exec.llmModelName, status: turnStatus, followUps: append([]bus.InboundMessage(nil), ts.followUps...), }, nil diff --git a/pkg/agent/pipeline_llm.go b/pkg/agent/pipeline_llm.go index 9fbe9b74..3dda03b3 100644 --- a/pkg/agent/pipeline_llm.go +++ b/pkg/agent/pipeline_llm.go @@ -37,9 +37,10 @@ func (p *Pipeline) CallLLM( // PreLLM: graceful terminal handling exec.gracefulTerminal, _ = ts.gracefulInterruptRequested() exec.providerToolDefs = ts.agent.Tools.ToProviderDefs() + exec.providerToolDefs = filterToolsByTurnProfile(exec.providerToolDefs, ts.profile) // Native web search support - webSearchEnabled := al.cfg.Tools.IsToolEnabled("web") + webSearchEnabled := al.cfg.Tools.IsToolEnabled("web") && turnProfileToolAllowed(ts.profile, "web_search") exec.useNativeSearch = webSearchEnabled && al.cfg.Tools.Web.PreferNative && func() bool { if ns, ok := ts.agent.Provider.(providers.NativeSearchCapable); ok { @@ -47,7 +48,6 @@ func (p *Pipeline) CallLLM( } return false }() - if exec.useNativeSearch { filtered := make([]providers.ToolDefinition, 0, len(exec.providerToolDefs)) for _, td := range exec.providerToolDefs { @@ -73,14 +73,7 @@ func (p *Pipeline) CallLLM( if exec.useNativeSearch { exec.llmOpts["native_search"] = true } - if ts.agent.ThinkingLevel != ThinkingOff { - if tc, ok := ts.agent.Provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() { - exec.llmOpts["thinking_level"] = string(ts.agent.ThinkingLevel) - } else { - logger.WarnCF("agent", "thinking_level is set but current provider does not support it, ignoring", - map[string]any{"agent_id": ts.agent.ID, "thinking_level": string(ts.agent.ThinkingLevel)}) - } - } + applyTurnThinkingOptions(exec, ts.agent, exec.activeProvider, true) exec.llmModel = exec.activeModel @@ -98,15 +91,27 @@ func (p *Pipeline) CallLLM( switch decision.normalizedAction() { case HookActionContinue, HookActionModify: if llmReq != nil { + prevModel := exec.llmModel exec.llmModel = llmReq.Model exec.callMessages = llmReq.Messages - exec.providerToolDefs = llmReq.Tools + exec.providerToolDefs = filterToolsByTurnProfile(llmReq.Tools, ts.profile) exec.llmOpts = llmReq.Options + nativeSearchAllowed := exec.useNativeSearch && + turnProfileToolAllowed(ts.profile, "web_search") + if !nativeSearchAllowed { + delete(exec.llmOpts, "native_search") + } + if strings.TrimSpace(exec.llmModel) != "" && exec.llmModel != prevModel { + p.applyBeforeLLMModelRewrite(ts, exec) + applyTurnThinkingOptions(exec, ts.agent, exec.activeProvider, true) + } } case HookActionAbortTurn: + cancelConfiguredStreamingLLM(turnCtx, exec) exec.abortedByHook = true return ControlBreak, nil case HookActionHardAbort: + cancelConfiguredStreamingLLM(turnCtx, exec) _ = ts.requestHardAbort() exec.abortedByHardAbort = true return ControlBreak, nil @@ -155,16 +160,44 @@ func (p *Pipeline) CallLLM( al.activeRequests.Add(1) defer al.activeRequests.Done() + if response, handled, streamErr := p.tryConfiguredStreamingLLM( + providerCtx, + ts, + exec, + messagesForCall, + toolDefsForCall, + ); handled { + return response, streamErr + } + if len(exec.activeCandidates) > 1 && p.Fallback != nil { - fbResult, fbErr := p.Fallback.Execute( + fbResult, fbErr := p.Fallback.ExecuteCandidate( providerCtx, exec.activeCandidates, - func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { - candidateProvider := exec.activeProvider - if cp, ok := ts.agent.CandidateProviders[providers.ModelKey(provider, model)]; ok { - candidateProvider = cp + func(ctx context.Context, candidate providers.FallbackCandidate) (*providers.LLMResponse, error) { + candidateProvider, err := providerForFallbackCandidate( + ts.agent, + exec.activeProvider, + exec.activeCandidates, + candidate.Provider, + candidate.Model, + ) + if err != nil { + return nil, err } - return candidateProvider.Chat(ctx, messagesForCall, toolDefsForCall, model, exec.llmOpts) + callOpts := shallowCloneLLMOptions(exec.llmOpts) + delete(callOpts, "thinking_level") + candidateCfg := resolveActiveModelConfig( + p.Cfg, + ts.agent.Workspace, + []providers.FallbackCandidate{candidate}, + candidate.Model, + p.Cfg.Agents.Defaults.Provider, + ) + candidateThinking := thinkingSettingsFromModelConfig(candidateCfg) + applyThinkingOption(callOpts, candidateProvider, candidateThinking, true, ts.agent.ID) + exec.suppressReasoning = shouldSuppressReasoningFor(candidateThinking) + return candidateProvider.Chat(ctx, messagesForCall, toolDefsForCall, candidate.Model, callOpts) }, ) if fbErr != nil { @@ -178,6 +211,16 @@ func (p *Pipeline) CallLLM( map[string]any{"agent_id": ts.agent.ID, "iteration": iteration}, ) } + for _, candidate := range exec.activeCandidates { + if candidate.StableKey() != fbResult.IdentityKey { + continue + } + exec.llmModelName = resolvedCandidateModelName( + []providers.FallbackCandidate{candidate}, + exec.llmModelName, + ) + break + } return fbResult.Response, nil } return exec.activeProvider.Chat(providerCtx, messagesForCall, toolDefsForCall, exec.llmModel, exec.llmOpts) @@ -203,6 +246,9 @@ func (p *Pipeline) CallLLM( exec.abortedByHardAbort = true return ControlBreak, nil } + if isConfiguredStreamingVisibleError(err) { + break + } // Retry without media if vision is unsupported if hasMediaRefs(exec.callMessages) && isVisionUnsupportedError(err) && retry < maxRetries { @@ -375,7 +421,7 @@ func (p *Pipeline) CallLLM( ) buildMessages := func(trimmedHistory []providers.Message) []providers.Message { fullHistory := append(append([]providers.Message(nil), trimmedHistory...), protectedTurnTail...) - rebuildPromptReq := promptBuildRequestForTurn(ts, fullHistory, exec.summary, "", nil) + rebuildPromptReq := promptBuildRequestForTurn(ts, fullHistory, exec.summary, "", nil, p.Cfg) rebuildPromptReq.ActiveSkills = append([]string(nil), contextualSkills...) return ts.agent.ContextBuilder.BuildMessagesFromPrompt(rebuildPromptReq) } @@ -466,9 +512,11 @@ func (p *Pipeline) CallLLM( exec.response = llmResp.Response } case HookActionAbortTurn: + cancelConfiguredStreamingLLM(turnCtx, exec) exec.abortedByHook = true return ControlBreak, nil case HookActionHardAbort: + cancelConfiguredStreamingLLM(turnCtx, exec) _ = ts.requestHardAbort() exec.abortedByHardAbort = true return ControlBreak, nil @@ -483,16 +531,31 @@ func (p *Pipeline) CallLLM( } } + if exec.suppressReasoning { + exec.response.Reasoning = "" + exec.response.ReasoningContent = "" + exec.response.ReasoningDetails = nil + } reasoningContent := responseReasoningContent(exec.response) shouldPublishPicoToolCallInterim := ts.channel == "pico" && len(exec.response.ToolCalls) > 0 if shouldPublishPicoToolCallInterim { // Pico tool-call turns publish their reasoning/content/tool summary as a // structured sequence after the tool-call payload is normalized below. } else if ts.channel == "pico" { - // Publish pico thoughts before the turn context is canceled at return time. - // The async variant can race with turn teardown and intermittently drop the - // thought message in CI even though the LLM produced reasoning content. - al.publishPicoReasoning(turnCtx, reasoningContent, ts.chatID) + if exec.streamingPublisher != nil && exec.streamingPublisher.ReasoningPublished() { + if err := exec.streamingPublisher.FinalizeReasoning(turnCtx, reasoningContent); err != nil { + logger.WarnCF("agent", "Failed to finalize streamed pico reasoning", map[string]any{ + "channel": ts.channel, + "chat_id": ts.chatID, + "error": err.Error(), + }) + } + } else { + // Publish pico thoughts before the turn context is canceled at return time. + // The async variant can race with turn teardown and intermittently drop the + // thought message in CI even though the LLM produced reasoning content. + al.publishPicoReasoning(turnCtx, reasoningContent, ts.chatID, ts.sessionKey, exec.llmModelName) + } } else { go al.handleReasoning( turnCtx, @@ -534,6 +597,7 @@ func (p *Pipeline) CallLLM( responseContent = exec.response.ReasoningContent } if steerMsgs := al.dequeueSteeringMessagesForScope(ts.sessionKey); len(steerMsgs) > 0 { + cancelConfiguredStreamingLLM(turnCtx, exec) logger.InfoCF("agent", "Steering arrived after direct LLM response; continuing turn", map[string]any{ "agent_id": ts.agent.ID, @@ -543,6 +607,7 @@ func (p *Pipeline) CallLLM( exec.pendingMessages = append(exec.pendingMessages, steerMsgs...) return ControlContinue, nil } + exec.finalContent = responseContent logger.InfoCF("agent", "LLM response without tool calls (direct answer)", map[string]any{ @@ -552,6 +617,7 @@ func (p *Pipeline) CallLLM( }) return ControlBreak, nil } + cancelConfiguredStreamingLLM(turnCtx, exec) // Tool-call path: normalize and prepare for tool execution exec.normalizedToolCalls = make([]providers.ToolCall, 0, len(exec.response.ToolCalls)) @@ -575,6 +641,7 @@ func (p *Pipeline) CallLLM( assistantMsg := providers.Message{ Role: "assistant", Content: exec.response.Content, + ModelName: exec.llmModelName, ReasoningContent: reasoningContent, } for _, tc := range exec.normalizedToolCalls { @@ -618,6 +685,7 @@ func (p *Pipeline) CallLLM( al.publishPicoToolCallInterim( turnCtx, ts, + exec.llmModelName, reasoningContent, exec.response.Content, assistantMsg.ToolCalls, @@ -626,3 +694,44 @@ func (p *Pipeline) CallLLM( return ControlToolLoop, nil } + +func (p *Pipeline) applyBeforeLLMModelRewrite(ts *turnState, exec *turnExecution) { + if p == nil || ts == nil || ts.agent == nil || exec == nil { + return + } + rawModel := strings.TrimSpace(exec.llmModel) + if rawModel == "" { + return + } + + defaultProvider := "openai" + if p.Cfg != nil { + if provider := strings.TrimSpace(p.Cfg.Agents.Defaults.Provider); provider != "" { + defaultProvider = provider + } + } + defaultProvider = effectiveDefaultProvider(defaultProvider) + candidates := resolveModelCandidates(p.Cfg, defaultProvider, rawModel, nil) + exec.activeCandidates = candidates + exec.activeModel = resolvedCandidateModel(candidates, rawModel) + exec.llmModel = exec.activeModel + exec.activeModelConfig = resolveActiveModelConfig(p.Cfg, ts.agent.Workspace, candidates, rawModel, defaultProvider) +} + +func providerForFallbackCandidate( + agent *AgentInstance, + activeProvider providers.LLMProvider, + activeCandidates []providers.FallbackCandidate, + provider string, + model string, +) (providers.LLMProvider, error) { + if agent != nil { + if cp, ok := agent.CandidateProviders[providers.ModelKey(provider, model)]; ok && cp != nil { + return cp, nil + } + } + if activeProvider == nil { + return nil, fmt.Errorf("fallback model %q has no active provider", model) + } + return activeProvider, nil +} diff --git a/pkg/agent/pipeline_setup.go b/pkg/agent/pipeline_setup.go index cf03526b..9d5033ce 100644 --- a/pkg/agent/pipeline_setup.go +++ b/pkg/agent/pipeline_setup.go @@ -36,14 +36,14 @@ func (p *Pipeline) SetupTurn(ctx context.Context, ts *turnState) (*turnExecution contextualSkills = ts.agent.ContextBuilder.ResolveActiveSkillsForContext(ts.activeSkills) } ts.recordSkillContextSnapshot(skillContextTriggerInitialBuild, contextualSkills) - initialPromptReq := promptBuildRequestForTurn(ts, history, summary, ts.userMessage, ts.media) + initialPromptReq := promptBuildRequestForTurn(ts, history, summary, ts.userMessage, ts.media, cfg) initialPromptReq.ActiveSkills = append([]string(nil), contextualSkills...) messages := ts.agent.ContextBuilder.BuildMessagesFromPrompt(initialPromptReq) messages = resolveMediaRefs(messages, p.MediaStore, maxMediaSize) if !ts.opts.NoHistory { - toolDefs := ts.agent.Tools.ToProviderDefs() + toolDefs := filterToolsByTurnProfile(ts.agent.Tools.ToProviderDefs(), ts.profile) if isOverContextBudget(ts.agent.ContextWindow, messages, toolDefs, ts.agent.MaxTokens) { logger.WarnCF("agent", "Proactive compression: context budget exceeded before LLM call", map[string]any{"session_key": ts.sessionKey}) @@ -71,7 +71,7 @@ func (p *Pipeline) SetupTurn(ctx context.Context, ts *turnState) (*turnExecution history, messages, fit = trimHistoryToFitContextWindow( history, func(trimmedHistory []providers.Message) []providers.Message { - rebuildPromptReq := promptBuildRequestForTurn(ts, trimmedHistory, summary, ts.userMessage, ts.media) + rebuildPromptReq := promptBuildRequestForTurn(ts, trimmedHistory, summary, ts.userMessage, ts.media, cfg) rebuildPromptReq.ActiveSkills = append([]string(nil), contextualSkills...) rebuilt := ts.agent.ContextBuilder.BuildMessagesFromPrompt(rebuildPromptReq) return resolveMediaRefs(rebuilt, p.MediaStore, maxMediaSize) @@ -117,6 +117,11 @@ func (p *Pipeline) SetupTurn(ctx context.Context, ts *turnState) (*turnExecution if usedLight && ts.agent.LightProvider != nil { activeProvider = ts.agent.LightProvider } + activeModelName := strings.TrimSpace(ts.agent.Model) + if usedLight { + activeModelName = strings.TrimSpace(sideQuestionModelName(ts.agent, true)) + } + activeModelName = resolvedCandidateModelName(activeCandidates, activeModelName) exec := newTurnExecution( ts.agent, @@ -127,6 +132,14 @@ func (p *Pipeline) SetupTurn(ctx context.Context, ts *turnState) (*turnExecution ) exec.activeCandidates = activeCandidates exec.activeModel = activeModel + exec.activeModelConfig = resolveActiveModelConfig( + p.Cfg, + ts.agent.Workspace, + activeCandidates, + activeModel, + p.Cfg.Agents.Defaults.Provider, + ) + exec.llmModelName = activeModelName exec.activeProvider = activeProvider exec.usedLight = usedLight diff --git a/pkg/agent/pipeline_streaming.go b/pkg/agent/pipeline_streaming.go new file mode 100644 index 00000000..115cbffc --- /dev/null +++ b/pkg/agent/pipeline_streaming.go @@ -0,0 +1,489 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "reflect" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/providers" +) + +func (p *Pipeline) tryConfiguredStreamingLLM( + ctx context.Context, + ts *turnState, + exec *turnExecution, + messagesForCall []providers.Message, + toolDefsForCall []providers.ToolDefinition, +) (*providers.LLMResponse, bool, error) { + exec.streamingPublisher = nil + exec.streamingFallback = false + if !p.configuredStreamingEligible(ts, exec) { + return nil, false, nil + } + streamProvider, ok := exec.activeProvider.(providers.StreamingProvider) + if !ok { + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.activeModel, + "reason": "provider_not_streaming", + }) + return nil, false, nil + } + + streamer, ok := p.Bus.GetStreamer(ctx, ts.channel, ts.chatID, ts.sessionKey) + if !ok || streamer == nil { + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "chat_id": ts.chatID, + "model": exec.activeModel, + "reason": "streamer_unavailable", + }) + return nil, false, nil + } + + publisher := &streamingChunkPublisher{ + streamer: streamer, + channel: ts.channel, + chatID: ts.chatID, + modelName: exec.llmModelName, + } + + logger.DebugCF("agent", "configured streaming enabled", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "chat_id": ts.chatID, + "model": exec.llmModel, + }) + + chunkCount := 0 + firstChunkAt := time.Time{} + lastChunkAt := time.Time{} + recordChunk := func() { + now := time.Now() + chunkCount++ + if firstChunkAt.IsZero() { + firstChunkAt = now + } + lastChunkAt = now + } + var response *providers.LLMResponse + var streamErr error + if eventProvider, ok := exec.activeProvider.(providers.StreamingEventProvider); ok { + response, streamErr = eventProvider.ChatStreamEvents( + ctx, + messagesForCall, + toolDefsForCall, + exec.llmModel, + exec.llmOpts, + func(chunk providers.StreamChunk) { + recordChunk() + if !exec.suppressReasoning && strings.TrimSpace(chunk.ReasoningContent) != "" { + publisher.UpdateReasoning(ctx, chunk.ReasoningContent) + } + if strings.TrimSpace(chunk.Content) != "" { + publisher.Update(ctx, chunk.Content) + } + }, + ) + } else { + response, streamErr = streamProvider.ChatStream( + ctx, + messagesForCall, + toolDefsForCall, + exec.llmModel, + exec.llmOpts, + func(accumulated string) { + recordChunk() + publisher.Update(ctx, accumulated) + }, + ) + } + logConfiguredStreamingSummary(ts, exec, chunkCount, firstChunkAt, lastChunkAt, streamErr) + if streamErr == nil { + if updateErr := publisher.Err(); updateErr != nil { + logFields := map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.llmModel, + "error": updateErr.Error(), + } + if publisher.Published() { + logger.WarnCF("agent", "ChatStream update failed after visible output", logFields) + return nil, true, configuredStreamingVisibleError{err: updateErr} + } + logger.WarnCF("agent", "ChatStream update failed before visible output; retrying with Chat", logFields) + publisher.Cancel(ctx) + fallbackResponse, err := exec.activeProvider.Chat( + ctx, + messagesForCall, + toolDefsForCall, + exec.llmModel, + exec.llmOpts, + ) + if err == nil && fallbackResponse != nil { + exec.streamingFallback = true + } + return fallbackResponse, true, err + } + } + if streamErr != nil { + if !publisher.Published() { + logger.WarnCF("agent", "ChatStream failed before visible output; retrying with Chat", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.llmModel, + "error": streamErr.Error(), + }) + publisher.Cancel(ctx) + fallbackResponse, err := exec.activeProvider.Chat( + ctx, + messagesForCall, + toolDefsForCall, + exec.llmModel, + exec.llmOpts, + ) + if err == nil && fallbackResponse != nil { + exec.streamingFallback = true + } + return fallbackResponse, true, err + } + return nil, true, configuredStreamingVisibleError{err: streamErr} + } + + if response != nil { + exec.streamingPublisher = publisher + } + + return response, true, nil +} + +func logConfiguredStreamingSummary( + ts *turnState, + exec *turnExecution, + chunkCount int, + firstChunkAt time.Time, + lastChunkAt time.Time, + streamErr error, +) { + fields := map[string]any{ + "chunks": chunkCount, + } + if ts != nil { + fields["agent_id"] = ts.agent.ID + fields["channel"] = ts.channel + } + if exec != nil { + fields["model"] = exec.llmModel + } + if !firstChunkAt.IsZero() && !lastChunkAt.IsZero() { + fields["chunk_span_ms"] = lastChunkAt.Sub(firstChunkAt).Milliseconds() + } + if streamErr != nil { + fields["error"] = streamErr.Error() + } + logger.DebugCF("agent", "configured streaming completed", fields) +} + +type configuredStreamingVisibleError struct { + err error +} + +func (e configuredStreamingVisibleError) Error() string { + if e.err == nil { + return "configured streaming failed after visible output" + } + return e.err.Error() +} + +func (e configuredStreamingVisibleError) Unwrap() error { + return e.err +} + +func isConfiguredStreamingVisibleError(err error) bool { + var visibleErr configuredStreamingVisibleError + return errors.As(err, &visibleErr) +} + +func finalizeConfiguredStreamingLLM( + ctx context.Context, + ts *turnState, + exec *turnExecution, + content string, + contextUsage *bus.ContextUsage, +) error { + if exec == nil || exec.streamingPublisher == nil { + return nil + } + publisher := exec.streamingPublisher + exec.streamingPublisher = nil + visibleBeforeFinalize := publisher.Published() + if err := publisher.Finalize(ctx, content, contextUsage); err != nil { + if visibleBeforeFinalize { + logger.WarnCF("agent", "stream final flush failed after visible output", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.llmModel, + "error": err.Error(), + }) + return configuredStreamingVisibleError{err: err} + } + publisher.Cancel(ctx) + logger.WarnCF("agent", "stream final flush failed", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.llmModel, + "error": err.Error(), + }) + return err + } + return nil +} + +func cancelConfiguredStreamingLLM(ctx context.Context, exec *turnExecution) { + if exec == nil || exec.streamingPublisher == nil { + return + } + publisher := exec.streamingPublisher + exec.streamingPublisher = nil + publisher.Cancel(ctx) +} + +func (p *Pipeline) configuredStreamingEligible(ts *turnState, exec *turnExecution) bool { + if p == nil || ts == nil || exec == nil || p.Bus == nil { + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "reason": "missing_pipeline_state", + }) + return false + } + if strings.TrimSpace(ts.channel) == "" || strings.TrimSpace(ts.chatID) == "" { + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "chat_id": ts.chatID, + "model": exec.activeModel, + "reason": "missing_channel_context", + }) + return false + } + if !ts.opts.SendResponse && !ts.opts.AllowInterimPicoPublish { + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "chat_id": ts.chatID, + "model": exec.activeModel, + "reason": "turn_output_disabled", + }) + return false + } + if len(exec.activeCandidates) != 1 { + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.activeModel, + "candidates": len(exec.activeCandidates), + "reason": "fallback_candidates_enabled", + }) + return false + } + if exec.activeModelConfig == nil || !exec.activeModelConfig.Streaming.Enabled { + modelName := "" + modelStreaming := false + if exec.activeModelConfig != nil { + modelName = exec.activeModelConfig.ModelName + modelStreaming = exec.activeModelConfig.Streaming.Enabled + } + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.activeModel, + "model_name": modelName, + "model_streaming": modelStreaming, + "has_model_config": exec.activeModelConfig != nil, + "reason": "model_streaming_disabled", + }) + return false + } + channelStreaming, ok := p.channelStreamingConfig(ts.channel) + if !ok || !channelStreaming.Enabled { + logger.DebugCF("agent", "configured streaming not used", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "model": exec.activeModel, + "channel_streaming": channelStreaming.Enabled, + "has_channel_config": ok, + "reason": "channel_streaming_disabled", + }) + return false + } + return true +} + +func (p *Pipeline) channelStreamingConfig(channelName string) (config.StreamingConfig, bool) { + if p == nil || p.Cfg == nil || p.Cfg.Channels == nil { + return config.StreamingConfig{}, false + } + ch := p.Cfg.Channels[channelName] + if ch == nil { + return config.StreamingConfig{}, false + } + decoded, err := ch.GetDecoded() + if err != nil { + logger.WarnCF("agent", "channel streaming config decode failed", map[string]any{ + "channel": channelName, + "error": err.Error(), + }) + return config.StreamingConfig{}, false + } + return streamingConfigFromDecodedSettings(decoded) +} + +func streamingConfigFromDecodedSettings(decoded any) (config.StreamingConfig, bool) { + value := reflect.ValueOf(decoded) + if !value.IsValid() { + return config.StreamingConfig{}, false + } + if value.Kind() == reflect.Ptr { + if value.IsNil() { + return config.StreamingConfig{}, false + } + value = value.Elem() + } + if value.Kind() != reflect.Struct { + return config.StreamingConfig{}, false + } + + field := value.FieldByName("Streaming") + if !field.IsValid() || !field.CanInterface() { + return config.StreamingConfig{}, false + } + streaming, ok := field.Interface().(config.StreamingConfig) + return streaming, ok +} + +type streamingChunkPublisher struct { + streamer bus.Streamer + channel string + chatID string + modelName string + published bool + reasoningPublished bool + err error +} + +func (p *streamingChunkPublisher) Update(ctx context.Context, accumulated string) { + if p == nil || p.streamer == nil || strings.TrimSpace(accumulated) == "" { + return + } + if setter, ok := p.streamer.(interface{ SetModelName(modelName string) }); ok { + setter.SetModelName(p.modelName) + } + if err := p.streamer.Update(ctx, accumulated); err != nil { + p.err = err + logger.WarnCF("agent", "stream update failed", map[string]any{ + "channel": p.channel, + "chat_id": p.chatID, + "error": err.Error(), + }) + return + } + p.published = true +} + +func (p *streamingChunkPublisher) UpdateReasoning(ctx context.Context, accumulated string) { + if p == nil || p.streamer == nil || strings.TrimSpace(accumulated) == "" { + return + } + if setter, ok := p.streamer.(interface{ SetModelName(modelName string) }); ok { + setter.SetModelName(p.modelName) + } + reasoningStreamer, ok := p.streamer.(bus.ReasoningStreamer) + if !ok { + return + } + if err := reasoningStreamer.UpdateReasoning(ctx, accumulated); err != nil { + p.err = err + logger.WarnCF("agent", "stream reasoning update failed", map[string]any{ + "channel": p.channel, + "chat_id": p.chatID, + "error": err.Error(), + }) + return + } + p.reasoningPublished = true +} + +func (p *streamingChunkPublisher) Published() bool { + return p != nil && p.published +} + +func (p *streamingChunkPublisher) ReasoningPublished() bool { + return p != nil && p.reasoningPublished +} + +func (p *streamingChunkPublisher) Err() error { + if p == nil { + return nil + } + return p.err +} + +func (p *streamingChunkPublisher) Finalize(ctx context.Context, content string, contextUsage *bus.ContextUsage) error { + if p == nil || p.streamer == nil { + return nil + } + if strings.TrimSpace(content) == "" && !p.published { + return nil + } + if setter, ok := p.streamer.(interface{ SetModelName(modelName string) }); ok { + setter.SetModelName(p.modelName) + } + var err error + if streamer, ok := p.streamer.(bus.ContextUsageStreamer); ok { + err = streamer.FinalizeWithContext(ctx, content, contextUsage) + } else { + err = p.streamer.Finalize(ctx, content) + } + if err != nil { + return fmt.Errorf("stream finalize: %w", err) + } + p.published = true + return nil +} + +func (p *streamingChunkPublisher) FinalizeReasoning(ctx context.Context, content string) error { + if p == nil || p.streamer == nil || !p.reasoningPublished || strings.TrimSpace(content) == "" { + return nil + } + reasoningStreamer, ok := p.streamer.(bus.ReasoningStreamer) + if !ok { + return nil + } + if err := reasoningStreamer.FinalizeReasoning(ctx, content); err != nil { + return fmt.Errorf("stream reasoning finalize: %w", err) + } + return nil +} + +func (p *streamingChunkPublisher) ClearFinalizedStreamMarker() { + if p == nil || p.streamer == nil { + return + } + if cleaner, ok := p.streamer.(interface{ ClearFinalizedStreamMarker() }); ok { + cleaner.ClearFinalizedStreamMarker() + } +} + +func (p *streamingChunkPublisher) Cancel(ctx context.Context) { + if p == nil || p.streamer == nil { + return + } + p.streamer.Cancel(ctx) +} diff --git a/pkg/agent/pipeline_streaming_test.go b/pkg/agent/pipeline_streaming_test.go new file mode 100644 index 00000000..9bba1b90 --- /dev/null +++ b/pkg/agent/pipeline_streaming_test.go @@ -0,0 +1,1169 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "io" + "os" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" +) + +type configuredStreamingProvider struct { + chatCalls int + streamCalls int + eventCalls int + chatModels []string + streamModels []string + + chatResponse *providers.LLMResponse + streamPlan []configuredStreamingCall + eventPlan []configuredStreamingEventCall +} + +type configuredStreamingCall struct { + chunks []string + response *providers.LLMResponse + err error +} + +type configuredStreamingEventCall struct { + chunks []providers.StreamChunk + response *providers.LLMResponse + err error +} + +func (p *configuredStreamingProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.chatCalls++ + p.chatModels = append(p.chatModels, model) + if p.chatResponse != nil { + return p.chatResponse, nil + } + return &providers.LLMResponse{Content: "chat response"}, nil +} + +func (p *configuredStreamingProvider) ChatStream( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, + onChunk func(accumulated string), +) (*providers.LLMResponse, error) { + p.streamCalls++ + p.streamModels = append(p.streamModels, model) + var plan configuredStreamingCall + if len(p.streamPlan) >= p.streamCalls { + plan = p.streamPlan[p.streamCalls-1] + } + for _, chunk := range plan.chunks { + onChunk(chunk) + } + if plan.err != nil { + return nil, plan.err + } + if plan.response != nil { + return plan.response, nil + } + return &providers.LLMResponse{Content: "stream response"}, nil +} + +func (p *configuredStreamingProvider) ChatStreamEvents( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, + onChunk func(providers.StreamChunk), +) (*providers.LLMResponse, error) { + p.eventCalls++ + p.streamCalls++ + p.streamModels = append(p.streamModels, model) + var plan configuredStreamingEventCall + if len(p.eventPlan) >= p.eventCalls { + plan = p.eventPlan[p.eventCalls-1] + } else if len(p.streamPlan) >= p.eventCalls { + legacyPlan := p.streamPlan[p.eventCalls-1] + plan.response = legacyPlan.response + plan.err = legacyPlan.err + for _, chunk := range legacyPlan.chunks { + plan.chunks = append(plan.chunks, providers.StreamChunk{Content: chunk}) + } + } + for _, chunk := range plan.chunks { + onChunk(chunk) + } + if plan.err != nil { + return nil, plan.err + } + if plan.response != nil { + return plan.response, nil + } + return &providers.LLMResponse{Content: "stream response"}, nil +} + +func (p *configuredStreamingProvider) GetDefaultModel() string { + return "mock-model" +} + +type configuredStreamingChatOnlyProvider struct { + chatCalls int +} + +func (p *configuredStreamingChatOnlyProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.chatCalls++ + return &providers.LLMResponse{Content: "chat only"}, nil +} + +func (p *configuredStreamingChatOnlyProvider) GetDefaultModel() string { + return "mock-model" +} + +type configuredStreamingDelegate struct { + streamer bus.Streamer +} + +func (d configuredStreamingDelegate) GetStreamer( + ctx context.Context, + channel, chatID, sessionKey string, +) (bus.Streamer, bool) { + if d.streamer == nil { + return nil, false + } + return d.streamer, true +} + +type recordingStreamer struct { + updates []string + finalized []string + reasoningUpdates []string + reasoningFinalized []string + events []string + canceled int +} + +func (s *recordingStreamer) Update(ctx context.Context, content string) error { + s.updates = append(s.updates, content) + s.events = append(s.events, "content:"+content) + return nil +} + +func (s *recordingStreamer) Finalize(ctx context.Context, content string) error { + s.finalized = append(s.finalized, content) + s.events = append(s.events, "final:"+content) + return nil +} + +func (s *recordingStreamer) UpdateReasoning(ctx context.Context, content string) error { + s.reasoningUpdates = append(s.reasoningUpdates, content) + s.events = append(s.events, "reasoning:"+content) + return nil +} + +func (s *recordingStreamer) FinalizeReasoning(ctx context.Context, content string) error { + s.reasoningFinalized = append(s.reasoningFinalized, content) + s.events = append(s.events, "reasoning-final:"+content) + return nil +} + +func (s *recordingStreamer) Cancel(context.Context) { + s.canceled++ +} + +type cleanableRecordingStreamer struct { + recordingStreamer + clearMarkers int +} + +func (s *cleanableRecordingStreamer) ClearFinalizedStreamMarker() { + s.clearMarkers++ +} + +type failingFinalizeStreamer struct { + recordingStreamer + err error +} + +func (s *failingFinalizeStreamer) Finalize(ctx context.Context, content string) error { + s.finalized = append(s.finalized, content) + return s.err +} + +type failingUpdateStreamer struct { + recordingStreamer + err error +} + +func (s *failingUpdateStreamer) Update(ctx context.Context, content string) error { + s.updates = append(s.updates, content) + return s.err +} + +type failNthUpdateStreamer struct { + recordingStreamer + failOn int + err error +} + +func (s *failNthUpdateStreamer) Update(ctx context.Context, content string) error { + s.updates = append(s.updates, content) + if len(s.updates) == s.failOn { + return s.err + } + return nil +} + +type configuredStreamingAfterHook struct { + content string + action HookAction +} + +func (h configuredStreamingAfterHook) BeforeLLM( + ctx context.Context, + req *LLMHookRequest, +) (*LLMHookRequest, HookDecision, error) { + return req, HookDecision{Action: HookActionContinue}, nil +} + +func (h configuredStreamingAfterHook) AfterLLM( + ctx context.Context, + resp *LLMHookResponse, +) (*LLMHookResponse, HookDecision, error) { + if h.action == HookActionAbortTurn || h.action == HookActionHardAbort { + return resp, HookDecision{Action: h.action}, nil + } + next := resp.Clone() + next.Response.Content = h.content + return next, HookDecision{Action: HookActionModify}, nil +} + +type configuredStreamingBeforeModelHook struct { + model string +} + +func (h configuredStreamingBeforeModelHook) BeforeLLM( + ctx context.Context, + req *LLMHookRequest, +) (*LLMHookRequest, HookDecision, error) { + next := req.Clone() + next.Model = h.model + return next, HookDecision{Action: HookActionModify}, nil +} + +func (h configuredStreamingBeforeModelHook) AfterLLM( + ctx context.Context, + resp *LLMHookResponse, +) (*LLMHookResponse, HookDecision, error) { + return resp, HookDecision{Action: HookActionContinue}, nil +} + +func TestConfiguredStreamingEligibilityGates(t *testing.T) { + tests := []struct { + name string + channel string + channelStreaming bool + modelStreaming bool + fallbacks []string + streamingProvider bool + streamDelegate bool + wantStreamCalls int + wantChatCalls int + }{ + { + name: "channel and model enabled streams", + channel: "pico", + channelStreaming: true, + modelStreaming: true, + streamingProvider: true, + streamDelegate: true, + wantStreamCalls: 1, + }, + { + name: "wecom channel and model enabled streams", + channel: "wecom", + channelStreaming: true, + modelStreaming: true, + streamingProvider: true, + streamDelegate: true, + wantStreamCalls: 1, + }, + { + name: "channel disabled uses chat", + channel: "pico", + modelStreaming: true, + streamingProvider: true, + streamDelegate: true, + wantChatCalls: 1, + }, + { + name: "model disabled uses chat", + channel: "pico", + channelStreaming: true, + streamingProvider: true, + streamDelegate: true, + wantChatCalls: 1, + }, + { + name: "provider without streaming uses chat", + channel: "pico", + channelStreaming: true, + modelStreaming: true, + streamDelegate: true, + wantChatCalls: 1, + }, + { + name: "multi candidate fallback uses chat", + channel: "pico", + channelStreaming: true, + modelStreaming: true, + fallbacks: []string{"fallback-model"}, + streamingProvider: true, + streamDelegate: true, + wantChatCalls: 1, + }, + { + name: "missing streamer uses chat", + channel: "pico", + channelStreaming: true, + modelStreaming: true, + streamingProvider: true, + wantChatCalls: 1, + }, + { + name: "omitted fields use chat", + channel: "pico", + streamingProvider: true, + streamDelegate: true, + wantChatCalls: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, tt.channelStreaming, tt.modelStreaming, tt.fallbacks) + msgBus := bus.NewMessageBus() + if tt.streamDelegate { + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: &recordingStreamer{}}) + } + + if tt.streamingProvider { + provider := &configuredStreamingProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + runConfiguredStreamingTurn(t, al, tt.channel) + if provider.streamCalls != tt.wantStreamCalls { + t.Fatalf("ChatStream calls = %d, want %d", provider.streamCalls, tt.wantStreamCalls) + } + if provider.chatCalls != tt.wantChatCalls { + t.Fatalf("Chat calls = %d, want %d", provider.chatCalls, tt.wantChatCalls) + } + return + } + + provider := &configuredStreamingChatOnlyProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + runConfiguredStreamingTurn(t, al, tt.channel) + if provider.chatCalls != tt.wantChatCalls { + t.Fatalf("Chat calls = %d, want %d", provider.chatCalls, tt.wantChatCalls) + } + }) + } +} + +func TestConfiguredStreamingPreChunkFailureFallsBackToChat(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: &recordingStreamer{}}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + err: errors.New("stream setup failed"), + }}, + chatResponse: &providers.LLMResponse{Content: "chat after stream failure"}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != "chat after stream failure" { + t.Fatalf("response = %q, want chat fallback response", got) + } + if provider.streamCalls != 1 || provider.chatCalls != 1 { + t.Fatalf("calls = stream:%d chat:%d, want stream:1 chat:1", provider.streamCalls, provider.chatCalls) + } + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "chat after stream failure" { + t.Fatalf("fallback outbound content = %q, want chat after stream failure", outbound.Content) + } + case <-time.After(time.Second): + t.Fatal("expected fallback outbound after pre-chunk stream failure") + } +} + +func TestConfiguredStreamingDisabledForInternalTurnWithoutUserVisibleOutput(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"internal stream"}, + response: &providers.LLMResponse{Content: "stream response"}, + }}, + chatResponse: &providers.LLMResponse{Content: "chat response"}, + } + al := NewAgentLoop(cfg, msgBus, provider) + opts := configuredStreamingProcessOptions("pico") + opts.SendResponse = false + opts.AllowInterimPicoPublish = false + + got, err := al.runAgentLoop(context.Background(), al.GetRegistry().GetDefaultAgent(), opts) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if got != "chat response" { + t.Fatalf("response = %q, want chat response", got) + } + if provider.streamCalls != 0 || provider.chatCalls != 1 { + t.Fatalf("calls = stream:%d chat:%d, want stream:0 chat:1", provider.streamCalls, provider.chatCalls) + } + if len(streamer.updates) != 0 || len(streamer.finalized) != 0 { + t.Fatalf("streamer updates=%v finalized=%v, want no streaming output", streamer.updates, streamer.finalized) + } +} + +func TestConfiguredStreamingVisibleSendResponseFalseRetainsFinalizedStreamMarker(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &cleanableRecordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"visible stream"}, + response: &providers.LLMResponse{Content: "stream response"}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != "stream response" { + t.Fatalf("response = %q, want stream response", got) + } + if streamer.clearMarkers != 0 { + t.Fatalf("clear markers = %d, want 0", streamer.clearMarkers) + } +} + +func TestConfiguredStreamingStreamsPicoReasoningBeforeAnswerContent(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + eventPlan: []configuredStreamingEventCall{{ + chunks: []providers.StreamChunk{ + {ReasoningContent: "thinking"}, + {ReasoningContent: "thinking more"}, + {Content: "answer"}, + }, + response: &providers.LLMResponse{ + Content: "answer", + ReasoningContent: "thinking more", + }, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + if got != "answer" { + t.Fatalf("response = %q, want answer", got) + } + if provider.eventCalls != 1 { + t.Fatalf("ChatStreamEvents calls = %d, want 1", provider.eventCalls) + } + if len(streamer.reasoningUpdates) != 2 { + t.Fatalf("reasoning updates = %v, want two streamed updates", streamer.reasoningUpdates) + } + if len(streamer.updates) != 1 || streamer.updates[0] != "answer" { + t.Fatalf("content updates = %v, want [answer]", streamer.updates) + } + if len(streamer.events) < 3 || + streamer.events[0] != "reasoning:thinking" || + streamer.events[1] != "reasoning:thinking more" || + streamer.events[2] != "content:answer" { + t.Fatalf("stream event order = %v, want reasoning before answer content", streamer.events) + } + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("expected streamed reasoning to avoid a later thought outbound, got %+v", outbound) + case <-time.After(50 * time.Millisecond): + } +} + +func TestConfiguredStreamingSuppressesPicoReasoningWhenThinkingOff(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + cfg.ModelList[0].ThinkingLevel = "off" + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + eventPlan: []configuredStreamingEventCall{{ + chunks: []providers.StreamChunk{ + {ReasoningContent: "thinking"}, + {Content: "answer"}, + }, + response: &providers.LLMResponse{ + Content: "answer", + ReasoningContent: "thinking", + }, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + if got != "answer" { + t.Fatalf("response = %q, want answer", got) + } + if len(streamer.reasoningUpdates) != 0 { + t.Fatalf("reasoning updates = %v, want none when thinking is off", streamer.reasoningUpdates) + } + if len(streamer.reasoningFinalized) != 0 { + t.Fatalf("reasoning finalized = %v, want none when thinking is off", streamer.reasoningFinalized) + } + if len(streamer.updates) != 1 || streamer.updates[0] != "answer" { + t.Fatalf("content updates = %v, want [answer]", streamer.updates) + } + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("expected no reasoning outbound when thinking is off, got %+v", outbound) + case <-time.After(50 * time.Millisecond): + } +} + +func TestConfiguredStreamingFinalFlushFailureAfterVisibleOutputReturnsErrorWithoutFallbackOrCancel(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &failingFinalizeStreamer{err: errors.New("final failed")} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"partial stream"}, + response: &providers.LLMResponse{Content: "stream response"}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + configuredStreamingProcessOptions("pico"), + ) + if err == nil { + t.Fatal("expected final flush failure after visible output to return an error") + } + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("unexpected fallback outbound after visible final flush failure: %#v", outbound) + default: + } + if streamer.canceled != 0 { + t.Fatalf("streamer canceled = %d, want 0 for already-visible final flush failure", streamer.canceled) + } +} + +func TestConfiguredStreamingFinalFlushFailureBeforeVisibleOutputPublishesFallback(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &failingFinalizeStreamer{err: errors.New("final failed")} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + response: &providers.LLMResponse{Content: "stream response"}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != "stream response" { + t.Fatalf("response = %q, want stream response", got) + } + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "stream response" { + t.Fatalf("fallback outbound content = %q, want stream response", outbound.Content) + } + if got := outbound.Context.Raw["model_name"]; got != "test-model" { + t.Fatalf("fallback outbound model_name = %q, want %q", got, "test-model") + } + case <-time.After(time.Second): + t.Fatal("expected fallback outbound after invisible final stream flush failure") + } + if streamer.canceled != 1 { + t.Fatalf("streamer canceled = %d, want 1", streamer.canceled) + } +} + +func TestConfiguredStreamingFinalFlushFailureBeforeVisibleOutputKeepsNormalOutbound(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &failingFinalizeStreamer{err: errors.New("final failed")} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + response: &providers.LLMResponse{Content: "stream response"}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + opts := configuredStreamingProcessOptions("pico") + opts.SendResponse = true + + got, err := al.runAgentLoop(context.Background(), al.GetRegistry().GetDefaultAgent(), opts) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if got != "stream response" { + t.Fatalf("response = %q, want stream response", got) + } + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "stream response" { + t.Fatalf("normal outbound content = %q, want stream response", outbound.Content) + } + case <-time.After(time.Second): + t.Fatal("expected normal outbound after invisible final stream flush failure") + } + if streamer.canceled != 1 { + t.Fatalf("streamer canceled = %d, want 1", streamer.canceled) + } +} + +func TestConfiguredStreamingUpdateFailureThenStreamErrorFallsBackToChat(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + msgBus := bus.NewMessageBus() + streamer := &failingUpdateStreamer{err: errors.New("draft failed")} + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"not visible"}, + err: errors.New("stream failed after invisible update"), + }}, + chatResponse: &providers.LLMResponse{Content: "chat fallback after invisible update"}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != "chat fallback after invisible update" { + t.Fatalf("response = %q, want chat fallback", got) + } + if provider.streamCalls != 1 || provider.chatCalls != 1 { + t.Fatalf("calls = stream:%d chat:%d, want stream:1 chat:1", provider.streamCalls, provider.chatCalls) + } + if streamer.canceled != 1 { + t.Fatalf("streamer canceled = %d, want 1", streamer.canceled) + } + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "chat fallback after invisible update" { + t.Fatalf("fallback outbound content = %q, want chat fallback after invisible update", outbound.Content) + } + case <-time.After(time.Second): + t.Fatal("expected fallback outbound after update failure and stream error") + } +} + +func TestConfiguredStreamingUpdateFailureThenStreamSuccessFallsBackToChat(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + msgBus := bus.NewMessageBus() + streamer := &failingUpdateStreamer{err: errors.New("draft failed")} + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"not visible"}, + response: &providers.LLMResponse{Content: "stream response"}, + }}, + chatResponse: &providers.LLMResponse{Content: "chat fallback after invisible update"}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != "chat fallback after invisible update" { + t.Fatalf("response = %q, want chat fallback", got) + } + if provider.streamCalls != 1 || provider.chatCalls != 1 { + t.Fatalf("calls = stream:%d chat:%d, want stream:1 chat:1", provider.streamCalls, provider.chatCalls) + } + if len(streamer.finalized) != 0 { + t.Fatalf("stream finalized = %v, want none", streamer.finalized) + } + select { + case outbound := <-msgBus.OutboundChan(): + if outbound.Content != "chat fallback after invisible update" { + t.Fatalf("fallback outbound content = %q, want chat fallback after invisible update", outbound.Content) + } + case <-time.After(time.Second): + t.Fatal("expected fallback outbound after update failure and stream success") + } +} + +func TestConfiguredStreamingLaterUpdateFailureThenStreamSuccessReturnsVisibleError(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + msgBus := bus.NewMessageBus() + streamer := &failNthUpdateStreamer{failOn: 2, err: errors.New("draft failed")} + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"visible chunk", "failed later chunk"}, + response: &providers.LLMResponse{Content: "stream response"}, + }}, + chatResponse: &providers.LLMResponse{Content: "chat fallback after later update failure"}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + configuredStreamingProcessOptions("pico"), + ) + if err == nil { + t.Fatal("expected post-visible update failure to return an error") + } + if provider.streamCalls != 1 || provider.chatCalls != 0 { + t.Fatalf("calls = stream:%d chat:%d, want stream:1 chat:0", provider.streamCalls, provider.chatCalls) + } + if streamer.canceled != 0 { + t.Fatalf("streamer canceled = %d, want 0", streamer.canceled) + } + if len(streamer.finalized) != 0 { + t.Fatalf("stream finalized = %v, want none", streamer.finalized) + } + select { + case outbound := <-msgBus.OutboundChan(): + t.Fatalf("unexpected fallback outbound after post-visible update failure: %#v", outbound) + default: + } +} + +func TestConfiguredStreamingBeforeLLMModelRewriteReevaluatesModelStreaming(t *testing.T) { + tests := []struct { + name string + initialModelStreaming bool + rewriteModel string + rewriteModelStreaming bool + fallbacks []string + wantStreamCalls int + wantChatCalls int + wantFinalizedResponses int + }{ + { + name: "rewrite to disabled model uses chat", + initialModelStreaming: true, + rewriteModel: "hook-disabled-model", + wantChatCalls: 1, + }, + { + name: "rewrite to enabled model streams", + rewriteModel: "hook-enabled-model", + rewriteModelStreaming: true, + fallbacks: []string{"fallback-model"}, + wantStreamCalls: 1, + wantFinalizedResponses: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, tt.initialModelStreaming, tt.fallbacks) + cfg.ModelList = append(cfg.ModelList, &config.ModelConfig{ + ModelName: tt.rewriteModel, + Provider: "openai", + Model: "openai/" + tt.rewriteModel, + Streaming: config.ModelStreamingConfig{Enabled: tt.rewriteModelStreaming}, + }) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"streamed after hook model rewrite"}, + response: &providers.LLMResponse{Content: "stream response"}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + if err := al.MountHook(NamedHook("rewrite-model", configuredStreamingBeforeModelHook{ + model: tt.rewriteModel, + })); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + got := runConfiguredStreamingTurn(t, al, "pico") + + if provider.streamCalls != tt.wantStreamCalls { + t.Fatalf("ChatStream calls = %d, want %d", provider.streamCalls, tt.wantStreamCalls) + } + if provider.chatCalls != tt.wantChatCalls { + t.Fatalf("Chat calls = %d, want %d", provider.chatCalls, tt.wantChatCalls) + } + if len(streamer.finalized) != tt.wantFinalizedResponses { + t.Fatalf("stream finalized = %v, want %d responses", streamer.finalized, tt.wantFinalizedResponses) + } + if tt.wantChatCalls == 1 && got != "chat response" { + t.Fatalf("response = %q, want chat response", got) + } + if tt.wantStreamCalls == 1 && got != "stream response" { + t.Fatalf("response = %q, want stream response", got) + } + wantResolvedModel := "openai/" + tt.rewriteModel + if tt.wantStreamCalls == 1 && + (len(provider.streamModels) != 1 || provider.streamModels[0] != wantResolvedModel) { + t.Fatalf("stream models = %v, want [%s]", provider.streamModels, wantResolvedModel) + } + if tt.wantChatCalls == 1 && (len(provider.chatModels) != 1 || provider.chatModels[0] != wantResolvedModel) { + t.Fatalf("chat models = %v, want [%s]", provider.chatModels, wantResolvedModel) + } + }) + } +} + +func TestConfiguredStreamingPostChunkFailureDoesNotFallBackToChat(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"partial"}, + err: errors.New("stream failed after chunk"), + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + configuredStreamingProcessOptions("pico"), + ) + if err == nil { + t.Fatal("expected post-chunk stream failure to return an error") + } + if provider.streamCalls != 1 || provider.chatCalls != 0 { + t.Fatalf("calls = stream:%d chat:%d, want stream:1 chat:0", provider.streamCalls, provider.chatCalls) + } + if len(streamer.updates) != 1 || streamer.updates[0] != "partial" { + t.Fatalf("stream updates = %v, want [partial]", streamer.updates) + } + if streamer.canceled != 0 { + t.Fatalf("streamer canceled = %d, want 0 for already-visible stream failure", streamer.canceled) + } +} + +func TestConfiguredStreamingPostChunkEOFDoesNotRetryOrCancelVisibleOutput(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"partial"}, + err: io.EOF, + }}, + chatResponse: &providers.LLMResponse{Content: "chat retry"}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + configuredStreamingProcessOptions("pico"), + ) + if err == nil { + t.Fatal("expected post-chunk EOF to return an error") + } + if provider.streamCalls != 1 || provider.chatCalls != 0 { + t.Fatalf("calls = stream:%d chat:%d, want stream:1 chat:0", provider.streamCalls, provider.chatCalls) + } + if len(streamer.updates) != 1 || streamer.updates[0] != "partial" { + t.Fatalf("stream updates = %v, want [partial]", streamer.updates) + } + if streamer.canceled != 0 { + t.Fatalf("streamer canceled = %d, want 0 for already-visible stream EOF", streamer.canceled) + } +} + +func TestConfiguredStreamingFinalizesAfterAfterLLMHookMutation(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"partial"}, + response: &providers.LLMResponse{Content: "original streamed response"}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + if err := al.MountHook(NamedHook("rewrite-stream-response", configuredStreamingAfterHook{ + content: "hooked final response", + })); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != "hooked final response" { + t.Fatalf("response = %q, want hook-modified response", got) + } + if len(streamer.finalized) != 1 || streamer.finalized[0] != "hooked final response" { + t.Fatalf("stream finalized = %v, want [hooked final response]", streamer.finalized) + } +} + +func TestConfiguredStreamingAfterLLMAbortCancelsPublishedStream(t *testing.T) { + tests := []struct { + name string + action HookAction + }{ + {name: "abort turn", action: HookActionAbortTurn}, + {name: "hard abort", action: HookActionHardAbort}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"partial before abort"}, + response: &providers.LLMResponse{Content: "should not be visible"}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + if err := al.MountHook(NamedHook("abort-stream-response", configuredStreamingAfterHook{ + action: tt.action, + })); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + _, _ = al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + configuredStreamingProcessOptions("pico"), + ) + + if streamer.canceled != 1 { + t.Fatalf("streamer canceled = %d, want 1", streamer.canceled) + } + if len(streamer.finalized) != 0 { + t.Fatalf("stream finalized = %v, want none", streamer.finalized) + } + }) + } +} + +func TestConfiguredStreamingFinalizesWithDefaultResponseWhenContentEmpty(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{{ + chunks: []string{"partial response"}, + response: &providers.LLMResponse{}, + }}, + } + al := NewAgentLoop(cfg, msgBus, provider) + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != defaultResponse { + t.Fatalf("response = %q, want default response", got) + } + if len(streamer.finalized) != 1 || streamer.finalized[0] != defaultResponse { + t.Fatalf("stream finalized = %v, want [%q]", streamer.finalized, defaultResponse) + } +} + +func TestConfiguredStreamingToolCallsUseCompleteStreamResponse(t *testing.T) { + cfg := newConfiguredStreamingTestConfig(t, true, true, nil) + streamer := &recordingStreamer{} + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(configuredStreamingDelegate{streamer: streamer}) + provider := &configuredStreamingProvider{ + streamPlan: []configuredStreamingCall{ + { + chunks: []string{"partial tool-call response"}, + response: &providers.LLMResponse{ + Content: "need a tool", + ToolCalls: []providers.ToolCall{{ + ID: "call-1", + Type: "function", + Name: "tool_limit_test_tool", + Arguments: map[string]any{"value": "x"}, + }}, + }, + }, + { + response: &providers.LLMResponse{Content: "tool call handled"}, + }, + }, + } + al := NewAgentLoop(cfg, msgBus, provider) + agent := al.GetRegistry().GetDefaultAgent() + agent.Tools.Register(&toolLimitTestTool{}) + + got := runConfiguredStreamingTurn(t, al, "pico") + + if got != "tool call handled" { + t.Fatalf("response = %q, want tool call handled", got) + } + if provider.streamCalls != 2 { + t.Fatalf("ChatStream calls = %d, want 2", provider.streamCalls) + } + if provider.chatCalls != 0 { + t.Fatalf("Chat calls = %d, want 0", provider.chatCalls) + } + if streamer.canceled != 1 { + t.Fatalf("streamer canceled = %d, want 1 for non-final tool-call response", streamer.canceled) + } + if len(streamer.finalized) != 1 || streamer.finalized[0] != "tool call handled" { + t.Fatalf("stream finalized = %v, want [tool call handled]", streamer.finalized) + } +} + +func newConfiguredStreamingTestConfig( + t *testing.T, + channelStreaming bool, + modelStreaming bool, + fallbacks []string, +) *config.Config { + t.Helper() + tmpDir, err := os.MkdirTemp("", "configured-streaming-agent-test-*") + if err != nil { + t.Fatalf("MkdirTemp() error = %v", err) + } + t.Cleanup(func() { os.RemoveAll(tmpDir) }) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + ModelFallbacks: append([]string(nil), fallbacks...), + MaxTokens: 4096, + MaxToolIterations: 3, + }, + }, + Channels: config.ChannelsConfig{ + "pico": newConfiguredStreamingPicoChannel(t, channelStreaming), + "wecom": newConfiguredStreamingWeComChannel(t, channelStreaming), + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Provider: "openai", + Model: "openai/test-model", + Streaming: config.ModelStreamingConfig{Enabled: modelStreaming}, + }}, + } + if len(fallbacks) > 0 { + cfg.ModelList = append(cfg.ModelList, &config.ModelConfig{ + ModelName: "fallback-model", + Provider: "openai", + Model: "openai/fallback-model", + Streaming: config.ModelStreamingConfig{Enabled: true}, + }) + } + if err := config.InitChannelList(cfg.Channels); err != nil { + t.Fatalf("InitChannelList() error = %v", err) + } + return cfg +} + +func newConfiguredStreamingWeComChannel(t *testing.T, enabled bool) *config.Channel { + t.Helper() + settings := config.WeComSettings{ + BotID: "bot-1", + Streaming: config.StreamingConfig{ + Enabled: enabled, + ThrottleSeconds: 1, + MinGrowthChars: 40, + }, + } + raw, err := json.Marshal(settings) + if err != nil { + t.Fatalf("Marshal settings error = %v", err) + } + return &config.Channel{ + Type: config.ChannelWeCom, + Enabled: true, + Settings: config.RawNode(raw), + } +} + +func newConfiguredStreamingPicoChannel(t *testing.T, enabled bool) *config.Channel { + t.Helper() + settings := config.PicoSettings{ + Streaming: config.StreamingConfig{ + Enabled: enabled, + ThrottleSeconds: 1, + MinGrowthChars: 40, + }, + } + settings.SetToken("test-token") + raw, err := json.Marshal(settings) + if err != nil { + t.Fatalf("Marshal settings error = %v", err) + } + return &config.Channel{ + Type: config.ChannelPico, + Enabled: true, + Settings: config.RawNode(raw), + } +} + +func configuredStreamingProcessOptions(channel string) processOptions { + return processOptions{ + SessionKey: "agent:main:" + channel + ":session-1", + Channel: channel, + ChatID: "session-1", + UserMessage: "hello", + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + AllowInterimPicoPublish: true, + NoHistory: true, + } +} + +func runConfiguredStreamingTurn(t *testing.T, al *AgentLoop, channel string) string { + t.Helper() + got, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + configuredStreamingProcessOptions(channel), + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + return got +} diff --git a/pkg/agent/prompt.go b/pkg/agent/prompt.go index 02c85036..1c66fe69 100644 --- a/pkg/agent/prompt.go +++ b/pkg/agent/prompt.go @@ -115,6 +115,13 @@ type PromptBuildRequest struct { ActiveSkills []string Overlays []PromptPart + + SuppressDefaultSystemPrompt bool + SuppressSkillContext bool + SuppressToolUseRule bool + AllowedSkills []string + AllowedTools []string + ToolUseFallback bool } type PromptContributor interface { diff --git a/pkg/agent/prompt_contributors.go b/pkg/agent/prompt_contributors.go index d6a2c09e..3db5b851 100644 --- a/pkg/agent/prompt_contributors.go +++ b/pkg/agent/prompt_contributors.go @@ -4,6 +4,8 @@ import ( "context" "fmt" "strings" + + "github.com/sipeed/picoclaw/pkg/tools" ) type toolDiscoveryPromptContributor struct { @@ -23,9 +25,17 @@ func (c toolDiscoveryPromptContributor) PromptSource() PromptSourceDescriptor { func (c toolDiscoveryPromptContributor) ContributePrompt( _ context.Context, - _ PromptBuildRequest, + req PromptBuildRequest, ) ([]PromptPart, error) { - content := formatToolDiscoveryRule(c.useBM25, c.useRegex) + if req.SuppressToolUseRule { + return nil, nil + } + useBM25 := c.useBM25 && promptAllowsTool(req, tools.BM25SearchToolName) + useRegex := c.useRegex && promptAllowsTool(req, tools.RegexSearchToolName) + if !useBM25 && !useRegex { + return nil, nil + } + content := formatToolDiscoveryRule(useBM25, useRegex) if strings.TrimSpace(content) == "" { return nil, nil } @@ -62,12 +72,19 @@ func (c mcpServerPromptContributor) PromptSource() PromptSourceDescriptor { func (c mcpServerPromptContributor) ContributePrompt( _ context.Context, - _ PromptBuildRequest, + req PromptBuildRequest, ) ([]PromptPart, error) { + if req.SuppressToolUseRule { + return nil, nil + } serverName := strings.TrimSpace(c.serverName) if serverName == "" || c.toolCount <= 0 { return nil, nil } + if len(req.AllowedTools) > 0 && + !promptAllowsToolPrefix(req, "mcp_"+promptSourceComponent(serverName)+"_") { + return nil, nil + } availability := "available as native tools" if c.deferred { @@ -110,8 +127,14 @@ func (c agentDiscoveryPromptContributor) PromptSource() PromptSourceDescriptor { func (c agentDiscoveryPromptContributor) ContributePrompt( _ context.Context, - _ PromptBuildRequest, + req PromptBuildRequest, ) ([]PromptPart, error) { + if req.SuppressToolUseRule { + return nil, nil + } + if !promptAllowsTool(req, "spawn") { + return nil, nil + } if c.discover == nil { return nil, nil } @@ -178,3 +201,28 @@ func promptSourceComponent(value string) string { } return result } + +func promptAllowsTool(req PromptBuildRequest, name string) bool { + if len(req.AllowedTools) == 0 { + return true + } + allowed := cleanAllowedSet(req.AllowedTools) + _, ok := allowed[strings.ToLower(strings.TrimSpace(name))] + return ok +} + +func promptAllowsToolPrefix(req PromptBuildRequest, prefix string) bool { + if len(req.AllowedTools) == 0 { + return true + } + prefix = strings.ToLower(strings.TrimSpace(prefix)) + if prefix == "" { + return false + } + for _, name := range req.AllowedTools { + if strings.HasPrefix(strings.ToLower(strings.TrimSpace(name)), prefix) { + return true + } + } + return false +} diff --git a/pkg/agent/prompt_test.go b/pkg/agent/prompt_test.go index b76b0040..605e5f4a 100644 --- a/pkg/agent/prompt_test.go +++ b/pkg/agent/prompt_test.go @@ -197,6 +197,75 @@ func TestContextBuilder_CollectsToolDiscoveryContributor(t *testing.T) { } } +func TestContextBuilder_SuppressesToolDiscoveryContributorWhenToolsUnavailable(t *testing.T) { + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + cb := NewContextBuilder(t.TempDir()).WithToolDiscovery(true, false) + + messages := cb.BuildMessagesFromPrompt(PromptBuildRequest{ + CurrentMessage: "hello", + SuppressToolUseRule: true, + }) + system := messages[0] + if strings.Contains(system.Content, "tool_search_tool_bm25") { + t.Fatalf("system prompt includes tool discovery despite tools being unavailable: %q", system.Content) + } + for _, part := range system.SystemParts { + if part.PromptSource == string(PromptSourceToolDiscovery) { + t.Fatalf("system parts include tool discovery despite tools being unavailable: %#v", part) + } + } +} + +func TestContextBuilder_SuppressesToolReferencesWhenToolsUnavailable(t *testing.T) { + workspace := t.TempDir() + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + writeTurnProfileSkill( + t, + workspace, + "research", + "---\ndescription: research skill\n---\n# research\n\nResearch carefully.", + ) + cb := NewContextBuilder(workspace) + + messages := cb.BuildMessagesFromPrompt(PromptBuildRequest{ + CurrentMessage: "hello", + SuppressToolUseRule: true, + }) + system := messages[0] + if strings.Contains(system.Content, "When using tools") || + strings.Contains(system.Content, "read_file tool") || + strings.Contains(system.Content, "update "+workspace+"/memory/MEMORY.md") { + t.Fatalf("system prompt includes tool references despite tools being unavailable: %q", system.Content) + } + if !strings.Contains(system.Content, "research") { + t.Fatalf("system prompt should keep non-tool skill catalog context, got: %q", system.Content) + } +} + +func TestContextBuilder_CustomToolAllowListSuppressesReadFileSkillInstruction(t *testing.T) { + workspace := t.TempDir() + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + writeTurnProfileSkill( + t, + workspace, + "research", + "---\ndescription: research skill\n---\n# research\n\nResearch carefully.", + ) + cb := NewContextBuilder(workspace) + + messages := cb.BuildMessagesFromPrompt(PromptBuildRequest{ + CurrentMessage: "hello", + AllowedTools: []string{"web_search"}, + }) + system := messages[0] + if strings.Contains(system.Content, "read_file tool") { + t.Fatalf("system prompt includes read_file skill instruction without read_file permission: %q", system.Content) + } + if !strings.Contains(system.Content, "research") { + t.Fatalf("system prompt should keep skill catalog context, got: %q", system.Content) + } +} + func TestContextBuilder_CollectsMCPServerContributor(t *testing.T) { t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) cb := NewContextBuilder(t.TempDir()) @@ -232,6 +301,111 @@ func TestContextBuilder_CollectsMCPServerContributor(t *testing.T) { } } +func TestContextBuilder_SuppressesMCPServerContributorWhenToolsUnavailable(t *testing.T) { + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + cb := NewContextBuilder(t.TempDir()) + err := cb.RegisterPromptContributor(mcpServerPromptContributor{ + serverName: "GitHub Server", + toolCount: 3, + deferred: false, + }) + if err != nil { + t.Fatalf("RegisterPromptContributor() error = %v", err) + } + + messages := cb.BuildMessagesFromPrompt(PromptBuildRequest{ + CurrentMessage: "hello", + SuppressToolUseRule: true, + }) + system := messages[0] + if strings.Contains(system.Content, "MCP server `GitHub Server` is connected") || + strings.Contains(system.Content, "available as native tools") { + t.Fatalf("system prompt includes MCP tooling despite tools being unavailable: %q", system.Content) + } + for _, part := range system.SystemParts { + if part.PromptSource == "mcp:github_server" { + t.Fatalf("system parts include MCP tooling despite tools being unavailable: %#v", part) + } + } +} + +func TestContextBuilder_SuppressesAgentDiscoveryContributorWhenToolsUnavailable(t *testing.T) { + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + cb := NewContextBuilder(t.TempDir()).WithAgentDiscovery( + "main", + func(agentID string) []AgentDescriptor { + return []AgentDescriptor{{ + ID: "helper", + Name: "Helper", + Description: "Helps with tasks", + }} + }, + ) + + messages := cb.BuildMessagesFromPrompt(PromptBuildRequest{ + CurrentMessage: "hello", + SuppressToolUseRule: true, + }) + system := messages[0] + if strings.Contains(system.Content, "Agent Discovery") || + strings.Contains(system.Content, "calling spawn") { + t.Fatalf("system prompt includes agent discovery despite tools being unavailable: %q", system.Content) + } + for _, part := range system.SystemParts { + if part.PromptSource == string(PromptSourceAgentDiscovery) { + t.Fatalf("system parts include agent discovery despite tools being unavailable: %#v", part) + } + } +} + +func TestContextBuilder_CustomToolAllowListSuppressesUnallowedToolContributors(t *testing.T) { + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + cb := NewContextBuilder(t.TempDir()). + WithToolDiscovery(true, true). + WithAgentDiscovery( + "main", + func(agentID string) []AgentDescriptor { + return []AgentDescriptor{{ + ID: "helper", + Name: "Helper", + Description: "Helps with tasks", + }} + }, + ) + err := cb.RegisterPromptContributor(mcpServerPromptContributor{ + serverName: "GitHub Server", + toolCount: 3, + deferred: false, + }) + if err != nil { + t.Fatalf("RegisterPromptContributor() error = %v", err) + } + + messages := cb.BuildMessagesFromPrompt(PromptBuildRequest{ + CurrentMessage: "hello", + AllowedTools: []string{"echo_text"}, + }) + system := messages[0] + blockedSnippets := []string{ + "tool_search_tool_bm25", + "tool_search_tool_regex", + "MCP server `GitHub Server` is connected", + "Agent Discovery", + "calling spawn", + } + for _, snippet := range blockedSnippets { + if strings.Contains(system.Content, snippet) { + t.Fatalf("system prompt includes unallowed tool contributor %q: %q", snippet, system.Content) + } + } + for _, part := range system.SystemParts { + switch part.PromptSource { + case string(PromptSourceToolDiscovery), string(PromptSourceAgentDiscovery), "mcp:github_server": + t.Fatalf("system parts include unallowed tool contributor: %#v", part) + } + } +} + type testPromptContributor struct { desc PromptSourceDescriptor part PromptPart diff --git a/pkg/agent/prompt_turn.go b/pkg/agent/prompt_turn.go index 588a8f00..053eb846 100644 --- a/pkg/agent/prompt_turn.go +++ b/pkg/agent/prompt_turn.go @@ -4,6 +4,7 @@ import ( "fmt" "strings" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers" ) @@ -13,8 +14,9 @@ func promptBuildRequestForTurn( summary string, currentMessage string, media []string, + cfg *config.Config, ) PromptBuildRequest { - return PromptBuildRequest{ + req := PromptBuildRequest{ History: history, Summary: summary, CurrentMessage: currentMessage, @@ -26,6 +28,92 @@ func promptBuildRequestForTurn( ActiveSkills: activeSkillNames(ts.agent, ts.opts), Overlays: promptOverlaysForOptions(ts.opts), } + hasCallableTools := true + if ts.profile.Enabled { + hasCallableTools = turnProfileHasCallableTools(ts.profile, ts.agent.Tools.ToProviderDefs()) || + turnProfileNativeSearchCallable(cfg, ts.profile, ts.agent) + } + if turnProfileSystemPromptOff(ts.profile) { + req.SuppressDefaultSystemPrompt = true + req.SuppressSkillContext = true + req.ToolUseFallback = hasCallableTools + } + if ts.profile.Enabled && !hasCallableTools { + req.SuppressToolUseRule = true + } + if turnProfileSkillsOff(ts.profile) { + req.SuppressSkillContext = true + } + if turnProfileCustomSkills(ts.profile) { + req.AllowedSkills = append([]string(nil), ts.profile.AllowedSkills...) + } + if ts.profile.Enabled && ts.profile.ToolsMode == config.TurnProfileModeCustom { + req.AllowedTools = append([]string(nil), ts.profile.AllowedTools...) + } + return req +} + +func turnProfileNativeSearchCallable( + cfg *config.Config, + profile config.EffectiveTurnProfile, + agent *AgentInstance, +) bool { + if cfg == nil || agent == nil { + return false + } + if !cfg.Tools.IsToolEnabled("web") || !cfg.Tools.Web.PreferNative { + return false + } + if !turnProfileToolAllowed(profile, "web_search") { + return false + } + nativeProvider, ok := agent.Provider.(providers.NativeSearchCapable) + return ok && nativeProvider.SupportsNativeSearch() +} + +func promptBuildRequestForProcessOptions( + agent *AgentInstance, + opts processOptions, + history []providers.Message, + summary string, + currentMessage string, + media []string, +) PromptBuildRequest { + req := PromptBuildRequest{ + History: history, + Summary: summary, + CurrentMessage: currentMessage, + Media: append([]string(nil), media...), + Channel: opts.Channel, + ChatID: opts.ChatID, + SenderID: opts.SenderID, + SenderDisplayName: opts.SenderDisplayName, + ActiveSkills: activeSkillNames(agent, opts), + Overlays: promptOverlaysForOptions(opts), + } + profile := opts.TurnProfile + hasCallableTools := true + if profile.Enabled && agent != nil { + hasCallableTools = turnProfileHasCallableTools(profile, agent.Tools.ToProviderDefs()) + } + if turnProfileSystemPromptOff(profile) { + req.SuppressDefaultSystemPrompt = true + req.SuppressSkillContext = true + req.ToolUseFallback = hasCallableTools + } + if profile.Enabled && !hasCallableTools { + req.SuppressToolUseRule = true + } + if turnProfileSkillsOff(profile) { + req.SuppressSkillContext = true + } + if turnProfileCustomSkills(profile) { + req.AllowedSkills = append([]string(nil), profile.AllowedSkills...) + } + if profile.Enabled && profile.ToolsMode == config.TurnProfileModeCustom { + req.AllowedTools = append([]string(nil), profile.AllowedTools...) + } + return req } func promptOverlaysForOptions(opts processOptions) []PromptPart { diff --git a/pkg/agent/steering_test.go b/pkg/agent/steering_test.go index 23d34840..569f2058 100644 --- a/pkg/agent/steering_test.go +++ b/pkg/agent/steering_test.go @@ -1618,6 +1618,8 @@ func TestAgentLoop_InterruptHard_RestoresSession(t *testing.T) { } defaultAgent.Sessions.SetHistory(sessionKey, originalHistory) + originalHistory = defaultAgent.Sessions.GetHistory(sessionKey) + runtimeCh, closeRuntimeEvents := subscribeRuntimeEventsForTest( t, al, diff --git a/pkg/agent/subturn.go b/pkg/agent/subturn.go index 86617d02..fb9d9406 100644 --- a/pkg/agent/subturn.go +++ b/pkg/agent/subturn.go @@ -378,6 +378,7 @@ func spawnSubTurn( Dispatch: dispatch, SenderID: parentTS.opts.Dispatch.SenderID(), SenderDisplayName: parentTS.opts.SenderDisplayName, + TurnProfile: parentTS.profile, SystemPromptOverride: cfg.ActualSystemPrompt, InitialSteeringMessages: cfg.InitialMessages, DefaultResponse: "", @@ -386,6 +387,9 @@ func spawnSubTurn( NoHistory: true, // SubTurns don't use session history SkipInitialSteeringPoll: true, } + if !opts.TurnProfile.Enabled { + opts.TurnProfile = parentTS.opts.TurnProfile + } // Create event scope for the child turn scope := al.newTurnEventScope( diff --git a/pkg/agent/thinking.go b/pkg/agent/thinking.go index 015b6928..087f5136 100644 --- a/pkg/agent/thinking.go +++ b/pkg/agent/thinking.go @@ -1,6 +1,12 @@ package agent -import "strings" +import ( + "strings" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/providers" +) // ThinkingLevel controls how the provider sends thinking parameters. // @@ -37,3 +43,87 @@ func parseThinkingLevel(level string) ThinkingLevel { return ThinkingOff } } + +func isConfiguredThinkingLevel(level string) bool { + switch strings.ToLower(strings.TrimSpace(level)) { + case "off", "low", "medium", "high", "xhigh", "adaptive": + return true + default: + return false + } +} + +type thinkingSettings struct { + level ThinkingLevel + configured bool +} + +func thinkingSettingsFromModelConfig(mc *config.ModelConfig) thinkingSettings { + if mc == nil || !isConfiguredThinkingLevel(mc.ThinkingLevel) { + return thinkingSettings{} + } + return thinkingSettings{ + level: parseThinkingLevel(mc.ThinkingLevel), + configured: true, + } +} + +func activeThinkingSettings(agent *AgentInstance, modelCfg *config.ModelConfig) thinkingSettings { + if settings := thinkingSettingsFromModelConfig(modelCfg); settings.configured { + return settings + } + if modelCfg == nil && agent != nil { + return thinkingSettings{ + level: agent.ThinkingLevel, + configured: agent.ThinkingLevelConfigured, + } + } + return thinkingSettings{} +} + +func applyThinkingOption( + opts map[string]any, + provider providers.LLMProvider, + settings thinkingSettings, + warnUnsupported bool, + agentID string, +) { + if opts == nil || !settings.configured { + return + } + if settings.level == ThinkingOff { + opts["thinking_level"] = string(settings.level) + return + } + if tc, ok := provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() { + opts["thinking_level"] = string(settings.level) + return + } + if warnUnsupported { + logger.WarnCF("agent", "thinking_level is set but current provider does not support it, ignoring", + map[string]any{"agent_id": agentID, "thinking_level": string(settings.level)}) + } +} + +func applyTurnThinkingOptions( + exec *turnExecution, + agent *AgentInstance, + provider providers.LLMProvider, + warnUnsupported bool, +) { + if exec == nil || exec.llmOpts == nil { + return + } + delete(exec.llmOpts, "thinking_level") + settings := activeThinkingSettings(agent, exec.activeModelConfig) + agentID := "" + if agent != nil { + agentID = agent.ID + } + applyThinkingOption(exec.llmOpts, provider, settings, warnUnsupported, agentID) + exec.suppressReasoning = shouldSuppressReasoningFor(settings) +} + +func shouldSuppressReasoningFor(settings thinkingSettings) bool { + return settings.configured && settings.level == ThinkingOff +} diff --git a/pkg/agent/turn_coord.go b/pkg/agent/turn_coord.go index 06034633..4e14c759 100644 --- a/pkg/agent/turn_coord.go +++ b/pkg/agent/turn_coord.go @@ -373,6 +373,11 @@ func (al *AgentLoop) askSideQuestion( if opts != nil { normalizeProcessOptionsInPlace(opts) + resolved, err := resolveTurnProfileOptions(al.GetConfig(), *opts) + if err != nil { + return "", err + } + *opts = resolved } var media []string @@ -399,16 +404,31 @@ func (al *AgentLoop) askSideQuestion( } } - messages := agent.ContextBuilder.BuildMessages( - history, - summary, - question, - media, - channel, - chatID, - senderID, - senderDisplayName, - ) + var promptReq PromptBuildRequest + if opts == nil { + promptReq = PromptBuildRequest{ + History: history, + Summary: summary, + CurrentMessage: question, + Media: append([]string(nil), media...), + Channel: channel, + ChatID: chatID, + SenderID: senderID, + SenderDisplayName: senderDisplayName, + } + } else { + promptReq = promptBuildRequestForProcessOptions( + agent, + *opts, + history, + summary, + question, + media, + ) + } + promptReq.SuppressToolUseRule = true + promptReq.ToolUseFallback = false + messages := agent.ContextBuilder.BuildMessagesFromPrompt(promptReq) maxMediaSize := al.GetConfig().Agents.Defaults.GetMaxMediaSize() messages = resolveMediaRefs(messages, al.mediaStore, maxMediaSize) @@ -423,6 +443,7 @@ func (al *AgentLoop) askSideQuestion( } hookModelChanged := false + sideSuppressReasoning := false callProvider := func( ctx context.Context, candidate providers.FallbackCandidate, @@ -430,7 +451,15 @@ func (al *AgentLoop) askSideQuestion( forceModel bool, callMessages []providers.Message, ) (*providers.LLMResponse, error) { - provider, providerModel, cleanup, err := al.isolatedSideQuestionProvider(agent, selectedModelName, candidate) + baseModelName := selectedModelName + if forceModel && strings.TrimSpace(model) != "" { + baseModelName = model + } + provider, providerModel, modelCfg, cleanup, err := al.isolatedSideQuestionProvider( + agent, + baseModelName, + candidate, + ) if err != nil { return nil, err } @@ -439,10 +468,12 @@ func (al *AgentLoop) askSideQuestion( model = providerModel } callOpts := llmOpts - if _, exists := callOpts["thinking_level"]; !exists && agent.ThinkingLevel != ThinkingOff { - if tc, ok := provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() { + settings := thinkingSettingsFromModelConfig(modelCfg) + sideSuppressReasoning = shouldSuppressReasoningFor(settings) + if _, exists := callOpts["thinking_level"]; !exists { + if settings.configured { callOpts = shallowCloneLLMOptions(llmOpts) - callOpts["thinking_level"] = string(agent.ThinkingLevel) + applyThinkingOption(callOpts, provider, settings, false, agent.ID) } } return provider.Chat(ctx, callMessages, nil, model, callOpts) @@ -476,6 +507,7 @@ func (al *AgentLoop) askSideQuestion( llmModel = llmReq.Model messages = llmReq.Messages llmOpts = llmReq.Options + delete(llmOpts, "native_search") } case HookActionAbortTurn: reason := decision.Reason @@ -500,18 +532,11 @@ func (al *AgentLoop) askSideQuestion( callSideLLM := func(callMessages []providers.Message) (*providers.LLMResponse, error) { if len(activeCandidates) > 1 && al.fallback != nil { - fbResult, err := al.fallback.Execute( + fbResult, err := al.fallback.ExecuteCandidate( ctx, activeCandidates, - func(ctx context.Context, providerName, model string) (*providers.LLMResponse, error) { - candidate := providers.FallbackCandidate{Provider: providerName, Model: model} - for _, activeCandidate := range activeCandidates { - if activeCandidate.Provider == providerName && activeCandidate.Model == model { - candidate = activeCandidate - break - } - } - return callProvider(ctx, candidate, model, false, callMessages) + func(ctx context.Context, candidate providers.FallbackCandidate) (*providers.LLMResponse, error) { + return callProvider(ctx, candidate, candidate.Model, false, callMessages) }, ) if err != nil { @@ -584,6 +609,11 @@ func (al *AgentLoop) askSideQuestion( return "", fmt.Errorf("hook aborted turn during after_llm: %s", reason) } } + if sideSuppressReasoning { + resp.Reasoning = "" + resp.ReasoningContent = "" + resp.ReasoningDetails = nil + } return sideQuestionResponseContent(resp), nil } @@ -592,14 +622,14 @@ func (al *AgentLoop) isolatedSideQuestionProvider( agent *AgentInstance, baseModelName string, candidate providers.FallbackCandidate, -) (providers.LLMProvider, string, func(), error) { +) (providers.LLMProvider, string, *config.ModelConfig, func(), error) { if agent == nil { - return nil, "", func() {}, fmt.Errorf("isolatedSideQuestionProvider: no agent available for /btw") + return nil, "", nil, func() {}, fmt.Errorf("isolatedSideQuestionProvider: no agent available for /btw") } modelCfg, err := al.sideQuestionModelConfig(agent, baseModelName, candidate) if err != nil { - return nil, "", func() {}, fmt.Errorf("isolatedSideQuestionProvider: %w", err) + return nil, "", nil, func() {}, fmt.Errorf("isolatedSideQuestionProvider: %w", err) } factory := al.providerFactory @@ -608,13 +638,13 @@ func (al *AgentLoop) isolatedSideQuestionProvider( } provider, modelID, err := factory(modelCfg) if err != nil { - return nil, "", func() {}, fmt.Errorf("isolatedSideQuestionProvider: %w", err) + return nil, "", nil, func() {}, fmt.Errorf("isolatedSideQuestionProvider: %w", err) } cleanup := func() { closeProviderIfStateful(provider) } - return provider, modelID, cleanup, nil + return provider, modelID, modelCfg, cleanup, nil } func (al *AgentLoop) sideQuestionModelConfig( @@ -626,7 +656,15 @@ func (al *AgentLoop) sideQuestionModelConfig( return nil, fmt.Errorf("sideQuestionModelConfig: no agent available for /btw") } - // If candidate has an identity key, use that + if name := modelAliasFromCandidateIdentityKey(candidate.IdentityKey); name != "" { + modelCfg, err := resolvedModelConfig(al.GetConfig(), name, agent.Workspace) + if err == nil { + return modelCfg, nil + } + // Fallback: create a minimal config if lookup fails + } + + // Older identity keys used provider/model; keep resolving those by model. if name := modelNameFromIdentityKey(candidate.IdentityKey); name != "" { modelCfg, err := resolvedModelConfig(al.GetConfig(), name, agent.Workspace) if err == nil { @@ -635,6 +673,18 @@ func (al *AgentLoop) sideQuestionModelConfig( // Fallback: create a minimal config if lookup fails } + if candidate.Provider != "" && candidate.Model != "" { + candidateRef := providers.NormalizeProvider(candidate.Provider) + "/" + candidate.Model + if modelCfg, err := resolvedModelConfig(al.GetConfig(), candidateRef, agent.Workspace); err == nil { + return modelCfg, nil + } + return &config.ModelConfig{ + ModelName: candidateRef, + Model: candidateRef, + Workspace: agent.Workspace, + }, nil + } + // Otherwise, clean up the base model name and use it baseModelName = strings.TrimSpace(baseModelName) modelCfg, err := resolvedModelConfig(al.GetConfig(), baseModelName, agent.Workspace) @@ -658,8 +708,5 @@ func (al *AgentLoop) sideQuestionModelConfig( // If candidate specifies a different provider/model, override clone := *modelCfg - if candidate.Provider != "" && candidate.Model != "" { - clone.Model = providers.NormalizeProvider(candidate.Provider) + "/" + candidate.Model - } return &clone, nil } diff --git a/pkg/agent/turn_coord_test.go b/pkg/agent/turn_coord_test.go index c7cdd8a3..dc2715af 100644 --- a/pkg/agent/turn_coord_test.go +++ b/pkg/agent/turn_coord_test.go @@ -10,6 +10,7 @@ import ( "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/session" ) @@ -37,8 +38,43 @@ func (p *simpleConvProvider) GetDefaultModel() string { return "simple-model" } +type sequenceProvider struct { + responses []*providers.LLMResponse + errors []error + callCount int + mu sync.Mutex +} + +func (p *sequenceProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.mu.Lock() + defer p.mu.Unlock() + + idx := p.callCount + p.callCount++ + + if idx < len(p.errors) && p.errors[idx] != nil { + return nil, p.errors[idx] + } + if idx < len(p.responses) && p.responses[idx] != nil { + return p.responses[idx], nil + } + return &providers.LLMResponse{Content: "ok", FinishReason: "stop"}, nil +} + +func (p *sequenceProvider) GetDefaultModel() string { + return "sequence-model" +} + type nativeSearchCaptureProvider struct { lastOpts map[string]any + messages []providers.Message + tools []providers.ToolDefinition } func (p *nativeSearchCaptureProvider) Chat( @@ -48,6 +84,8 @@ func (p *nativeSearchCaptureProvider) Chat( model string, opts map[string]any, ) (*providers.LLMResponse, error) { + p.messages = append([]providers.Message(nil), messages...) + p.tools = append([]providers.ToolDefinition(nil), tools...) p.lastOpts = make(map[string]any, len(opts)) for k, v := range opts { p.lastOpts[k] = v @@ -271,6 +309,152 @@ func TestPipeline_CallLLM_SimpleResponse(t *testing.T) { } } +func TestPipeline_SetupTurn_ModelNameDoesNotUseFallbackAliasBeforeFallback(t *testing.T) { + al, agent, cleanup := newTurnCoordTestLoop(t, &simpleConvProvider{}) + defer cleanup() + + agent.Model = "primary-model" + agent.Candidates = []providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4"}, + {Provider: "anthropic", Model: "claude-sonnet", IdentityKey: "model_name:fallback-model"}, + } + + pipeline := NewPipeline(al) + ts := newTurnState(agent, makeTestProcessOpts("test-session"), turnEventScope{ + turnID: "turn-1", + context: newTurnContext(nil, nil, nil), + }) + + exec, err := pipeline.SetupTurn(context.Background(), ts) + if err != nil { + t.Fatalf("SetupTurn failed: %v", err) + } + if exec.llmModelName != "primary-model" { + t.Fatalf("exec.llmModelName = %q, want %q", exec.llmModelName, "primary-model") + } +} + +func TestPipeline_CallLLM_UsesSuccessfulFallbackIdentityAlias(t *testing.T) { + provider := &sequenceProvider{ + errors: []error{ + errors.New("status: 429 - rate limit exceeded"), + nil, + }, + responses: []*providers.LLMResponse{ + nil, + {Content: "fallback answer", FinishReason: "stop"}, + }, + } + al, agent, cleanup := newTurnCoordTestLoop(t, provider) + defer cleanup() + + agent.Model = "primary-model" + agent.Candidates = []providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4", IdentityKey: "model_name:primary"}, + {Provider: "openai", Model: "gpt-5.4", IdentityKey: "model_name:secondary"}, + } + al.fallback = providers.NewFallbackChain(providers.NewCooldownTracker(), nil) + + pipeline := NewPipeline(al) + ts := newTurnState(agent, makeTestProcessOpts("test-session"), turnEventScope{ + turnID: "turn-1", + context: newTurnContext(nil, nil, nil), + }) + + exec, err := pipeline.SetupTurn(context.Background(), ts) + if err != nil { + t.Fatalf("SetupTurn failed: %v", err) + } + + ctrl, err := pipeline.CallLLM(context.Background(), context.Background(), ts, exec, 1) + if err != nil { + t.Fatalf("CallLLM failed: %v", err) + } + if ctrl != ControlBreak { + t.Fatalf("expected ControlBreak, got %v", ctrl) + } + if exec.llmModelName != "secondary" { + t.Fatalf("exec.llmModelName = %q, want %q", exec.llmModelName, "secondary") + } +} + +func TestPipeline_CallLLM_UsesSuccessfulFallbackDisplayNameWithoutAlias(t *testing.T) { + provider := &sequenceProvider{ + errors: []error{ + errors.New("status: 429 - rate limit exceeded"), + nil, + }, + responses: []*providers.LLMResponse{ + nil, + {Content: "fallback answer", FinishReason: "stop"}, + }, + } + al, agent, cleanup := newTurnCoordTestLoop(t, provider) + defer cleanup() + + agent.Model = "primary-model" + agent.Candidates = []providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4", IdentityKey: "model_name:primary", DisplayName: "primary-model"}, + {Provider: "anthropic", Model: "claude-sonnet", DisplayName: "anthropic/claude-sonnet"}, + } + al.fallback = providers.NewFallbackChain(providers.NewCooldownTracker(), nil) + + pipeline := NewPipeline(al) + ts := newTurnState(agent, makeTestProcessOpts("test-session"), turnEventScope{ + turnID: "turn-1", + context: newTurnContext(nil, nil, nil), + }) + + exec, err := pipeline.SetupTurn(context.Background(), ts) + if err != nil { + t.Fatalf("SetupTurn failed: %v", err) + } + + ctrl, err := pipeline.CallLLM(context.Background(), context.Background(), ts, exec, 1) + if err != nil { + t.Fatalf("CallLLM failed: %v", err) + } + if ctrl != ControlBreak { + t.Fatalf("expected ControlBreak, got %v", ctrl) + } + if exec.llmModelName != "anthropic/claude-sonnet" { + t.Fatalf("exec.llmModelName = %q, want %q", exec.llmModelName, "anthropic/claude-sonnet") + } +} + +func TestPipeline_SetupTurn_UsesLightCandidateDisplayName(t *testing.T) { + al, agent, cleanup := newTurnCoordTestLoop(t, &simpleConvProvider{}) + defer cleanup() + + agent.Model = "primary-model" + agent.Candidates = []providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4", IdentityKey: "model_name:primary", DisplayName: "primary-model"}, + } + agent.LightCandidates = []providers.FallbackCandidate{ + {Provider: "openai", Model: "gpt-5.4-mini", IdentityKey: "model_name:light-model", DisplayName: "light-model"}, + } + agent.Router = routing.New(routing.RouterConfig{LightModel: "light-model", Threshold: 1}) + + pipeline := NewPipeline(al) + opts := makeTestProcessOpts("test-session") + opts.UserMessage = "" + ts := newTurnState(agent, opts, turnEventScope{ + turnID: "turn-1", + context: newTurnContext(nil, nil, nil), + }) + + exec, err := pipeline.SetupTurn(context.Background(), ts) + if err != nil { + t.Fatalf("SetupTurn failed: %v", err) + } + if !exec.usedLight { + t.Fatal("expected light routing to be used") + } + if exec.llmModelName != "light-model" { + t.Fatalf("exec.llmModelName = %q, want %q", exec.llmModelName, "light-model") + } +} + func TestRunTurn_FinalizeSaveErrorEmitsErrorTurnEnd(t *testing.T) { al, agent, cleanup := newTurnCoordTestLoop(t, &simpleConvProvider{}) defer cleanup() diff --git a/pkg/agent/turn_profile_policy.go b/pkg/agent/turn_profile_policy.go new file mode 100644 index 00000000..140254e8 --- /dev/null +++ b/pkg/agent/turn_profile_policy.go @@ -0,0 +1,131 @@ +package agent + +import ( + "strings" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" +) + +func resolveTurnProfileOptions(cfg *config.Config, opts processOptions) (processOptions, error) { + if cfg == nil { + return opts, nil + } + profile, ok, err := cfg.Agents.Defaults.ResolveTurnProfile() + if err != nil { + return opts, err + } + if !ok { + return opts, nil + } + opts.TurnProfile = profile + if profile.HistoryMode == config.TurnProfileModeOff { + opts.NoHistory = true + opts.EnableSummary = false + } + return opts, nil +} + +func turnProfileSystemPromptOff(profile config.EffectiveTurnProfile) bool { + return profile.Enabled && profile.SystemPromptMode == config.TurnProfileModeOff +} + +func turnProfileSkillsOff(profile config.EffectiveTurnProfile) bool { + return profile.Enabled && profile.SkillsMode == config.TurnProfileModeOff +} + +func turnProfileCustomSkills(profile config.EffectiveTurnProfile) bool { + return profile.Enabled && profile.SkillsMode == config.TurnProfileModeCustom +} + +func turnProfileHasCallableTools( + profile config.EffectiveTurnProfile, + defs []providers.ToolDefinition, +) bool { + if !profile.Enabled { + return true + } + return len(filterToolsByTurnProfile(defs, profile)) > 0 +} + +func turnProfileToolAllowed(profile config.EffectiveTurnProfile, name string) bool { + if !profile.Enabled { + return true + } + switch profile.ToolsMode { + case config.TurnProfileModeOff: + return false + case config.TurnProfileModeCustom: + allowed := cleanAllowedSet(profile.AllowedTools) + if len(allowed) == 0 { + return false + } + _, ok := allowed[strings.ToLower(strings.TrimSpace(name))] + return ok + default: + return true + } +} + +func toolUseSystemPromptRule() string { + return "**ALWAYS use tools** - When you need to perform an action (schedule reminders, send messages, execute commands, etc.), you MUST call the appropriate tool. Do NOT just say you'll do it or pretend to do it." +} + +func filterNamesByTurnProfile(names []string, allowed []string) []string { + if len(names) == 0 { + return nil + } + allowedSet := cleanAllowedSet(allowed) + if len(allowedSet) == 0 { + return nil + } + out := make([]string, 0, len(names)) + for _, name := range names { + if _, ok := allowedSet[strings.ToLower(strings.TrimSpace(name))]; ok { + out = append(out, name) + } + } + return out +} + +func filterToolsByTurnProfile( + defs []providers.ToolDefinition, + profile config.EffectiveTurnProfile, +) []providers.ToolDefinition { + if !profile.Enabled { + return defs + } + switch profile.ToolsMode { + case config.TurnProfileModeOff: + return nil + case config.TurnProfileModeCustom: + allowed := cleanAllowedSet(profile.AllowedTools) + if len(allowed) == 0 { + return nil + } + filtered := make([]providers.ToolDefinition, 0, len(defs)) + for _, def := range defs { + if _, ok := allowed[strings.ToLower(strings.TrimSpace(def.Function.Name))]; ok { + filtered = append(filtered, def) + } + } + return filtered + default: + return defs + } +} + +func cleanAllowedSet(values []string) map[string]struct{} { + if len(values) == 0 { + return nil + } + out := make(map[string]struct{}, len(values)) + for _, value := range values { + value = strings.ToLower(strings.TrimSpace(value)) + if value == "" { + continue + } + out[value] = struct{}{} + } + return out +} diff --git a/pkg/agent/turn_profile_test.go b/pkg/agent/turn_profile_test.go new file mode 100644 index 00000000..eb79d523 --- /dev/null +++ b/pkg/agent/turn_profile_test.go @@ -0,0 +1,1177 @@ +package agent + +import ( + "context" + "os" + "path/filepath" + "reflect" + "strings" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/tools" +) + +type turnProfileCaptureProvider struct { + messages []providers.Message + tools []providers.ToolDefinition +} + +func (p *turnProfileCaptureProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.messages = append([]providers.Message(nil), messages...) + p.tools = append([]providers.ToolDefinition(nil), tools...) + return &providers.LLMResponse{Content: "profile response"}, nil +} + +func (p *turnProfileCaptureProvider) GetDefaultModel() string { + return "test-model" +} + +type turnProfileSideQuestionCaptureProvider struct { + messages []providers.Message +} + +func (p *turnProfileSideQuestionCaptureProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.messages = append([]providers.Message(nil), messages...) + return &providers.LLMResponse{Content: "side answer"}, nil +} + +func (p *turnProfileSideQuestionCaptureProvider) GetDefaultModel() string { + return "test-model" +} + +func newTurnProfileAgentLoop( + t *testing.T, + cfg *config.Config, + provider *turnProfileCaptureProvider, +) *AgentLoop { + t.Helper() + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + if cfg.Agents.Defaults.Workspace == "" { + cfg.Agents.Defaults.Workspace = t.TempDir() + } + if cfg.Agents.Defaults.ModelName == "" { + cfg.Agents.Defaults.ModelName = "test-model" + } + if cfg.Agents.Defaults.MaxTokens == 0 { + cfg.Agents.Defaults.MaxTokens = 4096 + } + if cfg.Agents.Defaults.MaxToolIterations == 0 { + cfg.Agents.Defaults.MaxToolIterations = 10 + } + return NewAgentLoop(cfg, bus.NewMessageBus(), provider) +} + +func writeTurnProfileSkill(t *testing.T, workspace, name, body string) { + t.Helper() + dir := filepath.Join(workspace, "skills", name) + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("mkdir skill: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte(body), 0o644); err != nil { + t.Fatalf("write skill: %v", err) + } +} + +func TestTurnProfile_DisabledPreservesDefaultHistoryAndPrompt(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: false, + History: config.TurnProfileBlock{ + Mode: config.TurnProfileModeOff, + }, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + agent := al.GetRegistry().GetDefaultAgent() + sessionKey := "agent:default:test-default" + ts := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + initialHistory := []providers.Message{ + {Role: "user", Content: "old user", CreatedAt: &ts}, + {Role: "assistant", Content: "old assistant", CreatedAt: &ts}, + } + agent.Sessions.SetHistory(sessionKey, initialHistory) + agent.Sessions.SetSummary(sessionKey, "old summary") + + got, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: sessionKey, + UserMessage: "new user", + DefaultResponse: defaultResponse, + EnableSummary: false, + }) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if got != "profile response" { + t.Fatalf("runAgentLoop() = %q, want profile response", got) + } + if len(provider.messages) != 4 { + t.Fatalf("provider messages len = %d, want system + history + user", len(provider.messages)) + } + if !reflect.DeepEqual(provider.messages[1:3], initialHistory) { + t.Fatalf("provider history = %#v, want %#v", provider.messages[1:3], initialHistory) + } + if !strings.Contains(provider.messages[0].Content, "CONTEXT_SUMMARY") { + t.Fatalf("system prompt missing summary in default mode:\n%s", provider.messages[0].Content) + } + history := agent.Sessions.GetHistory(sessionKey) + if len(history) != len(initialHistory)+2 { + t.Fatalf("history len = %d, want initial + user + assistant", len(history)) + } +} + +func TestTurnProfile_HistoryOffSuppressesHistoryAndPersistence(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + agent := al.GetRegistry().GetDefaultAgent() + sessionKey := "agent:default:test-history-off" + ts := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + initialHistory := []providers.Message{ + {Role: "user", Content: "old user", CreatedAt: &ts}, + {Role: "assistant", Content: "old assistant", CreatedAt: &ts}, + } + agent.Sessions.SetHistory(sessionKey, initialHistory) + agent.Sessions.SetSummary(sessionKey, "old summary") + + _, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: sessionKey, + UserMessage: "new user", + DefaultResponse: defaultResponse, + EnableSummary: true, + }) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if len(provider.messages) != 2 { + t.Fatalf("provider messages len = %d, want system + current user", len(provider.messages)) + } + if provider.messages[1].Content != "new user" { + t.Fatalf("current message = %q, want new user", provider.messages[1].Content) + } + if strings.Contains(provider.messages[0].Content, "old summary") { + t.Fatalf("system prompt includes suppressed summary:\n%s", provider.messages[0].Content) + } + history := agent.Sessions.GetHistory(sessionKey) + if !reflect.DeepEqual(history, initialHistory) { + t.Fatalf("history = %#v, want unchanged %#v", history, initialHistory) + } +} + +func TestTurnProfile_ProcessMessageUsesEnabledTurnProfile(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + + _, err := al.processMessage(context.Background(), bus.InboundMessage{ + Context: bus.InboundContext{ + Channel: "pico", + ChatID: "pico:sess-1", + ChatType: "direct", + SenderID: "pico-user", + }, + Content: "hello from pico", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if len(provider.messages) != 2 { + t.Fatalf("provider messages len = %d, want system + current user", len(provider.messages)) + } + if provider.messages[1].Content != "hello from pico" { + t.Fatalf("current message = %q, want hello from pico", provider.messages[1].Content) + } +} + +func TestTurnProfile_BtwCommandUsesEnabledTurnProfile(t *testing.T) { + workspace := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "openai/test-model", + }}, + } + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + sideProvider := &turnProfileSideQuestionCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), &turnProfileCaptureProvider{}) + al.providerFactory = func(mc *config.ModelConfig) (providers.LLMProvider, string, error) { + return sideProvider, "test-model", nil + } + + _, err := al.processMessage(context.Background(), bus.InboundMessage{ + Context: bus.InboundContext{ + Channel: "pico", + ChatID: "pico:btw", + ChatType: "direct", + SenderID: "pico-user", + }, + Content: "/btw explain privately", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if len(sideProvider.messages) < 2 { + t.Fatalf("side question messages len = %d, want system + user", len(sideProvider.messages)) + } + systemPrompt := sideProvider.messages[0].Content + blockedSnippets := []string{ + "ALWAYS use tools", + "When using tools", + "read_file tool", + "update " + workspace + "/memory/MEMORY.md", + } + for _, snippet := range blockedSnippets { + if strings.Contains(systemPrompt, snippet) { + t.Fatalf("side question system prompt includes %q despite tools.mode=off:\n%s", snippet, systemPrompt) + } + } +} + +func TestTurnProfile_BtwCommandDoesNotAddToolFallbackWhenSystemPromptOff(t *testing.T) { + workspace := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + SystemPrompt: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"echo_text"}, + }, + }, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "openai/test-model", + }}, + } + t.Setenv("PICOCLAW_BUILTIN_SKILLS", t.TempDir()) + sideProvider := &turnProfileSideQuestionCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), &turnProfileCaptureProvider{}) + al.RegisterTool(&echoTextTool{}) + al.providerFactory = func(mc *config.ModelConfig) (providers.LLMProvider, string, error) { + return sideProvider, "test-model", nil + } + + _, err := al.processMessage(context.Background(), bus.InboundMessage{ + Context: bus.InboundContext{ + Channel: "pico", + ChatID: "pico:btw-system-off", + ChatType: "direct", + SenderID: "pico-user", + }, + Content: "/btw explain privately", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + for _, msg := range sideProvider.messages { + if msg.Role == "system" && strings.Contains(msg.Content, toolUseSystemPromptRule()) { + t.Fatalf("side question system prompt includes tool fallback despite no tools:\n%s", msg.Content) + } + } +} + +func TestTurnProfile_BtwHookCannotReenableNativeSearchWhenToolsOff(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + ModelList: []*config.ModelConfig{{ + ModelName: "test-model", + Model: "openai/test-model", + }}, + } + provider := &nativeSearchCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + al.providerFactory = func(mc *config.ModelConfig) (providers.LLMProvider, string, error) { + return provider, "test-model", nil + } + if err := al.MountHook(NamedHook("enable-native-search", turnProfileEnableNativeSearchHook{})); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + _, err := al.processMessage(context.Background(), bus.InboundMessage{ + Context: bus.InboundContext{ + Channel: "pico", + ChatID: "pico:btw-native-search", + ChatType: "direct", + SenderID: "pico-user", + }, + Content: "/btw search privately", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if provider.lastOpts["turn_profile_test_hook"] != true { + t.Fatalf("BeforeLLM hook did not run for /btw: %#v", provider.lastOpts) + } + if provider.lastOpts["native_search"] == true { + t.Fatalf("native_search option enabled by /btw hook despite tools.mode=off: %#v", provider.lastOpts) + } +} + +func TestTurnProfile_SubTurnInheritsParentToolProfile(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"echo_text"}, + }, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + al.RegisterTool(&echoTextTool{}) + al.RegisterTool(&echoTextRewrittenTool{}) + agent := al.GetRegistry().GetDefaultAgent() + profile, ok, err := cfg.Agents.Defaults.ResolveTurnProfile() + if err != nil { + t.Fatalf("ResolveTurnProfile() error = %v", err) + } + if !ok { + t.Fatal("ResolveTurnProfile() did not return enabled profile") + } + parentOpts := processOptions{ + Dispatch: DispatchRequest{ + SessionKey: "agent:default:test-parent", + UserMessage: "parent", + }, + TurnProfile: profile, + } + parentTS := newTurnState(agent, parentOpts, turnEventScope{ + turnID: "parent-turn-profile", + }) + + _, err = spawnSubTurn(context.Background(), al, parentTS, SubTurnConfig{ + Model: "test-model", + SystemPrompt: "child task", + Timeout: time.Second, + }) + if err != nil { + t.Fatalf("spawnSubTurn() error = %v", err) + } + if len(provider.tools) != 1 { + t.Fatalf("child provider tools len = %d, want 1: %#v", len(provider.tools), provider.tools) + } + if provider.tools[0].Function.Name != "echo_text" { + t.Fatalf("child provider tool = %q, want echo_text", provider.tools[0].Function.Name) + } +} + +func TestTurnProfile_SystemPromptOffUsesExternalPromptOnly(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + SystemPrompt: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Skills: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + agent := al.GetRegistry().GetDefaultAgent() + + _, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "agent:default:test-external-prompt", + UserMessage: "hello", + DefaultResponse: defaultResponse, + SystemPromptOverride: "External prompt only.", + }) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if len(provider.messages) != 2 { + t.Fatalf("messages len = %d, want system + user", len(provider.messages)) + } + if strings.TrimSpace(provider.messages[0].Content) != "External prompt only." { + t.Fatalf("system prompt = %q, want external only", provider.messages[0].Content) + } +} + +func TestTurnProfile_SystemPromptOffBlankTurnStillSendsMessage(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + SystemPrompt: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Skills: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + + _, err := al.runAgentLoop(context.Background(), al.GetRegistry().GetDefaultAgent(), processOptions{ + SessionKey: "agent:default:test-blank-system-off", + UserMessage: "", + DefaultResponse: defaultResponse, + }) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if len(provider.messages) != 1 { + t.Fatalf( + "provider messages len = %d, want one blank user message: %#v", + len(provider.messages), + provider.messages, + ) + } + if provider.messages[0].Role != "user" || provider.messages[0].Content != "" { + t.Fatalf("provider message = %#v, want blank user message", provider.messages[0]) + } +} + +func TestTurnProfile_SystemPromptOffAddsToolFallbackWhenToolsVisible(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + SystemPrompt: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Skills: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"echo_text"}, + }, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + al.RegisterTool(&echoTextTool{}) + agent := al.GetRegistry().GetDefaultAgent() + + _, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "agent:default:test-tool-fallback", + UserMessage: "hello", + DefaultResponse: defaultResponse, + }) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if got := strings.TrimSpace(provider.messages[0].Content); got != toolUseSystemPromptRule() { + t.Fatalf("fallback prompt = %q, want existing tool rule %q", got, toolUseSystemPromptRule()) + } +} + +func TestTurnProfile_SkillsOffAndCustomControlCatalogAndActiveSkills(t *testing.T) { + workspace := t.TempDir() + writeTurnProfileSkill( + t, + workspace, + "shell", + "---\ndescription: shell skill\n---\n# shell\n\nUse shell carefully.", + ) + writeTurnProfileSkill( + t, + workspace, + "paint", + "---\ndescription: paint skill\n---\n# paint\n\nUse paint vividly.", + ) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{ + Mode: config.TurnProfileModeOff, + }, + Skills: config.TurnProfileBlock{ + Mode: config.TurnProfileModeOff, + }, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + agent := al.GetRegistry().GetDefaultAgent() + + _, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "agent:default:test-skills-off", + UserMessage: "hello", + DefaultResponse: defaultResponse, + ForcedSkills: []string{"shell"}, + }) + if err != nil { + t.Fatalf("runAgentLoop(no-skills) error = %v", err) + } + noSkillsPrompt := provider.messages[0].Content + if strings.Contains(noSkillsPrompt, "# Skills") || + strings.Contains(noSkillsPrompt, "# Active Skills") { + t.Fatalf("skills-off prompt includes skill context:\n%s", noSkillsPrompt) + } + + cfg.Agents.Defaults.TurnProfile.Skills = config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"shell"}, + } + provider = &turnProfileCaptureProvider{} + al = newTurnProfileAgentLoop(t, cfg, provider) + agent = al.GetRegistry().GetDefaultAgent() + + _, err = al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "agent:default:test-skills-custom", + UserMessage: "hello", + DefaultResponse: defaultResponse, + ForcedSkills: []string{"shell", "paint"}, + }) + if err != nil { + t.Fatalf("runAgentLoop(shell-only) error = %v", err) + } + customPrompt := provider.messages[0].Content + if !strings.Contains(customPrompt, "shell") || + !strings.Contains(customPrompt, "### Skill: shell") { + t.Fatalf("custom skills prompt missing allowed shell context:\n%s", customPrompt) + } + if strings.Contains(customPrompt, "paint") || + strings.Contains(customPrompt, "### Skill: paint") { + t.Fatalf("custom skills prompt includes disallowed paint context:\n%s", customPrompt) + } +} + +type turnProfileAddToolHook struct{} + +func (h turnProfileAddToolHook) BeforeLLM( + ctx context.Context, + req *LLMHookRequest, +) (*LLMHookRequest, HookDecision, error) { + next := req.Clone() + next.Tools = append(next.Tools, providers.ToolDefinition{ + Type: "function", + Function: providers.ToolFunctionDefinition{ + Name: "echo_text_rewritten", + Description: "hook-added tool", + Parameters: map[string]any{"type": "object"}, + }, + }) + return next, HookDecision{Action: HookActionModify}, nil +} + +type turnProfileEnableNativeSearchHook struct{} + +func (h turnProfileEnableNativeSearchHook) BeforeLLM( + ctx context.Context, + req *LLMHookRequest, +) (*LLMHookRequest, HookDecision, error) { + next := req.Clone() + if next.Options == nil { + next.Options = map[string]any{} + } + next.Options["turn_profile_test_hook"] = true + next.Options["native_search"] = true + return next, HookDecision{Action: HookActionModify}, nil +} + +func (h turnProfileEnableNativeSearchHook) AfterLLM( + ctx context.Context, + resp *LLMHookResponse, +) (*LLMHookResponse, HookDecision, error) { + return resp, HookDecision{Action: HookActionContinue}, nil +} + +type turnProfileRespondToolHook struct{} + +func (h turnProfileRespondToolHook) BeforeTool( + ctx context.Context, + req *ToolCallHookRequest, +) (*ToolCallHookRequest, HookDecision, error) { + next := req.Clone() + next.HookResult = &tools.ToolResult{ForLLM: "hook bypassed profile"} + return next, HookDecision{Action: HookActionRespond}, nil +} + +func (h turnProfileRespondToolHook) AfterTool( + ctx context.Context, + result *ToolResultHookResponse, +) (*ToolResultHookResponse, HookDecision, error) { + return result, HookDecision{Action: HookActionContinue}, nil +} + +type turnProfileToolCallProvider struct { + calls int + messages []providers.Message +} + +func (p *turnProfileToolCallProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.calls++ + p.messages = append([]providers.Message(nil), messages...) + if p.calls == 1 { + return &providers.LLMResponse{ + Content: "calling disallowed", + ToolCalls: []providers.ToolCall{{ + ID: "call_1", + Name: "echo_text_rewritten", + Arguments: map[string]any{"text": "blocked"}, + }}, + FinishReason: "tool_calls", + }, nil + } + return &providers.LLMResponse{Content: "done"}, nil +} + +func (p *turnProfileToolCallProvider) GetDefaultModel() string { + return "test-model" +} + +func TestTurnProfile_ToolsCustomFiltersProviderToolsAndHookAdditions(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"echo_text"}, + }, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + al.RegisterTool(&echoTextTool{}) + al.RegisterTool(&echoTextRewrittenTool{}) + if err := al.MountHook(NamedHook("add-disallowed-tool", turnProfileAddToolHook{})); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-tools-filter", + UserMessage: "hello", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if len(provider.tools) != 1 { + t.Fatalf("provider tools len = %d, want 1: %#v", len(provider.tools), provider.tools) + } + if provider.tools[0].Function.Name != "echo_text" { + t.Fatalf("provider tool = %q, want echo_text", provider.tools[0].Function.Name) + } +} + +func TestTurnProfile_ToolsOffDisablesProviderAndNativeSearchTools(t *testing.T) { + cfg := &config.Config{ + Tools: config.ToolsConfig{ + Web: config.WebToolsConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + PreferNative: true, + }, + }, + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + } + cfg.Agents.Defaults.Workspace = t.TempDir() + cfg.Agents.Defaults.ModelName = "test-model" + cfg.Agents.Defaults.MaxTokens = 4096 + cfg.Agents.Defaults.MaxToolIterations = 10 + provider := &nativeSearchCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-tools-off", + UserMessage: "hello", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if provider.lastOpts["native_search"] == true { + t.Fatalf("native_search option enabled despite tools.mode=off: %#v", provider.lastOpts) + } +} + +func TestTurnProfile_ToolsOffSuppressesToolUsePromptRule(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + al.RegisterTool(&echoTextTool{}) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-tools-off-prompt", + UserMessage: "hello", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if len(provider.tools) != 0 { + t.Fatalf("provider tools len = %d, want 0", len(provider.tools)) + } + if len(provider.messages) == 0 || provider.messages[0].Role != "system" { + t.Fatalf("first provider message = %#v, want system prompt", provider.messages) + } + if strings.Contains(provider.messages[0].Content, toolUseSystemPromptRule()) || + strings.Contains(provider.messages[0].Content, "**ALWAYS use tools**") { + t.Fatalf("tools-off system prompt still asks the model to use tools:\n%s", provider.messages[0].Content) + } +} + +func TestTurnProfile_ToolsCustomMissingToolSuppressesToolUsePromptRule(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"web_search"}, + }, + }, + }, + }, + } + provider := &turnProfileCaptureProvider{} + al := newTurnProfileAgentLoop(t, cfg, provider) + al.RegisterTool(&echoTextTool{}) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-tools-custom-missing-prompt", + UserMessage: "hello", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if len(provider.tools) != 0 { + t.Fatalf("provider tools len = %d, want 0", len(provider.tools)) + } + if strings.Contains(provider.messages[0].Content, toolUseSystemPromptRule()) || + strings.Contains(provider.messages[0].Content, "**ALWAYS use tools**") { + t.Fatalf( + "custom profile with no resolved tools still asks the model to use tools:\n%s", + provider.messages[0].Content, + ) + } +} + +func TestTurnProfile_ToolsCustomAllowsNativeWebSearch(t *testing.T) { + cfg := &config.Config{ + Tools: config.ToolsConfig{ + Web: config.WebToolsConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + PreferNative: true, + }, + }, + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"web_search"}, + }, + }, + }, + }, + } + provider := &nativeSearchCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-native-web-allowed", + UserMessage: "search", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if got, _ := provider.lastOpts["native_search"].(bool); !got { + t.Fatalf("native_search = %#v, want true", provider.lastOpts["native_search"]) + } +} + +func TestTurnProfile_SystemPromptOffAddsToolFallbackForNativeWebSearch(t *testing.T) { + cfg := &config.Config{ + Tools: config.ToolsConfig{ + Web: config.WebToolsConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + PreferNative: true, + }, + }, + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + SystemPrompt: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Skills: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"web_search"}, + }, + }, + }, + }, + } + provider := &nativeSearchCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-native-web-fallback", + UserMessage: "search", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if got, _ := provider.lastOpts["native_search"].(bool); !got { + t.Fatalf("native_search = %#v, want true", provider.lastOpts["native_search"]) + } + if len(provider.messages) == 0 || provider.messages[0].Content != toolUseSystemPromptRule() { + t.Fatalf("native-search-only prompt = %#v, want tool fallback", provider.messages) + } +} + +func TestTurnProfile_BeforeLLMHookCannotReenableNativeSearchWhenToolsOff(t *testing.T) { + cfg := &config.Config{ + Tools: config.ToolsConfig{ + Web: config.WebToolsConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + PreferNative: true, + }, + }, + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + }, + }, + }, + } + provider := &nativeSearchCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + if err := al.MountHook(NamedHook("enable-native-search", turnProfileEnableNativeSearchHook{})); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-native-web-hook-denied", + UserMessage: "search", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if provider.lastOpts["native_search"] == true { + t.Fatalf("native_search option enabled by hook despite tools.mode=off: %#v", provider.lastOpts) + } +} + +func TestTurnProfile_BeforeLLMHookCannotReenableNativeSearchWhenCustomToolsResolveEmpty( + t *testing.T, +) { + cfg := &config.Config{ + Tools: config.ToolsConfig{ + Web: config.WebToolsConfig{ + ToolConfig: config.ToolConfig{Enabled: true}, + PreferNative: true, + }, + }, + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: t.TempDir(), + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"missing_tool"}, + }, + }, + }, + }, + } + provider := &nativeSearchCaptureProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + if err := al.MountHook(NamedHook("enable-native-search", turnProfileEnableNativeSearchHook{})); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-native-web-hook-custom-empty", + UserMessage: "search", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if provider.lastOpts["native_search"] == true { + t.Fatalf( + "native_search option enabled by hook despite no resolved tools: %#v", + provider.lastOpts, + ) + } +} + +func TestTurnProfile_ToolExecutionRejectsDisallowedToolCalls(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"echo_text"}, + }, + }, + }, + }, + } + provider := &turnProfileToolCallProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + al.RegisterTool(&echoTextTool{}) + al.RegisterTool(&echoTextRewrittenTool{}) + + response, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-tool-exec-deny", + UserMessage: "run tool", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + if response != "done" { + t.Fatalf("response = %q, want done", response) + } + if provider.calls != 2 { + t.Fatalf("provider calls = %d, want 2", provider.calls) + } + var foundDeniedResult bool + for _, msg := range provider.messages { + if msg.Role == "tool" && + strings.Contains(msg.Content, "not allowed by the active turn profile") { + foundDeniedResult = true + break + } + } + if !foundDeniedResult { + t.Fatalf("second provider call did not include denied tool result: %#v", provider.messages) + } +} + +func TestTurnProfile_BeforeToolRespondCannotBypassDisallowedTool(t *testing.T) { + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + TurnProfile: config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"echo_text"}, + }, + }, + }, + }, + } + provider := &turnProfileToolCallProvider{} + al := NewAgentLoop(cfg, bus.NewMessageBus(), provider) + al.RegisterTool(&echoTextTool{}) + al.RegisterTool(&echoTextRewrittenTool{}) + if err := al.MountHook(NamedHook("respond-tool", turnProfileRespondToolHook{})); err != nil { + t.Fatalf("MountHook() error = %v", err) + } + + _, err := al.runAgentLoop( + context.Background(), + al.GetRegistry().GetDefaultAgent(), + processOptions{ + SessionKey: "agent:default:test-tool-hook-respond-denied", + UserMessage: "run tool", + DefaultResponse: defaultResponse, + }, + ) + if err != nil { + t.Fatalf("runAgentLoop() error = %v", err) + } + var foundDeniedResult bool + for _, msg := range provider.messages { + if msg.Role != "tool" { + continue + } + if strings.Contains(msg.Content, "hook bypassed profile") { + t.Fatalf("hook respond result bypassed turn profile: %#v", provider.messages) + } + if strings.Contains(msg.Content, "not allowed by the active turn profile") { + foundDeniedResult = true + } + } + if !foundDeniedResult { + t.Fatalf("second provider call did not include denied tool result: %#v", provider.messages) + } +} diff --git a/pkg/agent/turn_state.go b/pkg/agent/turn_state.go index 2df0c81c..3213ff55 100644 --- a/pkg/agent/turn_state.go +++ b/pkg/agent/turn_state.go @@ -11,6 +11,7 @@ import ( "time" "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/session" @@ -83,6 +84,7 @@ const ( type turnResult struct { finalContent string + modelName string status TurnEndStatus followUps []bus.InboundMessage } @@ -124,18 +126,23 @@ type turnExecution struct { iteration int // Per-iteration state set by Pipeline.PreLLM - activeCandidates []providers.FallbackCandidate - activeModel string - activeProvider providers.LLMProvider - usedLight bool + activeCandidates []providers.FallbackCandidate + activeModel string + activeModelConfig *config.ModelConfig + activeProvider providers.LLMProvider + usedLight bool // LLM call per-iteration state response *providers.LLMResponse normalizedToolCalls []providers.ToolCall allResponsesHandled bool + streamingPublisher *streamingChunkPublisher + streamingFallback bool + suppressReasoning bool callMessages []providers.Message providerToolDefs []providers.ToolDefinition llmModel string + llmModelName string llmOpts map[string]any gracefulTerminal bool useNativeSearch bool @@ -173,9 +180,10 @@ func newTurnExecution( type turnState struct { mu sync.RWMutex - agent *AgentInstance - opts processOptions - scope turnEventScope + agent *AgentInstance + opts processOptions + profile config.EffectiveTurnProfile + scope turnEventScope turnID string agentID string @@ -247,6 +255,7 @@ func newTurnState(agent *AgentInstance, opts processOptions, scope turnEventScop ts := &turnState{ agent: agent, opts: opts, + profile: opts.TurnProfile, scope: scope, turnID: scope.turnID, agentID: agent.ID, @@ -681,6 +690,21 @@ func (ts *turnState) restoreSession(agent *AgentInstance) error { return agent.Sessions.Save(ts.sessionKey) } +// messagesContentEqual compares two message slices by content only, ignoring CreatedAt. +// JSON roundtrip loses the monotonic clock portion of time.Time, so direct +// reflect.DeepEqual would always differ on messages that roundtripped through +// the JSONL store. +func messagesContentEqual(a, b []providers.Message) bool { + for i := range a { + aCopy, bCopy := a[i], b[i] + aCopy.CreatedAt, bCopy.CreatedAt = nil, nil + if !reflect.DeepEqual(aCopy, bCopy) { + return false + } + } + return true +} + func matchingTurnMessageTail(history, persisted []providers.Message) int { maxMatch := min(len(history), len(persisted)) for size := maxMatch; size > 0; size-- { diff --git a/pkg/audio/asr/asr.go b/pkg/audio/asr/asr.go index a7c93e57..94ede23e 100644 --- a/pkg/audio/asr/asr.go +++ b/pkg/audio/asr/asr.go @@ -29,12 +29,12 @@ func supportsAudioTranscription(modelCfg *config.ModelConfig) bool { protocol, _ := providers.ExtractProtocol(modelCfg) switch protocol { - case "openai", "azure", "azure-openai", + case "openai", "azure", "litellm", "openrouter", "groq", "zhipu", "gemini", "nvidia", "ollama", "moonshot", "shengsuanyun", "deepseek", "cerebras", - "vivgrid", "volcengine", "vllm", "qwen", "qwen-portal", "qwen-intl", "qwen-international", "dashscope-intl", - "qwen-us", "dashscope-us", "mistral", "avian", "minimax", "longcat", "modelscope", "novita", - "coding-plan", "alibaba-coding", "qwen-coding", "zai": + "vivgrid", "volcengine", "vllm", "qwen-portal", "qwen-intl", "qwen-us", + "mistral", "avian", "minimax", "longcat", "modelscope", "novita", + "alibaba-coding", "zai": // These protocols all go through the OpenAI-compatible or Azure provider path in // providers.CreateProviderFromConfig, so they are the only ones that can supply // the audio media payload shape expected by NewAudioModelTranscriber. @@ -53,9 +53,9 @@ func supportsWhisperTranscription(modelCfg *config.ModelConfig) bool { switch protocol { case "openai", "litellm", "openrouter", "groq", "zhipu", "gemini", "nvidia", "ollama", "moonshot", "shengsuanyun", "deepseek", "cerebras", - "vivgrid", "volcengine", "vllm", "qwen", "qwen-portal", "qwen-intl", "qwen-international", "dashscope-intl", - "qwen-us", "dashscope-us", "mistral", "avian", "minimax", "longcat", "modelscope", "novita", - "coding-plan", "alibaba-coding", "qwen-coding", "zai", "mimo": + "vivgrid", "volcengine", "vllm", "qwen-portal", "qwen-intl", "qwen-us", + "mistral", "avian", "minimax", "longcat", "modelscope", "novita", + "alibaba-coding", "zai", "mimo": return true default: return false diff --git a/pkg/bus/bus.go b/pkg/bus/bus.go index dee67d87..d076bd27 100644 --- a/pkg/bus/bus.go +++ b/pkg/bus/bus.go @@ -26,7 +26,7 @@ const defaultBusBufferSize = 64 type StreamDelegate interface { // GetStreamer returns a Streamer for the given channel+chatID if the channel // supports streaming. Returns nil, false if streaming is unavailable. - GetStreamer(ctx context.Context, channel, chatID string) (Streamer, bool) + GetStreamer(ctx context.Context, channel, chatID, sessionKey string) (Streamer, bool) } // Streamer pushes incremental content to a streaming-capable channel. @@ -37,6 +37,20 @@ type Streamer interface { Cancel(ctx context.Context) } +// ContextUsageStreamer can attach final context-window usage metadata when a +// streaming channel's final message replaces the normal outbound response. +type ContextUsageStreamer interface { + Streamer + FinalizeWithContext(ctx context.Context, content string, usage *ContextUsage) error +} + +// ReasoningStreamer can show incremental model reasoning/thought content +// separately from the final user-visible answer stream. +type ReasoningStreamer interface { + UpdateReasoning(ctx context.Context, content string) error + FinalizeReasoning(ctx context.Context, content string) error +} + type MessageBus struct { inbound chan InboundMessage outbound chan OutboundMessage @@ -182,10 +196,10 @@ func (mb *MessageBus) SetEventPublisher(p EventPublisher) { mb.eventPublisher.Store(p) } -// GetStreamer returns a Streamer for the given channel+chatID via the delegate. -func (mb *MessageBus) GetStreamer(ctx context.Context, channel, chatID string) (Streamer, bool) { +// GetStreamer returns a Streamer for the given channel+chatID+session via the delegate. +func (mb *MessageBus) GetStreamer(ctx context.Context, channel, chatID, sessionKey string) (Streamer, bool) { if d, ok := mb.streamDelegate.Load().(StreamDelegate); ok && d != nil { - return d.GetStreamer(ctx, channel, chatID) + return d.GetStreamer(ctx, channel, chatID, sessionKey) } return nil, false } diff --git a/pkg/channels/interfaces.go b/pkg/channels/interfaces.go index 0cfd435b..ad442195 100644 --- a/pkg/channels/interfaces.go +++ b/pkg/channels/interfaces.go @@ -20,6 +20,17 @@ type MessageEditor interface { EditMessage(ctx context.Context, chatID string, messageID string, content string) error } +// MessageEditorWithPayload extends MessageEditor for channels that can update +// structured message metadata in addition to plain text content. +type MessageEditorWithPayload interface { + EditMessageWithPayload( + ctx context.Context, + chatID string, + messageID string, + payload map[string]any, + ) error +} + // MessageDeleter — channels that can delete a message by ID. type MessageDeleter interface { DeleteMessage(ctx context.Context, chatID string, messageID string) error diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index d345a5d0..03f65e11 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -41,6 +41,8 @@ const ( janitorInterval = 10 * time.Second typingStopTTL = 5 * time.Minute placeholderTTL = 10 * time.Minute + + streamAuxiliaryTombstoneTTL = 30 * time.Second ) // typingEntry wraps a typing stop function with a creation timestamp for TTL eviction. @@ -82,22 +84,23 @@ type channelWorker struct { } type Manager struct { - channels map[string]Channel - workers map[string]*channelWorker - bus *bus.MessageBus - runtimeEvents runtimeevents.Bus - config *config.Config - mediaStore media.MediaStore - dispatchTask *asyncTask - mux *dynamicServeMux - httpServer *http.Server - httpListeners []net.Listener - mu sync.RWMutex - placeholders sync.Map // "channel:chatID" → placeholderID (string) - typingStops sync.Map // "channel:chatID" → func() - reactionUndos sync.Map // "channel:chatID" → reactionEntry - streamActive sync.Map // "channel:chatID" → true (set when streamer.Finalize sent the message) - channelHashes map[string]string // channel name → config hash + channels map[string]Channel + workers map[string]*channelWorker + bus *bus.MessageBus + runtimeEvents runtimeevents.Bus + config *config.Config + mediaStore media.MediaStore + dispatchTask *asyncTask + mux *dynamicServeMux + httpServer *http.Server + httpListeners []net.Listener + mu sync.RWMutex + placeholders sync.Map // "channel:chatID" → placeholderID (string) + typingStops sync.Map // "channel:chatID" → func() + reactionUndos sync.Map // "channel:chatID" → reactionEntry + streamActive sync.Map // streamSuppressionKey → true (set when streamer.Finalize sent the message) + streamAuxiliaryTombstones sync.Map // streamSuppressionKey → time.Time (drops late auxiliary messages after stream final) + channelHashes map[string]string // channel name → config hash } type mediaStoreSetter interface { @@ -166,6 +169,20 @@ func outboundMessageIsToolFeedback(msg bus.OutboundMessage) bool { return strings.EqualFold(strings.TrimSpace(msg.Context.Raw["message_kind"]), "tool_feedback") } +func outboundMessageHasAuxiliaryKind(msg bus.OutboundMessage) bool { + if len(msg.Context.Raw) == 0 { + return false + } + return strings.TrimSpace(msg.Context.Raw["message_kind"]) != "" +} + +func outboundMessageIsFinal(msg bus.OutboundMessage) bool { + if len(msg.Context.Raw) == 0 { + return false + } + return strings.EqualFold(strings.TrimSpace(msg.Context.Raw["outbound_kind"]), "final") +} + func outboundMessageBypassesPlaceholderEdit(msg bus.OutboundMessage) bool { if len(msg.Context.Raw) == 0 { return false @@ -174,6 +191,19 @@ func outboundMessageBypassesPlaceholderEdit(msg bus.OutboundMessage) bool { return strings.EqualFold(kind, "thought") || strings.EqualFold(kind, "tool_calls") } +func outboundMessageEditPayload(msg bus.OutboundMessage, content string) map[string]any { + payload := map[string]any{ + "content": content, + } + if len(msg.Context.Raw) == 0 { + return payload + } + if modelName := strings.TrimSpace(msg.Context.Raw["model_name"]); modelName != "" { + payload["model_name"] = modelName + } + return payload +} + func outboundMediaChannel(msg bus.OutboundMediaMessage) string { return msg.Context.Channel } @@ -182,6 +212,14 @@ func outboundMediaChatID(msg bus.OutboundMediaMessage) string { return msg.ChatID } +func streamSuppressionKey(channel, chatID, sessionKey string) string { + key := channel + ":" + chatID + if strings.TrimSpace(sessionKey) == "" { + return key + } + return key + ":" + sessionKey +} + func trackedToolFeedbackMessageChatID(ch Channel, chatID string, outboundCtx *bus.InboundContext) string { if resolver, ok := ch.(toolFeedbackMessageTargetResolver); ok { if resolved := strings.TrimSpace(resolver.ToolFeedbackMessageChatID(chatID, outboundCtx)); resolved != "" { @@ -324,6 +362,7 @@ func (m *Manager) RecordReactionUndo(channel, chatID string, undo func()) { func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMessage, ch Channel) ([]string, bool) { chatID := outboundMessageChatID(msg) key := name + ":" + chatID + streamKey := streamSuppressionKey(name, chatID, msg.SessionKey) // 1. Stop typing if v, loaded := m.typingStops.LoadAndDelete(key); loaded { @@ -340,38 +379,67 @@ func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMess } isToolFeedback := outboundMessageIsToolFeedback(msg) + isAuxiliaryMessage := outboundMessageHasAuxiliaryKind(msg) + isFinalMessage := outboundMessageIsFinal(msg) separateToolFeedbackMessages := m.toolFeedbackSeparateMessagesEnabled() - // 3. If a stream already finalized this chat, stale tool feedback must be - // dropped without consuming the final-response marker. Streaming finalization - // bypasses the worker queue, so older queued feedback can arrive before the - // normal final outbound message that cleans up the marker and placeholder. - if isToolFeedback { - if _, loaded := m.streamActive.Load(key); loaded { + // 3. If a stream already finalized this chat, stale auxiliary messages must + // be dropped without consuming the final-response marker. Streaming + // finalization bypasses the worker queue, so older queued feedback/thoughts + // can arrive before the normal final outbound message that cleans up the + // marker and placeholder. + if isAuxiliaryMessage { + if _, loaded := m.streamActive.Load(streamKey); loaded { + return nil, true + } + if m.streamAuxiliaryTombstoneActive(streamKey) { return nil, true } } - // 4. If a stream already finalized this message, delete the placeholder and skip send - if _, loaded := m.streamActive.LoadAndDelete(key); loaded { - if v, loaded := m.placeholders.LoadAndDelete(key); loaded { - if entry, ok := v.(placeholderEntry); ok && entry.id != "" { - // Prefer deleting the placeholder (cleaner UX than editing to same content) - if deleter, ok := ch.(MessageDeleter); ok { - deleter.DeleteMessage(ctx, chatID, entry.id) // best effort - } else if editor, ok := ch.(MessageEditor); ok { - editor.EditMessage(ctx, chatID, entry.id, msg.Content) // fallback + // 4. If a stream already finalized this turn, skip only the duplicate final + // outbound. Earlier queued visible messages must still be delivered. + if isFinalMessage { + if _, loaded := m.streamActive.LoadAndDelete(streamKey); loaded { + if v, loaded := m.placeholders.LoadAndDelete(key); loaded { + if entry, ok := v.(placeholderEntry); ok && entry.id != "" { + // Prefer deleting the placeholder (cleaner UX than editing to same content) + if deleter, ok := ch.(MessageDeleter); ok { + deleter.DeleteMessage(ctx, chatID, entry.id) // best effort + } else if editor, ok := ch.(MessageEditor); ok { + if payloadEditor, ok := ch.(MessageEditorWithPayload); ok { + _ = payloadEditor.EditMessageWithPayload( + ctx, + chatID, + entry.id, + outboundMessageEditPayload(msg, msg.Content), + ) + } else { + editor.EditMessage(ctx, chatID, entry.id, msg.Content) // fallback + } + } } } - } - if !isToolFeedback { - if separateToolFeedbackMessages { - clearTrackedToolFeedbackMessage(ch, chatID, &msg.Context) - } else { - dismissTrackedToolFeedbackMessage(ctx, ch, chatID, &msg.Context) + if !isToolFeedback { + if separateToolFeedbackMessages { + clearTrackedToolFeedbackMessage(ch, chatID, &msg.Context) + } else { + dismissTrackedToolFeedbackMessage(ctx, ch, chatID, &msg.Context) + } } + return nil, true } - return nil, true + } + + if _, loaded := m.streamActive.Load(streamKey); loaded { + return nil, false + } + if m.streamActiveForChat(name, chatID) { + return nil, false + } + + if !isAuxiliaryMessage { + m.streamAuxiliaryTombstones.Delete(streamKey) } if separateToolFeedbackMessages { @@ -400,7 +468,18 @@ func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMess trackedContent = prepareToolFeedbackMessageContent(ch, msg.Content) content = InitialAnimatedToolFeedbackContent(trackedContent) } - if err := editor.EditMessage(ctx, chatID, entry.id, content); err == nil { + err := func() error { + if payloadEditor, ok := ch.(MessageEditorWithPayload); ok { + return payloadEditor.EditMessageWithPayload( + ctx, + chatID, + entry.id, + outboundMessageEditPayload(msg, content), + ) + } + return editor.EditMessage(ctx, chatID, entry.id, content) + }() + if err == nil { trackedChatID := trackedToolFeedbackMessageChatID(ch, chatID, &msg.Context) if tracker, ok := ch.(toolFeedbackMessageTracker); ok && isToolFeedback { tracker.RecordToolFeedbackMessage(trackedChatID, entry.id, trackedContent) @@ -424,6 +503,7 @@ func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMess func (m *Manager) preSendMedia(ctx context.Context, name string, msg bus.OutboundMediaMessage, ch Channel) { chatID := outboundMediaChatID(msg) key := name + ":" + chatID + streamKey := streamSuppressionKey(name, chatID, msg.SessionKey) // 1. Stop typing if v, loaded := m.typingStops.LoadAndDelete(key); loaded { @@ -439,8 +519,9 @@ func (m *Manager) preSendMedia(ctx context.Context, name string, msg bus.Outboun } } - // 3. Clear any finalized stream marker for this chat before media delivery. - m.streamActive.LoadAndDelete(key) + // 3. Clear any finalized stream markers for this chat before media delivery. + m.streamActive.LoadAndDelete(streamKey) + m.streamAuxiliaryTombstones.Delete(streamKey) if m.toolFeedbackSeparateMessagesEnabled() { clearTrackedToolFeedbackMessage(ch, chatID, &msg.Context) @@ -507,7 +588,7 @@ func (m *Manager) SetMediaStore(store media.MediaStore) { // GetStreamer implements bus.StreamDelegate. // It checks if the named channel supports streaming and returns a Streamer. -func (m *Manager) GetStreamer(ctx context.Context, channelName, chatID string) (bus.Streamer, bool) { +func (m *Manager) GetStreamer(ctx context.Context, channelName, chatID, sessionKey string) (bus.Streamer, bool) { m.mu.RLock() ch, exists := m.channels[channelName] m.mu.RUnlock() @@ -531,51 +612,323 @@ func (m *Manager) GetStreamer(ctx context.Context, channelName, chatID string) ( } // Mark streamActive on Finalize so preSend knows to clean up the placeholder - key := channelName + ":" + chatID - return &finalizeHookStreamer{ - Streamer: streamer, - onFinalize: func(finalizeCtx context.Context) { - if m.toolFeedbackSeparateMessagesEnabled() { - clearTrackedToolFeedbackMessage( - ch, - chatID, - &bus.InboundContext{ - Channel: channelName, - ChatID: chatID, - }, - ) - } else { - dismissTrackedToolFeedbackMessage( - finalizeCtx, - ch, - chatID, - &bus.InboundContext{ - Channel: channelName, - ChatID: chatID, - }, - ) + // and late auxiliary messages cannot leak after streaming produced a final. + streamKey := streamSuppressionKey(channelName, chatID, sessionKey) + placeholderKey := channelName + ":" + chatID + clearMarker := func() { + m.streamActive.Delete(streamKey) + } + onFinalize := func(finalizeCtx context.Context, finalContent string) { + if m.toolFeedbackSeparateMessagesEnabled() { + clearTrackedToolFeedbackMessage( + ch, + chatID, + &bus.InboundContext{ + Channel: channelName, + ChatID: chatID, + }, + ) + } else { + dismissTrackedToolFeedbackMessage( + finalizeCtx, + ch, + chatID, + &bus.InboundContext{ + Channel: channelName, + ChatID: chatID, + }, + ) + } + if v, loaded := m.placeholders.LoadAndDelete(placeholderKey); loaded { + if entry, ok := v.(placeholderEntry); ok && entry.id != "" { + if deleter, ok := ch.(MessageDeleter); ok { + deleter.DeleteMessage(finalizeCtx, chatID, entry.id) // best effort + } else if editor, ok := ch.(MessageEditor); ok { + editor.EditMessage(finalizeCtx, chatID, entry.id, finalContent) // best effort fallback + } } - m.streamActive.Store(key, true) - }, + } + m.streamActive.Store(streamKey, true) + m.streamAuxiliaryTombstones.Store(streamKey, time.Now()) + } + + if m.config != nil && m.config.Agents.Defaults.SplitOnMarker { + return &splitMarkerStreamer{ + current: streamer, + reasoning: reasoningStreamerFrom(streamer), + begin: func(beginCtx context.Context) (bus.Streamer, error) { return sc.BeginStream(beginCtx, chatID) }, + onFinalize: onFinalize, + clearMarker: clearMarker, + }, true + } + + return &finalizeHookStreamer{ + Streamer: streamer, + clearMarker: clearMarker, + onFinalize: onFinalize, }, true } +func reasoningStreamerFrom(streamer bus.Streamer) bus.ReasoningStreamer { + if reasoningStreamer, ok := streamer.(bus.ReasoningStreamer); ok { + return reasoningStreamer + } + return nil +} + +type modelNameStreamer interface { + SetModelName(modelName string) +} + +func setStreamerModelName(streamer any, modelName string) { + setter, ok := streamer.(modelNameStreamer) + if !ok { + return + } + setter.SetModelName(modelName) +} + +// splitMarkerStreamer turns accumulated streaming text containing +// MessageSplitMarker into separate channel stream messages. +type splitMarkerStreamer struct { + mu sync.Mutex + current bus.Streamer + reasoning bus.ReasoningStreamer + begin func(context.Context) (bus.Streamer, error) + completedParts int + finalized bool + onFinalize func(context.Context, string) + clearMarker func() + modelName string +} + +func (s *splitMarkerStreamer) Update(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + return s.updateLocked(ctx, content) +} + +func (s *splitMarkerStreamer) Finalize(ctx context.Context, content string) error { + return s.FinalizeWithContext(ctx, content, nil) +} + +func (s *splitMarkerStreamer) FinalizeWithContext(ctx context.Context, content string, usage *bus.ContextUsage) error { + s.mu.Lock() + defer s.mu.Unlock() + if err := s.finalizeLocked(ctx, content, usage); err != nil { + return err + } + s.runFinalizeHook(ctx, content) + return nil +} + +func (s *splitMarkerStreamer) UpdateReasoning(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.reasoning == nil { + return nil + } + setStreamerModelName(s.reasoning, s.modelName) + return s.reasoning.UpdateReasoning(ctx, content) +} + +func (s *splitMarkerStreamer) FinalizeReasoning(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.reasoning == nil { + return nil + } + setStreamerModelName(s.reasoning, s.modelName) + return s.reasoning.FinalizeReasoning(ctx, content) +} + +func (s *splitMarkerStreamer) SetModelName(modelName string) { + s.mu.Lock() + defer s.mu.Unlock() + s.modelName = strings.TrimSpace(modelName) + setStreamerModelName(s.current, s.modelName) + setStreamerModelName(s.reasoning, s.modelName) +} + +func (s *splitMarkerStreamer) Cancel(ctx context.Context) { + s.mu.Lock() + defer s.mu.Unlock() + if s.current != nil { + s.current.Cancel(ctx) + } +} + +func (s *splitMarkerStreamer) ClearFinalizedStreamMarker() { + if s.clearMarker != nil { + s.clearMarker() + } +} + +func (s *splitMarkerStreamer) updateLocked(ctx context.Context, content string) error { + parts := strings.Split(content, MessageSplitMarker) + completedLimit := len(parts) - 1 + if err := s.finalizeCompletedPartsLocked(ctx, parts, completedLimit, nil); err != nil { + return err + } + active := strings.TrimSpace(parts[len(parts)-1]) + if active == "" { + return nil + } + if err := s.ensureCurrentLocked(ctx); err != nil { + return err + } + return s.current.Update(ctx, active) +} + +func (s *splitMarkerStreamer) finalizeLocked(ctx context.Context, content string, usage *bus.ContextUsage) error { + parts := strings.Split(content, MessageSplitMarker) + return s.finalizeCompletedPartsLocked(ctx, parts, len(parts), usage) +} + +func (s *splitMarkerStreamer) finalizeCompletedPartsLocked( + ctx context.Context, + parts []string, + limit int, + usage *bus.ContextUsage, +) error { + for s.completedParts < limit { + content := strings.TrimSpace(parts[s.completedParts]) + isLast := s.completedParts == limit-1 + if content != "" { + if err := s.ensureCurrentLocked(ctx); err != nil { + return err + } + if isLast && usage != nil { + if contextStreamer, ok := s.current.(bus.ContextUsageStreamer); ok { + if err := contextStreamer.FinalizeWithContext(ctx, content, usage); err != nil { + return err + } + } else if err := s.current.Finalize(ctx, content); err != nil { + return err + } + } else if err := s.current.Finalize(ctx, content); err != nil { + return err + } + s.current = nil + } + s.completedParts++ + } + return nil +} + +func (s *splitMarkerStreamer) ensureCurrentLocked(ctx context.Context) error { + if s.current != nil { + return nil + } + if s.begin == nil { + return fmt.Errorf("streamer is not initialized") + } + streamer, err := s.begin(ctx) + if err != nil { + return err + } + s.current = streamer + setStreamerModelName(s.current, s.modelName) + return nil +} + +func (s *splitMarkerStreamer) runFinalizeHook(ctx context.Context, content string) { + if s.finalized { + return + } + s.finalized = true + if s.onFinalize != nil { + s.onFinalize(ctx, content) + } +} + +func (m *Manager) streamAuxiliaryTombstoneActive(key string) bool { + v, ok := m.streamAuxiliaryTombstones.Load(key) + if !ok { + return false + } + createdAt, ok := v.(time.Time) + if !ok || time.Since(createdAt) > streamAuxiliaryTombstoneTTL { + m.streamAuxiliaryTombstones.Delete(key) + return false + } + return true +} + +func (m *Manager) streamActiveForChat(channel, chatID string) bool { + chatKey := streamSuppressionKey(channel, chatID, "") + found := false + m.streamActive.Range(func(key, _ any) bool { + keyString, ok := key.(string) + if !ok { + return true + } + if keyString == chatKey || strings.HasPrefix(keyString, chatKey+":") { + found = true + return false + } + return true + }) + return found +} + // finalizeHookStreamer wraps a Streamer to run a hook on Finalize. type finalizeHookStreamer struct { Streamer - onFinalize func(context.Context) + onFinalize func(context.Context, string) + clearMarker func() } func (s *finalizeHookStreamer) Finalize(ctx context.Context, content string) error { if err := s.Streamer.Finalize(ctx, content); err != nil { return err } - if s.onFinalize != nil { - s.onFinalize(ctx) + s.runFinalizeHook(ctx, content) + return nil +} + +func (s *finalizeHookStreamer) FinalizeWithContext(ctx context.Context, content string, usage *bus.ContextUsage) error { + if streamer, ok := s.Streamer.(bus.ContextUsageStreamer); ok { + if err := streamer.FinalizeWithContext(ctx, content, usage); err != nil { + return err + } + } else if err := s.Streamer.Finalize(ctx, content); err != nil { + return err + } + s.runFinalizeHook(ctx, content) + return nil +} + +func (s *finalizeHookStreamer) UpdateReasoning(ctx context.Context, content string) error { + if streamer, ok := s.Streamer.(bus.ReasoningStreamer); ok { + return streamer.UpdateReasoning(ctx, content) } return nil } +func (s *finalizeHookStreamer) FinalizeReasoning(ctx context.Context, content string) error { + if streamer, ok := s.Streamer.(bus.ReasoningStreamer); ok { + return streamer.FinalizeReasoning(ctx, content) + } + return nil +} + +func (s *finalizeHookStreamer) SetModelName(modelName string) { + setStreamerModelName(s.Streamer, strings.TrimSpace(modelName)) +} + +func (s *finalizeHookStreamer) runFinalizeHook(ctx context.Context, content string) { + if s.onFinalize != nil { + s.onFinalize(ctx, content) + } +} + +func (s *finalizeHookStreamer) ClearFinalizedStreamMarker() { + if s.clearMarker != nil { + s.clearMarker() + } +} + // initChannel is a helper that looks up a factory by type name and creates the channel. // typeName is the channel type used for factory lookup (e.g., "telegram"). // channelName is the config map key used as the channel's runtime name (e.g., "my_telegram"). @@ -1063,7 +1416,11 @@ func (m *Manager) runWorker(ctx context.Context, name string, w *channelWorker) // Step 1: Try marker-based splitting if enabled. // Tool feedback must stay a single message, so it skips marker splitting. - if m.config != nil && m.config.Agents.Defaults.SplitOnMarker && !outboundMessageIsToolFeedback(msg) { + // Stream-final duplicate responses must also stay intact so preSend can + // consume the whole final message before any marker chunk leaks. + if m.finalizedStreamActiveForMessage(name, msg) { + chunks = []string{msg.Content} + } else if m.config != nil && m.config.Agents.Defaults.SplitOnMarker && !outboundMessageIsToolFeedback(msg) { if markerChunks := SplitByMarker(msg.Content); len(markerChunks) > 1 { for _, chunk := range markerChunks { chunkMsg := msg @@ -1090,6 +1447,18 @@ func (m *Manager) runWorker(ctx context.Context, name string, w *channelWorker) } } +func (m *Manager) finalizedStreamActiveForMessage(channelName string, msg bus.OutboundMessage) bool { + if m == nil || !outboundMessageIsFinal(msg) { + return false + } + chatID := outboundMessageChatID(msg) + if strings.TrimSpace(channelName) == "" || strings.TrimSpace(chatID) == "" { + return false + } + _, active := m.streamActive.Load(streamSuppressionKey(channelName, chatID, msg.SessionKey)) + return active +} + // splitOutboundMessageContent splits regular outbound content by maxLen, but // keeps tool feedback in a single message by truncating the explanation body. func splitOutboundMessageContent(msg bus.OutboundMessage, maxLen int) []string { @@ -1389,9 +1758,9 @@ func (m *Manager) sendMediaWithRetry( return nil, lastErr } -// runTTLJanitor periodically scans the typingStops and placeholders maps -// and evicts entries that have exceeded their TTL. This prevents memory -// accumulation when outbound paths fail to trigger preSend (e.g. LLM errors). +// runTTLJanitor periodically scans the typingStops, placeholders, and stream +// tombstone maps and evicts entries that have exceeded their TTL. This prevents +// memory accumulation when outbound paths fail to trigger preSend (e.g. LLM errors). func (m *Manager) runTTLJanitor(ctx context.Context) { ticker := time.NewTicker(janitorInterval) defer ticker.Stop() @@ -1429,6 +1798,12 @@ func (m *Manager) runTTLJanitor(ctx context.Context) { } return true }) + m.streamAuxiliaryTombstones.Range(func(key, value any) bool { + if createdAt, ok := value.(time.Time); !ok || now.Sub(createdAt) > streamAuxiliaryTombstoneTTL { + m.streamAuxiliaryTombstones.Delete(key) + } + return true + }) } } } diff --git a/pkg/channels/manager_test.go b/pkg/channels/manager_test.go index 8c2f6ecf..a7ceac67 100644 --- a/pkg/channels/manager_test.go +++ b/pkg/channels/manager_test.go @@ -104,7 +104,8 @@ func (m *mockDeletingMediaChannel) DismissToolFeedbackMessage(_ context.Context, } type mockStreamer struct { - finalizeFn func(context.Context, string) error + finalizeFn func(context.Context, string) error + finalizeWithContextFn func(context.Context, string, *bus.ContextUsage) error } func (m *mockStreamer) Update(context.Context, string) error { return nil } @@ -116,15 +117,82 @@ func (m *mockStreamer) Finalize(ctx context.Context, content string) error { return nil } +func (m *mockStreamer) FinalizeWithContext(ctx context.Context, content string, usage *bus.ContextUsage) error { + if m.finalizeWithContextFn != nil { + return m.finalizeWithContextFn(ctx, content, usage) + } + return m.Finalize(ctx, content) +} + func (m *mockStreamer) Cancel(context.Context) {} +type mockReasoningStreamer struct { + mockStreamer + reasoningUpdates []string + reasoningFinal string +} + +func (m *mockReasoningStreamer) UpdateReasoning(_ context.Context, content string) error { + m.reasoningUpdates = append(m.reasoningUpdates, content) + return nil +} + +func (m *mockReasoningStreamer) FinalizeReasoning(_ context.Context, content string) error { + m.reasoningFinal = content + return nil +} + +type modelTrackingReasoningStreamer struct { + mockReasoningStreamer + modelNames []string +} + +func (m *modelTrackingReasoningStreamer) SetModelName(modelName string) { + m.modelNames = append(m.modelNames, strings.TrimSpace(modelName)) +} + +type recordingStreamSegment struct { + updates []string + finals []string + finalUsage *bus.ContextUsage + canceledCount int + modelNames []string +} + +func (s *recordingStreamSegment) Update(_ context.Context, content string) error { + s.updates = append(s.updates, content) + return nil +} + +func (s *recordingStreamSegment) Finalize(ctx context.Context, content string) error { + return s.FinalizeWithContext(ctx, content, nil) +} + +func (s *recordingStreamSegment) FinalizeWithContext(_ context.Context, content string, usage *bus.ContextUsage) error { + s.finals = append(s.finals, content) + s.finalUsage = usage + return nil +} + +func (s *recordingStreamSegment) Cancel(context.Context) { + s.canceledCount++ +} + +func (s *recordingStreamSegment) SetModelName(modelName string) { + s.modelNames = append(s.modelNames, strings.TrimSpace(modelName)) +} + type mockStreamingChannel struct { mockMessageEditor streamer Streamer + beginStreamFn func(context.Context, string) (Streamer, error) resolveChatIDFn func(chatID string, outboundCtx *bus.InboundContext) string } -func (m *mockStreamingChannel) BeginStream(context.Context, string) (Streamer, error) { +func (m *mockStreamingChannel) BeginStream(ctx context.Context, chatID string) (Streamer, error) { + if m.beginStreamFn != nil { + return m.beginStreamFn(ctx, chatID) + } if m.streamer == nil { return nil, errors.New("missing streamer") } @@ -1476,6 +1544,9 @@ func TestPreSend_StaleToolFeedbackDoesNotConsumeStreamActiveMarker(t *testing.T) Context: bus.InboundContext{ Channel: "test", ChatID: "123", + Raw: map[string]string{ + "outbound_kind": "final", + }, }, }) @@ -1494,6 +1565,206 @@ func TestPreSend_StaleToolFeedbackDoesNotConsumeStreamActiveMarker(t *testing.T) } } +func TestPreSend_StaleThoughtDoesNotConsumeStreamActiveMarker(t *testing.T) { + m := newTestManager() + m.streamActive.Store("test:123", true) + m.streamAuxiliaryTombstones.Store("test:123", time.Now()) + m.RecordPlaceholder("test", "123", "placeholder-1") + + var editedContent string + ch := &mockMessageEditor{ + editFn: func(_ context.Context, chatID, messageID, content string) error { + if chatID != "123" || messageID != "placeholder-1" { + t.Fatalf("unexpected edit target: %s/%s", chatID, messageID) + } + editedContent = content + return nil + }, + } + + thought := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "late reasoning", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + Raw: map[string]string{ + "message_kind": "thought", + }, + }, + }) + + msgIDs, handled := m.preSend(context.Background(), "test", thought, ch) + if !handled { + t.Fatal("expected stale thought to be dropped after stream finalize") + } + if len(msgIDs) != 0 { + t.Fatalf("expected no delivered message IDs for stale thought, got %v", msgIDs) + } + if _, ok := m.streamActive.Load("test:123"); !ok { + t.Fatal("expected streamActive marker to remain for the final outbound message") + } + if _, ok := m.placeholders.Load("test:123"); !ok { + t.Fatal("expected placeholder cleanup to remain deferred to the final outbound message") + } + if ch.editedMessages != 0 { + t.Fatalf("expected no placeholder edit for stale thought, got %d edits", ch.editedMessages) + } + + finalMsg := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "final streamed reply", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + Raw: map[string]string{ + "outbound_kind": "final", + }, + }, + }) + + _, handled = m.preSend(context.Background(), "test", finalMsg, ch) + if !handled { + t.Fatal("expected final outbound message to consume streamActive marker") + } + if _, ok := m.streamActive.Load("test:123"); ok { + t.Fatal("expected streamActive marker to be cleared by final outbound message") + } + if _, ok := m.placeholders.Load("test:123"); ok { + t.Fatal("expected placeholder to be cleaned up by final outbound message") + } + if editedContent != "final streamed reply" { + t.Fatalf("editedContent = %q, want final streamed reply", editedContent) + } + + lateThought := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "later reasoning", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + Raw: map[string]string{ + "message_kind": "thought", + }, + }, + }) + msgIDs, handled = m.preSend(context.Background(), "test", lateThought, ch) + if !handled { + t.Fatal("expected tombstone to drop late thought after final outbound was suppressed") + } + if len(msgIDs) != 0 { + t.Fatalf("expected no delivered message IDs for late thought, got %v", msgIDs) + } +} + +func TestPreSend_StreamActiveDoesNotConsumeEarlierVisibleMessage(t *testing.T) { + m := newTestManager() + m.streamActive.Store("test:123", true) + m.streamAuxiliaryTombstones.Store("test:123", time.Now()) + m.RecordPlaceholder("test", "123", "placeholder-1") + + editCalls := 0 + ch := &mockMessageEditor{ + editFn: func(_ context.Context, chatID, messageID, content string) error { + editCalls++ + if chatID != "123" || messageID != "placeholder-1" || content != "final streamed reply" { + t.Fatalf("unexpected placeholder edit for %s/%s: %q", chatID, messageID, content) + } + return nil + }, + } + + earlierVisible := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "earlier visible message", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + }, + }) + _, handled := m.preSend(context.Background(), "test", earlierVisible, ch) + if handled { + t.Fatal("expected earlier visible message to be delivered normally") + } + if editCalls != 0 { + t.Fatalf("placeholder edits after earlier visible message = %d, want 0", editCalls) + } + if _, ok := m.streamActive.Load("test:123"); !ok { + t.Fatal("expected streamActive marker to remain for final outbound") + } + if _, ok := m.streamAuxiliaryTombstones.Load("test:123"); !ok { + t.Fatal("expected auxiliary tombstone to remain") + } + if _, ok := m.placeholders.Load("test:123"); !ok { + t.Fatal("expected placeholder cleanup to remain deferred to final outbound") + } + + finalMsg := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "final streamed reply", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + Raw: map[string]string{ + "outbound_kind": "final", + }, + }, + }) + _, handled = m.preSend(context.Background(), "test", finalMsg, ch) + if !handled { + t.Fatal("expected final outbound message to consume streamActive marker") + } + if _, ok := m.streamActive.Load("test:123"); ok { + t.Fatal("expected streamActive marker to be cleared by final outbound message") + } + if editCalls != 1 { + t.Fatalf("placeholder edits after final outbound = %d, want 1", editCalls) + } +} + +func TestPreSend_StreamActiveDoesNotConsumeOtherSessionFinal(t *testing.T) { + m := newTestManager() + m.streamActive.Store("test:123", true) + m.RecordPlaceholder("test", "123", "placeholder-1") + + ch := &mockMessageEditor{ + editFn: func(_ context.Context, _, _, _ string) error { + t.Fatal("placeholder edit should remain deferred for the streaming session") + return nil + }, + } + + otherSessionFinal := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + SessionKey: "session-other", + Content: "other session final", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + Raw: map[string]string{ + "outbound_kind": "final", + }, + }, + }) + + _, handled := m.preSend(context.Background(), "test", otherSessionFinal, ch) + if handled { + t.Fatal("expected final outbound from a different session to be delivered normally") + } + if _, ok := m.streamActive.Load("test:123"); !ok { + t.Fatal("expected streaming marker to remain for the streaming session") + } + if _, ok := m.placeholders.Load("test:123"); !ok { + t.Fatal("expected placeholder cleanup to remain deferred to the streaming session") + } +} + func TestPreSendMedia_LeavesTrackedMessageForChannelSend(t *testing.T) { m := newTestManager() ch := &mockDeletingMediaChannel{} @@ -1610,7 +1881,7 @@ func TestGetStreamer_FinalizeDismissesTrackedToolFeedback(t *testing.T) { } m.channels["test"] = ch - streamer, ok := m.GetStreamer(context.Background(), "test", "123") + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") if !ok { t.Fatal("expected streamer to be available") } @@ -1625,6 +1896,400 @@ func TestGetStreamer_FinalizeDismissesTrackedToolFeedback(t *testing.T) { } } +func TestGetStreamer_FinalizeCleansPlaceholderImmediately(t *testing.T) { + m := newTestManager() + m.RecordPlaceholder("test", "123", "placeholder-1") + var editedContent string + editCalls := 0 + ch := &mockStreamingChannel{ + mockMessageEditor: mockMessageEditor{ + editFn: func(_ context.Context, chatID, messageID, content string) error { + if chatID != "123" || messageID != "placeholder-1" { + t.Fatalf("unexpected edit target: %s/%s", chatID, messageID) + } + editCalls++ + editedContent = content + return nil + }, + }, + streamer: &mockStreamer{}, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") + if !ok { + t.Fatal("expected streamer to be available") + } + if err := streamer.Finalize(context.Background(), "final reply"); err != nil { + t.Fatalf("Finalize() error = %v", err) + } + if editedContent != "final reply" { + t.Fatalf("edited placeholder content = %q, want final reply", editedContent) + } + if _, placeholderExists := m.placeholders.Load("test:123"); placeholderExists { + t.Fatal("expected placeholder to be cleaned up during finalize") + } + if _, streamActiveExists := m.streamActive.Load("test:123"); !streamActiveExists { + t.Fatal("expected streamActive marker to be recorded after finalize") + } + cleaner, ok := streamer.(interface{ ClearFinalizedStreamMarker() }) + if !ok { + t.Fatal("expected streamer to expose marker cleanup") + } + cleaner.ClearFinalizedStreamMarker() + if _, streamActiveExists := m.streamActive.Load("test:123"); streamActiveExists { + t.Fatal("expected streamActive marker to be cleared") + } + if _, ok := m.streamAuxiliaryTombstones.Load("test:123"); !ok { + t.Fatal("expected auxiliary tombstone to remain after final marker cleanup") + } + + lateThought := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "late reasoning", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + Raw: map[string]string{ + "message_kind": "thought", + }, + }, + }) + msgIDs, handled := m.preSend(context.Background(), "test", lateThought, ch) + if !handled { + t.Fatal("expected auxiliary tombstone to drop late thought") + } + if len(msgIDs) != 0 { + t.Fatalf("expected no delivered message IDs for late thought, got %v", msgIDs) + } + if editCalls != 1 { + t.Fatalf("expected late thought not to edit placeholder, got %d edits", editCalls) + } + + finalOutbound := testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "visible final reply", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + }, + }) + _, handled = m.preSend(context.Background(), "test", finalOutbound, ch) + if handled { + t.Fatal("expected cleared final marker to let normal outbound send") + } + if _, ok := m.streamAuxiliaryTombstones.Load("test:123"); ok { + t.Fatal("expected normal outbound to clear auxiliary tombstone") + } +} + +func TestGetStreamer_FinalizeCleansPlaceholderWithSessionKey(t *testing.T) { + m := newTestManager() + m.RecordPlaceholder("test", "123", "placeholder-1") + ch := &mockStreamingChannel{ + mockMessageEditor: mockMessageEditor{ + editFn: func(_ context.Context, chatID, messageID, content string) error { + if chatID != "123" || messageID != "placeholder-1" || content != "final reply" { + t.Fatalf("unexpected edit for %s/%s: %q", chatID, messageID, content) + } + return nil + }, + }, + streamer: &mockStreamer{}, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "session-1") + if !ok { + t.Fatal("expected streamer to be available") + } + if err := streamer.Finalize(context.Background(), "final reply"); err != nil { + t.Fatalf("Finalize() error = %v", err) + } + if _, placeholderExists := m.placeholders.Load("test:123"); placeholderExists { + t.Fatal("expected placeholder to be cleaned up during finalize") + } + if _, streamActiveExists := m.streamActive.Load("test:123:session-1"); !streamActiveExists { + t.Fatal("expected session streamActive marker to be recorded after finalize") + } +} + +func TestGetStreamer_PreservesContextUsageStreamer(t *testing.T) { + m := newTestManager() + var gotUsage *bus.ContextUsage + ch := &mockStreamingChannel{ + streamer: &mockStreamer{ + finalizeWithContextFn: func(_ context.Context, content string, usage *bus.ContextUsage) error { + if content != "final reply" { + t.Fatalf("unexpected finalize content: %q", content) + } + gotUsage = usage + return nil + }, + }, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") + if !ok { + t.Fatal("expected streamer to be available") + } + contextStreamer, ok := streamer.(bus.ContextUsageStreamer) + if !ok { + t.Fatal("manager-wrapped streamer should preserve ContextUsageStreamer") + } + usage := &bus.ContextUsage{UsedTokens: 10, TotalTokens: 100, CompressAtTokens: 80, UsedPercent: 10} + if err := contextStreamer.FinalizeWithContext(context.Background(), "final reply", usage); err != nil { + t.Fatalf("FinalizeWithContext() error = %v", err) + } + if gotUsage != usage { + t.Fatalf("context usage = %#v, want original usage", gotUsage) + } + if _, ok := m.streamActive.Load("test:123"); !ok { + t.Fatal("expected streamActive marker to be recorded after finalize with context") + } +} + +func TestGetStreamer_PreservesReasoningStreamer(t *testing.T) { + m := newTestManager() + inner := &mockReasoningStreamer{} + ch := &mockStreamingChannel{ + streamer: inner, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") + if !ok { + t.Fatal("expected streamer to be available") + } + reasoningStreamer, ok := streamer.(bus.ReasoningStreamer) + if !ok { + t.Fatal("manager-wrapped streamer should preserve ReasoningStreamer") + } + if err := reasoningStreamer.UpdateReasoning(context.Background(), "thinking"); err != nil { + t.Fatalf("UpdateReasoning() error = %v", err) + } + if err := reasoningStreamer.FinalizeReasoning(context.Background(), "final thought"); err != nil { + t.Fatalf("FinalizeReasoning() error = %v", err) + } + if got := inner.reasoningUpdates; len(got) != 1 || got[0] != "thinking" { + t.Fatalf("reasoning updates = %v, want [thinking]", got) + } + if inner.reasoningFinal != "final thought" { + t.Fatalf("reasoning final = %q, want final thought", inner.reasoningFinal) + } +} + +func TestGetStreamer_PreservesModelNameSetter(t *testing.T) { + m := newTestManager() + inner := &modelTrackingReasoningStreamer{} + ch := &mockStreamingChannel{ + streamer: inner, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") + if !ok { + t.Fatal("expected streamer to be available") + } + setter, ok := streamer.(interface{ SetModelName(modelName string) }) + if !ok { + t.Fatal("manager-wrapped streamer should preserve SetModelName") + } + setter.SetModelName("gpt-5.4") + if err := streamer.Update(context.Background(), "hello"); err != nil { + t.Fatalf("Update() error = %v", err) + } + reasoningStreamer, ok := streamer.(bus.ReasoningStreamer) + if !ok { + t.Fatal("manager-wrapped streamer should preserve ReasoningStreamer") + } + setter.SetModelName("gpt-5.4") + if err := reasoningStreamer.UpdateReasoning(context.Background(), "thinking"); err != nil { + t.Fatalf("UpdateReasoning() error = %v", err) + } + if len(inner.modelNames) != 2 { + t.Fatalf("model name calls = %v, want 2 forwarded calls", inner.modelNames) + } + if inner.modelNames[0] != "gpt-5.4" || inner.modelNames[1] != "gpt-5.4" { + t.Fatalf("model name calls = %v, want both forwarded as gpt-5.4", inner.modelNames) + } +} + +func TestGetStreamer_SplitOnMarkerStreamsSeparateSegments(t *testing.T) { + m := newTestManager() + m.config = &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + SplitOnMarker: true, + }, + }, + } + + var segments []*recordingStreamSegment + ch := &mockStreamingChannel{ + beginStreamFn: func(context.Context, string) (Streamer, error) { + segment := &recordingStreamSegment{} + segments = append(segments, segment) + return segment, nil + }, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "session-1") + if !ok { + t.Fatal("expected streamer to be available") + } + contextStreamer, ok := streamer.(bus.ContextUsageStreamer) + if !ok { + t.Fatal("split streamer should preserve ContextUsageStreamer") + } + + if err := streamer.Update(context.Background(), "hello"); err != nil { + t.Fatalf("Update(first) error = %v", err) + } + if err := streamer.Update(context.Background(), "hello<|[SPLIT]|>world"); err != nil { + t.Fatalf("Update(split) error = %v", err) + } + if err := streamer.Update(context.Background(), "hello<|[SPLIT]|>world!"); err != nil { + t.Fatalf("Update(second segment) error = %v", err) + } + usage := &bus.ContextUsage{UsedTokens: 10, TotalTokens: 100} + if err := contextStreamer.FinalizeWithContext( + context.Background(), + "hello<|[SPLIT]|>world!", + usage, + ); err != nil { + t.Fatalf("FinalizeWithContext() error = %v", err) + } + + if len(segments) != 2 { + t.Fatalf("segments = %d, want 2", len(segments)) + } + if got := segments[0].updates; len(got) != 1 || got[0] != "hello" { + t.Fatalf("segment 0 updates = %v, want [hello]", got) + } + if got := segments[0].finals; len(got) != 1 || got[0] != "hello" { + t.Fatalf("segment 0 finals = %v, want [hello]", got) + } + if got := segments[1].updates; len(got) != 2 || got[0] != "world" || got[1] != "world!" { + t.Fatalf("segment 1 updates = %v, want [world world!]", got) + } + if got := segments[1].finals; len(got) != 1 || got[0] != "world!" { + t.Fatalf("segment 1 finals = %v, want [world!]", got) + } + if segments[1].finalUsage != usage { + t.Fatalf("final usage = %#v, want original usage", segments[1].finalUsage) + } + if _, ok := m.streamActive.Load("test:123:session-1"); !ok { + t.Fatal("expected streamActive marker to be recorded after split stream finalize") + } +} + +func TestGetStreamer_SplitOnMarkerKeepsReasoningOnInitialStreamer(t *testing.T) { + m := newTestManager() + m.config = &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + SplitOnMarker: true, + }, + }, + } + + initial := &mockReasoningStreamer{} + next := &recordingStreamSegment{} + callCount := 0 + ch := &mockStreamingChannel{ + beginStreamFn: func(context.Context, string) (Streamer, error) { + callCount++ + if callCount == 1 { + return initial, nil + } + return next, nil + }, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") + if !ok { + t.Fatal("expected streamer to be available") + } + if err := streamer.Update(context.Background(), "hello<|[SPLIT]|>world"); err != nil { + t.Fatalf("Update() error = %v", err) + } + reasoningStreamer, ok := streamer.(bus.ReasoningStreamer) + if !ok { + t.Fatal("split streamer should preserve ReasoningStreamer") + } + if err := reasoningStreamer.UpdateReasoning(context.Background(), "thinking"); err != nil { + t.Fatalf("UpdateReasoning() error = %v", err) + } + if err := reasoningStreamer.FinalizeReasoning(context.Background(), "final thought"); err != nil { + t.Fatalf("FinalizeReasoning() error = %v", err) + } + + if got := initial.reasoningUpdates; len(got) != 1 || got[0] != "thinking" { + t.Fatalf("initial reasoning updates = %v, want [thinking]", got) + } + if initial.reasoningFinal != "final thought" { + t.Fatalf("initial reasoning final = %q, want final thought", initial.reasoningFinal) + } +} + +func TestGetStreamer_SplitOnMarkerPreservesModelNameSetter(t *testing.T) { + m := newTestManager() + m.config = &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + SplitOnMarker: true, + }, + }, + } + + initial := &modelTrackingReasoningStreamer{} + next := &recordingStreamSegment{} + callCount := 0 + ch := &mockStreamingChannel{ + beginStreamFn: func(context.Context, string) (Streamer, error) { + callCount++ + if callCount == 1 { + return initial, nil + } + return next, nil + }, + } + m.channels["test"] = ch + + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") + if !ok { + t.Fatal("expected streamer to be available") + } + setter, ok := streamer.(interface{ SetModelName(modelName string) }) + if !ok { + t.Fatal("split streamer should preserve SetModelName") + } + setter.SetModelName("gpt-5.4-mini") + if err := streamer.Update(context.Background(), "hello<|[SPLIT]|>world"); err != nil { + t.Fatalf("Update() error = %v", err) + } + reasoningStreamer, ok := streamer.(bus.ReasoningStreamer) + if !ok { + t.Fatal("split streamer should preserve ReasoningStreamer") + } + if err := reasoningStreamer.UpdateReasoning(context.Background(), "thinking"); err != nil { + t.Fatalf("UpdateReasoning() error = %v", err) + } + + if len(initial.modelNames) == 0 || initial.modelNames[0] != "gpt-5.4-mini" { + t.Fatalf("initial model names = %v, want forwarded gpt-5.4-mini", initial.modelNames) + } + if len(next.modelNames) == 0 || next.modelNames[0] != "gpt-5.4-mini" { + t.Fatalf("next model names = %v, want forwarded gpt-5.4-mini", next.modelNames) + } +} + func TestGetStreamer_FinalizeSeparateMessagesClearsTrackedToolFeedback(t *testing.T) { m := newTestManager() m.config = &config.Config{ @@ -1650,7 +2315,7 @@ func TestGetStreamer_FinalizeSeparateMessagesClearsTrackedToolFeedback(t *testin } m.channels["test"] = ch - streamer, ok := m.GetStreamer(context.Background(), "test", "123") + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") if !ok { t.Fatal("expected streamer to be available") } @@ -1692,7 +2357,7 @@ func TestGetStreamer_FinalizeDismissesResolvedTrackedToolFeedback(t *testing.T) } m.channels["test"] = ch - streamer, ok := m.GetStreamer(context.Background(), "test", "-100123/42") + streamer, ok := m.GetStreamer(context.Background(), "test", "-100123/42", "") if !ok { t.Fatal("expected streamer to be available") } @@ -1761,7 +2426,7 @@ func TestGetStreamer_FinalizeFailureDoesNotDismissTrackedToolFeedback(t *testing } m.channels["test"] = ch - streamer, ok := m.GetStreamer(context.Background(), "test", "123") + streamer, ok := m.GetStreamer(context.Background(), "test", "123", "") if !ok { t.Fatal("expected streamer to be available") } @@ -1839,6 +2504,68 @@ func TestRunWorker_ToolFeedbackSkipsMarkerSplitting(t *testing.T) { } } +func TestRunWorker_FinalizedStreamSuppressesMarkerSplitBeforeSending(t *testing.T) { + m := newTestManager() + m.config = &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + SplitOnMarker: true, + }, + }, + } + + var ( + mu sync.Mutex + received []string + ) + ch := &mockChannel{ + sendFn: func(_ context.Context, msg bus.OutboundMessage) error { + mu.Lock() + received = append(received, msg.Content) + mu.Unlock() + return nil + }, + } + + w := &channelWorker{ + ch: ch, + queue: make(chan bus.OutboundMessage, 1), + done: make(chan struct{}), + limiter: rate.NewLimiter(rate.Inf, 1), + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go m.runWorker(ctx, "test", w) + + streamKey := streamSuppressionKey("test", "123", "session-1") + m.streamActive.Store(streamKey, true) + w.queue <- testOutboundMessage(bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + SessionKey: "session-1", + Content: "streamed full reply<|[SPLIT]|>duplicate chunk", + Context: bus.InboundContext{ + Channel: "test", + ChatID: "123", + Raw: map[string]string{ + "outbound_kind": "final", + }, + }, + }) + + time.Sleep(100 * time.Millisecond) + + mu.Lock() + defer mu.Unlock() + if len(received) != 0 { + t.Fatalf("received split duplicate messages = %v, want none", received) + } + if _, ok := m.streamActive.Load(streamKey); ok { + t.Fatal("expected finalized stream marker to be consumed") + } +} + func TestPreSend_PlaceholderEditFails_FallsThrough(t *testing.T) { m := newTestManager() diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index fefb64ee..3c393898 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -325,6 +325,9 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri PayloadKeyContent: content, "message_id": msgID, } + if modelName := strings.TrimSpace(msg.Context.Raw[PayloadKeyModelName]); modelName != "" { + payload[PayloadKeyModelName] = modelName + } switch { case isThought: payload[PayloadKeyKind] = MessageKindThought @@ -359,6 +362,15 @@ func (c *PicoChannel) EditMessage(ctx context.Context, chatID string, messageID return c.editMessage(ctx, chatID, messageID, content, nil) } +func (c *PicoChannel) EditMessageWithPayload( + ctx context.Context, + chatID string, + messageID string, + payload map[string]any, +) error { + return c.editMessagePayload(ctx, chatID, messageID, payload, nil) +} + // DeleteMessage implements channels.MessageDeleter. func (c *PicoChannel) DeleteMessage(ctx context.Context, chatID string, messageID string) error { outMsg := newMessage(TypeMessageDelete, map[string]any{ @@ -419,14 +431,23 @@ func (c *PicoChannel) finalizeTrackedToolFeedbackMessage( ctx context.Context, chatID string, content string, - editFn func(context.Context, string, string, string, *bus.ContextUsage) error, + editFn func(context.Context, string, string, map[string]any, *bus.ContextUsage) error, + payload map[string]any, contextUsage *bus.ContextUsage, ) ([]string, bool) { msgID, baseContent, ok := c.takeToolFeedbackMessage(chatID) if !ok || editFn == nil { return nil, false } - if err := editFn(ctx, chatID, msgID, content, contextUsage); err != nil { + if payload == nil { + payload = map[string]any{ + PayloadKeyContent: content, + } + } + if _, ok := payload[PayloadKeyContent]; !ok { + payload[PayloadKeyContent] = content + } + if err := editFn(ctx, chatID, msgID, payload, contextUsage); err != nil { c.RecordToolFeedbackMessage(chatID, msgID, baseContent) return nil, false } @@ -437,7 +458,20 @@ func (c *PicoChannel) FinalizeToolFeedbackMessage(ctx context.Context, msg bus.O if !outboundMessageFinalizesTrackedToolFeedback(msg) { return nil, false } - return c.finalizeTrackedToolFeedbackMessage(ctx, msg.ChatID, msg.Content, c.editMessage, msg.ContextUsage) + payload := map[string]any{ + PayloadKeyContent: msg.Content, + } + if modelName := strings.TrimSpace(msg.Context.Raw[PayloadKeyModelName]); modelName != "" { + payload[PayloadKeyModelName] = modelName + } + return c.finalizeTrackedToolFeedbackMessage( + ctx, + msg.ChatID, + msg.Content, + c.editMessagePayload, + payload, + msg.ContextUsage, + ) } // StartTyping implements channels.TypingCapable. @@ -464,8 +498,9 @@ func (c *PicoChannel) SendPlaceholder(ctx context.Context, chatID string) (strin msgID := uuid.New().String() outMsg := newMessage(TypeMessageCreate, map[string]any{ - PayloadKeyContent: text, - "message_id": msgID, + PayloadKeyContent: text, + PayloadKeyPlaceholder: true, + "message_id": msgID, }) if err := c.broadcastToSession(chatID, outMsg); err != nil { @@ -475,6 +510,221 @@ func (c *PicoChannel) SendPlaceholder(ctx context.Context, chatID string) (strin return msgID, nil } +// BeginStream implements channels.StreamingCapable for Pico WebUI. +func (c *PicoChannel) BeginStream(ctx context.Context, chatID string) (channels.Streamer, error) { + if c == nil || c.config == nil || !c.config.Streaming.Enabled { + return nil, fmt.Errorf("streaming disabled in config") + } + if !c.IsRunning() { + return nil, channels.ErrNotRunning + } + streamCfg := c.config.Streaming.WithDefaults(0, 1) + return &picoStreamer{ + channel: c, + chatID: chatID, + throttleInterval: time.Duration(streamCfg.ThrottleSeconds) * time.Second, + minGrowth: streamCfg.MinGrowthChars, + }, nil +} + +type picoStreamer struct { + channel *PicoChannel + chatID string + modelName string + messageID string + reasoningID string + throttleInterval time.Duration + minGrowth int + lastLen int + lastAt time.Time + lastContent string + reasoningLastLen int + reasoningLastAt time.Time + reasoningContent string + mu sync.Mutex +} + +func (s *picoStreamer) SetModelName(modelName string) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.modelName = strings.TrimSpace(modelName) +} + +func (s *picoStreamer) Update(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + return s.updateLocked(ctx, content, false, nil) +} + +func (s *picoStreamer) Finalize(ctx context.Context, content string) error { + return s.FinalizeWithContext(ctx, content, nil) +} + +func (s *picoStreamer) FinalizeWithContext(ctx context.Context, content string, contextUsage *bus.ContextUsage) error { + s.mu.Lock() + defer s.mu.Unlock() + return s.updateLocked(ctx, content, true, contextUsage) +} + +func (s *picoStreamer) UpdateReasoning(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + return s.updateReasoningLocked(ctx, content, false) +} + +func (s *picoStreamer) FinalizeReasoning(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + return s.updateReasoningLocked(ctx, content, true) +} + +func (s *picoStreamer) Cancel(ctx context.Context) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if s.channel == nil || s.messageID == "" { + if s.channel != nil && s.reasoningID != "" { + _ = s.channel.DeleteMessage(ctx, s.chatID, s.reasoningID) + s.reasoningID = "" + } + return + } + _ = s.channel.DeleteMessage(ctx, s.chatID, s.messageID) + s.messageID = "" + if s.reasoningID != "" { + _ = s.channel.DeleteMessage(ctx, s.chatID, s.reasoningID) + s.reasoningID = "" + } +} + +func (s *picoStreamer) updateLocked( + ctx context.Context, + content string, + force bool, + contextUsage *bus.ContextUsage, +) error { + if s == nil || s.channel == nil { + return fmt.Errorf("streamer is not initialized") + } + if strings.TrimSpace(content) == "" && s.messageID == "" { + return nil + } + + now := time.Now() + contentLen := len([]rune(content)) + if s.messageID != "" && !force { + growth := contentLen - s.lastLen + if now.Sub(s.lastAt) < s.throttleInterval || growth < s.minGrowth { + return nil + } + } + + return s.sendLocked(ctx, content, contextUsage) +} + +func (s *picoStreamer) updateReasoningLocked(ctx context.Context, content string, force bool) error { + if s == nil || s.channel == nil { + return fmt.Errorf("streamer is not initialized") + } + if strings.TrimSpace(content) == "" && s.reasoningID == "" { + return nil + } + + now := time.Now() + contentLen := len([]rune(content)) + if s.reasoningID != "" && !force { + growth := contentLen - s.reasoningLastLen + if now.Sub(s.reasoningLastAt) < s.throttleInterval || growth < s.minGrowth { + return nil + } + } + + return s.sendReasoningLocked(ctx, content) +} + +func (s *picoStreamer) sendLocked(ctx context.Context, content string, contextUsage *bus.ContextUsage) error { + now := time.Now() + contentLen := len([]rune(content)) + + if s.messageID == "" { + s.messageID = uuid.New().String() + payload := map[string]any{ + PayloadKeyContent: content, + "message_id": s.messageID, + } + if s.modelName != "" { + payload[PayloadKeyModelName] = s.modelName + } + setContextUsagePayload(payload, contextUsage) + outMsg := newMessage(TypeMessageCreate, payload) + if err := s.channel.broadcastToSession(s.chatID, outMsg); err != nil { + return err + } + } else if content != s.lastContent || contextUsage != nil { + payload := map[string]any{ + PayloadKeyContent: content, + "message_id": s.messageID, + } + if s.modelName != "" { + payload[PayloadKeyModelName] = s.modelName + } + if err := s.channel.editMessagePayload(ctx, s.chatID, s.messageID, payload, contextUsage); err != nil { + return err + } + } + + s.lastContent = content + s.lastLen = contentLen + s.lastAt = now + return nil +} + +func (s *picoStreamer) sendReasoningLocked(ctx context.Context, content string) error { + now := time.Now() + contentLen := len([]rune(content)) + + if s.reasoningID == "" { + s.reasoningID = uuid.New().String() + payload := map[string]any{ + PayloadKeyContent: content, + "message_id": s.reasoningID, + PayloadKeyKind: MessageKindThought, + PayloadKeyThought: true, + } + if s.modelName != "" { + payload[PayloadKeyModelName] = s.modelName + } + outMsg := newMessage(TypeMessageCreate, payload) + if err := s.channel.broadcastToSession(s.chatID, outMsg); err != nil { + return err + } + } else if content != s.reasoningContent { + payload := map[string]any{ + PayloadKeyContent: content, + "message_id": s.reasoningID, + PayloadKeyKind: MessageKindThought, + PayloadKeyThought: true, + } + if s.modelName != "" { + payload[PayloadKeyModelName] = s.modelName + } + outMsg := newMessage(TypeMessageUpdate, payload) + if err := s.channel.broadcastToSession(s.chatID, outMsg); err != nil { + return err + } + } + + s.reasoningContent = content + s.reasoningLastLen = contentLen + s.reasoningLastAt = now + return nil +} + // SendMedia implements channels.MediaSender for the Pico web UI. // Media is delivered as a normal assistant message carrying structured // attachments plus an authenticated same-origin download URL. @@ -554,6 +804,9 @@ func (c *PicoChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessag "attachments": attachments, "message_id": msgID, }) + if modelName := strings.TrimSpace(msg.Context.Raw[PayloadKeyModelName]); modelName != "" { + outMsg.Payload[PayloadKeyModelName] = modelName + } if err := c.broadcastToSession(msg.ChatID, outMsg); err != nil { return nil, err @@ -1168,11 +1421,30 @@ func (c *PicoChannel) editMessage( content string, contextUsage *bus.ContextUsage, ) error { - payload := map[string]any{ - "message_id": messageID, - "content": content, + return c.editMessagePayload(ctx, chatID, messageID, map[string]any{ + PayloadKeyContent: content, + }, contextUsage) +} + +func (c *PicoChannel) editMessagePayload( + ctx context.Context, + chatID string, + messageID string, + payload map[string]any, + contextUsage *bus.ContextUsage, +) error { + if payload == nil { + payload = map[string]any{} } - setContextUsagePayload(payload, contextUsage) - outMsg := newMessage(TypeMessageUpdate, payload) + normalized := make(map[string]any, len(payload)+1) + for key, value := range payload { + normalized[key] = value + } + if _, ok := normalized[PayloadKeyContent]; !ok { + normalized[PayloadKeyContent] = "" + } + normalized["message_id"] = messageID + setContextUsagePayload(normalized, contextUsage) + outMsg := newMessage(TypeMessageUpdate, normalized) return c.broadcastToSession(chatID, outMsg) } diff --git a/pkg/channels/pico/pico_test.go b/pkg/channels/pico/pico_test.go index bbe73a22..a793d7ad 100644 --- a/pkg/channels/pico/pico_test.go +++ b/pkg/channels/pico/pico_test.go @@ -36,6 +36,39 @@ func newTestPicoChannel(t *testing.T) *PicoChannel { return ch } +func TestHandleMessageSend_ForwardsMessageMetadata(t *testing.T) { + msgBus := bus.NewMessageBus() + bc := &config.Channel{Type: config.ChannelPico, Enabled: true} + cfg := &config.PicoSettings{} + cfg.SetToken("test-token") + ch, err := NewPicoChannel(bc, cfg, msgBus) + if err != nil { + t.Fatalf("NewPicoChannel: %v", err) + } + ch.ctx = context.Background() + + ch.handleMessageSend(&picoConn{id: "conn-1", sessionID: "sess-1"}, PicoMessage{ + Type: TypeMessageSend, + ID: "msg-1", + SessionID: "sess-1", + Payload: map[string]any{ + PayloadKeyContent: "hello", + }, + }) + + select { + case inbound := <-msgBus.InboundChan(): + if inbound.Content != "hello" { + t.Fatalf("content = %q, want hello", inbound.Content) + } + if got := inbound.Context.Raw["session_id"]; got != "sess-1" { + t.Fatalf("session_id raw = %q, want sess-1", got) + } + case <-time.After(time.Second): + t.Fatal("expected inbound pico message") + } +} + func TestFinalizeTrackedToolFeedbackMessage_StopsTrackingBeforeEdit(t *testing.T) { ch := &PicoChannel{ progress: channels.NewToolFeedbackAnimator(nil), @@ -46,12 +79,15 @@ func TestFinalizeTrackedToolFeedbackMessage_StopsTrackingBeforeEdit(t *testing.T context.Background(), "pico:chat-1", "final reply", - func(_ context.Context, chatID, messageID, content string, contextUsage *bus.ContextUsage) error { + func(_ context.Context, chatID, messageID string, payload map[string]any, contextUsage *bus.ContextUsage) error { if _, ok := ch.currentToolFeedbackMessage(chatID); ok { t.Fatal("expected tracked tool feedback to be stopped before edit") } - if chatID != "pico:chat-1" || messageID != "msg-1" || content != "final reply" { - t.Fatalf("unexpected edit args: %s %s %s", chatID, messageID, content) + if chatID != "pico:chat-1" || messageID != "msg-1" { + t.Fatalf("unexpected edit args: %s %s", chatID, messageID) + } + if got := payload[PayloadKeyContent]; got != "final reply" { + t.Fatalf("unexpected content payload: %#v", got) } if contextUsage != nil { t.Fatalf("unexpected context usage: %+v", contextUsage) @@ -59,6 +95,7 @@ func TestFinalizeTrackedToolFeedbackMessage_StopsTrackingBeforeEdit(t *testing.T return nil }, nil, + nil, ) if !handled { t.Fatal("expected finalizeTrackedToolFeedbackMessage to handle tracked message") @@ -115,7 +152,8 @@ func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) { Channel: "pico", ChatID: "pico:sess-1", Raw: map[string]string{ - "message_kind": MessageKindThought, + "message_kind": MessageKindThought, + PayloadKeyModelName: "gpt-5.4-mini", }, }, }); err != nil { @@ -134,6 +172,9 @@ func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) { if got := payload[PayloadKeyKind]; got != MessageKindThought { t.Fatalf("thought kind = %#v, want %q", got, MessageKindThought) } + if got := payload[PayloadKeyModelName]; got != "gpt-5.4-mini" { + t.Fatalf("thought model_name = %#v, want %q", got, "gpt-5.4-mini") + } if got := payload["message_id"]; got == "msg-progress" || got == nil || got == "" { t.Fatalf("thought message_id = %#v, want new non-progress id", got) } @@ -151,6 +192,9 @@ func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) { Context: bus.InboundContext{ Channel: "pico", ChatID: "pico:sess-1", + Raw: map[string]string{ + PayloadKeyModelName: "gpt-5.4", + }, }, ContextUsage: &bus.ContextUsage{ UsedTokens: 321, @@ -174,6 +218,9 @@ func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) { if got := payload[PayloadKeyContent]; got != "final reply" { t.Fatalf("final content = %#v, want %q", got, "final reply") } + if got := payload[PayloadKeyModelName]; got != "gpt-5.4" { + t.Fatalf("final model_name = %#v, want %q", got, "gpt-5.4") + } rawUsage, ok := payload["context_usage"].(map[string]any) if !ok { t.Fatalf("final context_usage = %#v, want map payload", payload["context_usage"]) @@ -193,6 +240,54 @@ func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) { } } +func TestSend_ToolCallsMessageIncludesModelName(t *testing.T) { + ch := newTestPicoChannel(t) + + if err := ch.Start(context.Background()); err != nil { + t.Fatalf("Start() error = %v", err) + } + defer ch.Stop(context.Background()) + + clientConn, received, cleanup := newTestPicoWebSocket(t) + defer cleanup() + ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"}) + + if _, err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "pico:sess-1", + Content: "", + Context: bus.InboundContext{ + Channel: "pico", + ChatID: "pico:sess-1", + Raw: map[string]string{ + "message_kind": MessageKindToolCalls, + PayloadKeyModelName: "gpt-5.4", + PayloadKeyToolCalls: `[{"id":"call_1","type":"function","function":{"name":"read_file","arguments":"{\"path\":\"README.md\"}"}}]`, + }, + }, + }); err != nil { + t.Fatalf("Send(tool_calls) error = %v", err) + } + + select { + case msg := <-received: + if msg.Type != TypeMessageCreate { + t.Fatalf("tool_calls message type = %q, want %q", msg.Type, TypeMessageCreate) + } + payload := msg.Payload + if got := payload[PayloadKeyKind]; got != MessageKindToolCalls { + t.Fatalf("tool_calls kind = %#v, want %q", got, MessageKindToolCalls) + } + if got := payload[PayloadKeyModelName]; got != "gpt-5.4" { + t.Fatalf("tool_calls model_name = %#v, want %q", got, "gpt-5.4") + } + if _, ok := payload[PayloadKeyToolCalls].([]any); !ok { + t.Fatalf("tool_calls payload = %#v, want parsed array", payload[PayloadKeyToolCalls]) + } + case <-time.After(time.Second): + t.Fatal("expected tool_calls message to be delivered") + } +} + func TestSendPlaceholder_EmitsNormalMessageWithoutKind(t *testing.T) { ch := newTestPicoChannel(t) ch.bc.Placeholder.Enabled = true @@ -226,6 +321,9 @@ func TestSendPlaceholder_EmitsNormalMessageWithoutKind(t *testing.T) { if got := payload[PayloadKeyContent]; got != "Thinking..." { t.Fatalf("placeholder content = %#v, want %q", got, "Thinking...") } + if got := payload[PayloadKeyPlaceholder]; got != true { + t.Fatalf("placeholder marker = %#v, want true", got) + } if got, ok := payload[PayloadKeyKind]; ok { t.Fatalf("placeholder kind = %#v, want absent", got) } @@ -234,6 +332,303 @@ func TestSendPlaceholder_EmitsNormalMessageWithoutKind(t *testing.T) { } } +func TestBeginStream_CreatesAndUpdatesSameMessage(t *testing.T) { + ch := newTestPicoChannel(t) + ch.config.Streaming = config.StreamingConfig{ + Enabled: true, + ThrottleSeconds: 1, + MinGrowthChars: 1, + } + if err := ch.Start(context.Background()); err != nil { + t.Fatalf("Start() error = %v", err) + } + defer ch.Stop(context.Background()) + + clientConn, received, cleanup := newTestPicoWebSocket(t) + defer cleanup() + ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"}) + + streamer, err := ch.BeginStream(context.Background(), "pico:sess-1") + if err != nil { + t.Fatalf("BeginStream() error = %v", err) + } + if setter, ok := streamer.(interface{ SetModelName(modelName string) }); ok { + setter.SetModelName("gpt-5.4") + } + if err := streamer.Update(context.Background(), "hello"); err != nil { + t.Fatalf("Update(first) error = %v", err) + } + first := mustReceivePicoMessage(t, received) + if first.Type != TypeMessageCreate { + t.Fatalf("first type = %q, want %q", first.Type, TypeMessageCreate) + } + msgID, _ := first.Payload["message_id"].(string) + if msgID == "" { + t.Fatalf("first message_id = %#v, want non-empty", first.Payload["message_id"]) + } + if got := first.Payload[PayloadKeyContent]; got != "hello" { + t.Fatalf("first content = %#v, want hello", got) + } + if got := first.Payload[PayloadKeyModelName]; got != "gpt-5.4" { + t.Fatalf("first model_name = %#v, want %q", got, "gpt-5.4") + } + + rawStreamer := streamer.(*picoStreamer) + rawStreamer.mu.Lock() + rawStreamer.lastAt = time.Now().Add(-2 * time.Second) + rawStreamer.mu.Unlock() + secondContent := "hello world with enough growth to pass the default streaming threshold" + if err := streamer.Update(context.Background(), secondContent); err != nil { + t.Fatalf("Update(second) error = %v", err) + } + second := mustReceivePicoMessage(t, received) + if second.Type != TypeMessageUpdate { + t.Fatalf("second type = %q, want %q", second.Type, TypeMessageUpdate) + } + if got := second.Payload["message_id"]; got != msgID { + t.Fatalf("second message_id = %#v, want %q", got, msgID) + } + if got := second.Payload[PayloadKeyContent]; got != secondContent { + t.Fatalf("second content = %#v, want %q", got, secondContent) + } + if got := second.Payload[PayloadKeyModelName]; got != "gpt-5.4" { + t.Fatalf("second model_name = %#v, want %q", got, "gpt-5.4") + } +} + +func TestBeginStream_DefaultStreamingShowsSmallIncrements(t *testing.T) { + ch := newTestPicoChannel(t) + ch.config.Streaming = config.StreamingConfig{Enabled: true} + if err := ch.Start(context.Background()); err != nil { + t.Fatalf("Start() error = %v", err) + } + defer ch.Stop(context.Background()) + + clientConn, received, cleanup := newTestPicoWebSocket(t) + defer cleanup() + ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"}) + + streamer, err := ch.BeginStream(context.Background(), "pico:sess-1") + if err != nil { + t.Fatalf("BeginStream() error = %v", err) + } + if err := streamer.Update(context.Background(), "h"); err != nil { + t.Fatalf("Update(first) error = %v", err) + } + first := mustReceivePicoMessage(t, received) + if first.Type != TypeMessageCreate { + t.Fatalf("first type = %q, want %q", first.Type, TypeMessageCreate) + } + msgID, _ := first.Payload["message_id"].(string) + if msgID == "" { + t.Fatalf("first message_id = %#v, want non-empty", first.Payload["message_id"]) + } + + if err := streamer.Update(context.Background(), "he"); err != nil { + t.Fatalf("Update(second) error = %v", err) + } + second := mustReceivePicoMessage(t, received) + if second.Type != TypeMessageUpdate { + t.Fatalf("second type = %q, want %q", second.Type, TypeMessageUpdate) + } + if got := second.Payload["message_id"]; got != msgID { + t.Fatalf("second message_id = %#v, want %q", got, msgID) + } + if got := second.Payload[PayloadKeyContent]; got != "he" { + t.Fatalf("second content = %#v, want he", got) + } +} + +func TestBeginStream_StreamsReasoningAsThoughtUpdates(t *testing.T) { + ch := newTestPicoChannel(t) + ch.config.Streaming = config.StreamingConfig{Enabled: true} + if err := ch.Start(context.Background()); err != nil { + t.Fatalf("Start() error = %v", err) + } + defer ch.Stop(context.Background()) + + clientConn, received, cleanup := newTestPicoWebSocket(t) + defer cleanup() + ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"}) + + streamer, err := ch.BeginStream(context.Background(), "pico:sess-1") + if err != nil { + t.Fatalf("BeginStream() error = %v", err) + } + reasoningStreamer, ok := streamer.(bus.ReasoningStreamer) + if !ok { + t.Fatal("pico stream should support reasoning updates") + } + if setter, ok := streamer.(interface{ SetModelName(modelName string) }); ok { + setter.SetModelName("gpt-5.4-mini") + } + if err := reasoningStreamer.UpdateReasoning(context.Background(), "thinking"); err != nil { + t.Fatalf("UpdateReasoning(first) error = %v", err) + } + first := mustReceivePicoMessage(t, received) + if first.Type != TypeMessageCreate { + t.Fatalf("first type = %q, want %q", first.Type, TypeMessageCreate) + } + msgID, _ := first.Payload["message_id"].(string) + if msgID == "" { + t.Fatalf("first message_id = %#v, want non-empty", first.Payload["message_id"]) + } + if got := first.Payload[PayloadKeyKind]; got != MessageKindThought { + t.Fatalf("first kind = %#v, want %q", got, MessageKindThought) + } + if got := first.Payload[PayloadKeyContent]; got != "thinking" { + t.Fatalf("first content = %#v, want thinking", got) + } + if got := first.Payload[PayloadKeyModelName]; got != "gpt-5.4-mini" { + t.Fatalf("first model_name = %#v, want %q", got, "gpt-5.4-mini") + } + + if err := reasoningStreamer.UpdateReasoning(context.Background(), "thinking more"); err != nil { + t.Fatalf("UpdateReasoning(second) error = %v", err) + } + second := mustReceivePicoMessage(t, received) + if second.Type != TypeMessageUpdate { + t.Fatalf("second type = %q, want %q", second.Type, TypeMessageUpdate) + } + if got := second.Payload["message_id"]; got != msgID { + t.Fatalf("second message_id = %#v, want %q", got, msgID) + } + if got := second.Payload[PayloadKeyKind]; got != MessageKindThought { + t.Fatalf("second kind = %#v, want %q", got, MessageKindThought) + } + if got := second.Payload[PayloadKeyContent]; got != "thinking more" { + t.Fatalf("second content = %#v, want thinking more", got) + } + if got := second.Payload[PayloadKeyModelName]; got != "gpt-5.4-mini" { + t.Fatalf("second model_name = %#v, want %q", got, "gpt-5.4-mini") + } +} + +func TestBeginStream_ThrottlesIntermediateUpdatesAndFinalFlushes(t *testing.T) { + ch := newTestPicoChannel(t) + ch.config.Streaming = config.StreamingConfig{ + Enabled: true, + ThrottleSeconds: 60, + MinGrowthChars: 100, + } + if err := ch.Start(context.Background()); err != nil { + t.Fatalf("Start() error = %v", err) + } + defer ch.Stop(context.Background()) + + clientConn, received, cleanup := newTestPicoWebSocket(t) + defer cleanup() + ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"}) + + streamer, err := ch.BeginStream(context.Background(), "pico:sess-1") + if err != nil { + t.Fatalf("BeginStream() error = %v", err) + } + if err := streamer.Update(context.Background(), "first"); err != nil { + t.Fatalf("Update(first) error = %v", err) + } + if err := streamer.Update(context.Background(), "first plus short growth"); err != nil { + t.Fatalf("Update(throttled) error = %v", err) + } + if err := streamer.Update(context.Background(), "first"+strings.Repeat("x", 120)); err != nil { + t.Fatalf("Update(enough growth too soon) error = %v", err) + } + + first := mustReceivePicoMessage(t, received) + if first.Type != TypeMessageCreate { + t.Fatalf("first type = %q, want %q", first.Type, TypeMessageCreate) + } + msgID, _ := first.Payload["message_id"].(string) + assertNoPicoMessage(t, received) + + rawStreamer := streamer.(*picoStreamer) + rawStreamer.mu.Lock() + rawStreamer.lastAt = time.Now().Add(-61 * time.Second) + rawStreamer.mu.Unlock() + if err := streamer.Update(context.Background(), "first plus small growth"); err != nil { + t.Fatalf("Update(enough time too little growth) error = %v", err) + } + assertNoPicoMessage(t, received) + + if err := streamer.Finalize(context.Background(), "first plus final text"); err != nil { + t.Fatalf("Finalize() error = %v", err) + } + final := mustReceivePicoMessage(t, received) + if final.Type != TypeMessageUpdate { + t.Fatalf("final type = %q, want %q", final.Type, TypeMessageUpdate) + } + if got := final.Payload["message_id"]; got != msgID { + t.Fatalf("final message_id = %#v, want %q", got, msgID) + } + if got := final.Payload[PayloadKeyContent]; got != "first plus final text" { + t.Fatalf("final content = %#v, want final text", got) + } + assertNoPicoMessage(t, received) +} + +func TestBeginStream_FinalizeIncludesContextUsage(t *testing.T) { + ch := newTestPicoChannel(t) + ch.config.Streaming = config.StreamingConfig{ + Enabled: true, + ThrottleSeconds: 0, + MinGrowthChars: 0, + } + if err := ch.Start(context.Background()); err != nil { + t.Fatalf("Start() error = %v", err) + } + defer ch.Stop(context.Background()) + + clientConn, received, cleanup := newTestPicoWebSocket(t) + defer cleanup() + ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"}) + + streamer, err := ch.BeginStream(context.Background(), "pico:sess-1") + if err != nil { + t.Fatalf("BeginStream() error = %v", err) + } + if setter, ok := streamer.(interface{ SetModelName(modelName string) }); ok { + setter.SetModelName("gpt-5.4") + } + if err := streamer.Update(context.Background(), "partial"); err != nil { + t.Fatalf("Update() error = %v", err) + } + first := mustReceivePicoMessage(t, received) + msgID, _ := first.Payload["message_id"].(string) + + contextStreamer, ok := streamer.(interface { + FinalizeWithContext(ctx context.Context, content string, usage *bus.ContextUsage) error + }) + if !ok { + t.Fatal("streamer should support FinalizeWithContext") + } + if err := contextStreamer.FinalizeWithContext(context.Background(), "final", &bus.ContextUsage{ + UsedTokens: 10, + TotalTokens: 100, + CompressAtTokens: 80, + UsedPercent: 10, + }); err != nil { + t.Fatalf("FinalizeWithContext() error = %v", err) + } + + final := mustReceivePicoMessage(t, received) + if final.Type != TypeMessageUpdate { + t.Fatalf("final type = %q, want %q", final.Type, TypeMessageUpdate) + } + if got := final.Payload["message_id"]; got != msgID { + t.Fatalf("final message_id = %#v, want %q", got, msgID) + } + if got := final.Payload[PayloadKeyModelName]; got != "gpt-5.4" { + t.Fatalf("final model_name = %#v, want %q", got, "gpt-5.4") + } + rawUsage, ok := final.Payload["context_usage"].(map[string]any) + if !ok { + t.Fatalf("final context_usage = %#v, want map", final.Payload["context_usage"]) + } + if got := rawUsage["used_tokens"]; got != float64(10) { + t.Fatalf("used_tokens = %#v, want 10", got) + } +} + func TestCreateAndAddConnection_RespectsMaxConnectionsConcurrently(t *testing.T) { ch := newTestPicoChannel(t) @@ -491,6 +886,26 @@ func TestHandleMediaDownload_ServesStoredFile(t *testing.T) { } } +func mustReceivePicoMessage(t *testing.T, received <-chan PicoMessage) PicoMessage { + t.Helper() + select { + case msg := <-received: + return msg + case <-time.After(time.Second): + t.Fatal("expected pico message") + } + return PicoMessage{} +} + +func assertNoPicoMessage(t *testing.T, received <-chan PicoMessage) { + t.Helper() + select { + case msg := <-received: + t.Fatalf("unexpected pico message: %+v", msg) + case <-time.After(150 * time.Millisecond): + } +} + func (c *PicoChannel) addConnForTest(pc *picoConn) { c.connsMu.Lock() defer c.connsMu.Unlock() diff --git a/pkg/channels/pico/protocol.go b/pkg/channels/pico/protocol.go index 6e3a5ca8..0c809430 100644 --- a/pkg/channels/pico/protocol.go +++ b/pkg/channels/pico/protocol.go @@ -22,10 +22,12 @@ const ( TypeError = "error" TypePong = "pong" - PayloadKeyContent = "content" - PayloadKeyThought = "thought" - PayloadKeyKind = "kind" - PayloadKeyToolCalls = "tool_calls" + PayloadKeyContent = "content" + PayloadKeyThought = "thought" + PayloadKeyKind = "kind" + PayloadKeyPlaceholder = "placeholder" + PayloadKeyToolCalls = "tool_calls" + PayloadKeyModelName = "model_name" MessageKindThought = "thought" MessageKindToolCalls = "tool_calls" diff --git a/pkg/channels/telegram/telegram.go b/pkg/channels/telegram/telegram.go index 0965bced..45672e5e 100644 --- a/pkg/channels/telegram/telegram.go +++ b/pkg/channels/telegram/telegram.go @@ -1471,7 +1471,7 @@ func (c *TelegramChannel) BeginStream(ctx context.Context, chatID string) (chann return nil, err } - streamCfg := c.tgCfg.Streaming + streamCfg := c.tgCfg.Streaming.WithDefaults(3, 200) return &telegramStreamer{ bot: c.bot, chatID: cid, @@ -1483,8 +1483,8 @@ func (c *TelegramChannel) BeginStream(ctx context.Context, chatID string) (chann } // telegramStreamer streams partial LLM output via Telegram's sendMessageDraft API. -// On first API error (e.g. bot lacks forum mode), it silently degrades: Update -// becomes a no-op, while Finalize still delivers the final message. +// Draft update failures are returned to the agent, which decides whether the +// stream was already visible enough to keep or should fall back to Chat(). type telegramStreamer struct { bot *telego.Bot chatID int64 @@ -1495,6 +1495,7 @@ type telegramStreamer struct { lastLen int lastAt time.Time failed bool + draftTouched bool mu sync.Mutex } @@ -1503,7 +1504,7 @@ func (s *telegramStreamer) Update(ctx context.Context, content string) error { defer s.mu.Unlock() if s.failed { - return nil + return fmt.Errorf("telegram streaming disabled after previous draft failure") } // Throttle: skip if not enough time or content has passed @@ -1514,6 +1515,7 @@ func (s *telegramStreamer) Update(ctx context.Context, content string) error { } htmlContent := markdownToTelegramHTML(content) + s.draftTouched = true err := s.bot.SendMessageDraft(ctx, &telego.SendMessageDraftParams{ ChatID: s.chatID, @@ -1523,12 +1525,11 @@ func (s *telegramStreamer) Update(ctx context.Context, content string) error { ParseMode: telego.ModeHTML, }) if err != nil { - // First error → degrade silently (e.g. no forum mode) logger.WarnCF("telegram", "sendMessageDraft failed, disabling streaming", map[string]any{ "error": err.Error(), }) s.failed = true - return nil // don't propagate — Finalize will still deliver + return fmt.Errorf("telegram draft update: %w", err) } s.lastLen = len(content) @@ -1554,11 +1555,33 @@ func (s *telegramStreamer) Finalize(ctx context.Context, content string) error { return fmt.Errorf("telegram finalize: %w", err) } } + s.Cancel(ctx) return nil } func (s *telegramStreamer) Cancel(ctx context.Context) { - // Draft auto-expires on Telegram's side; nothing to clean up. + s.mu.Lock() + defer s.mu.Unlock() + s.clearDraft(ctx) +} + +func (s *telegramStreamer) clearDraft(ctx context.Context) { + if !s.draftTouched { + return + } + if err := s.bot.SendMessageDraft(ctx, &telego.SendMessageDraftParams{ + ChatID: s.chatID, + MessageThreadID: s.threadID, + DraftID: s.draftID, + Text: " ", + }); err != nil { + logger.DebugCF("telegram", "failed to clear streaming draft", map[string]any{ + "chat_id": s.chatID, + "error": err.Error(), + }) + } + s.lastLen = 0 + s.draftTouched = false } // cryptoRandInt returns a non-zero random int using crypto/rand. diff --git a/pkg/channels/telegram/telegram_test.go b/pkg/channels/telegram/telegram_test.go index 14d02506..0ebde132 100644 --- a/pkg/channels/telegram/telegram_test.go +++ b/pkg/channels/telegram/telegram_test.go @@ -762,6 +762,120 @@ func TestBeginStream_UpdateUsesForumThreadID(t *testing.T) { assert.Equal(t, "partial", params.Text) } +func TestBeginStream_UsesDefaultThrottleWhenOnlyEnabled(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return &ta.Response{Ok: true, Result: []byte("true")}, nil + }, + } + ch := newTestChannel(t, caller) + ch.tgCfg.Streaming = config.StreamingConfig{Enabled: true} + + streamer, err := ch.BeginStream(context.Background(), "12345") + require.NoError(t, err) + require.NoError(t, streamer.Update(context.Background(), "partial")) + require.NoError(t, streamer.Update(context.Background(), "partial plus one")) + + require.Len(t, caller.calls, 1, "second small update should be throttled by defaults") +} + +func TestBeginStream_UpdateReturnsErrorWhenDraftFails(t *testing.T) { + callCount := 0 + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + callCount++ + if callCount == 1 { + return nil, errors.New("draft unsupported") + } + return &ta.Response{Ok: true, Result: []byte("true")}, nil + }, + } + ch := newTestChannel(t, caller) + ch.tgCfg.Streaming = config.StreamingConfig{Enabled: true} + + streamer, err := ch.BeginStream(context.Background(), "12345") + require.NoError(t, err) + + err = streamer.Update(context.Background(), "partial") + require.Error(t, err) + assert.Contains(t, err.Error(), "draft unsupported") + + streamer.Cancel(context.Background()) + require.Len(t, caller.calls, 2) + assert.Contains(t, caller.calls[1].URL, "sendMessageDraft") + + var params struct { + ChatID int64 `json:"chat_id"` + DraftID int `json:"draft_id"` + Text string `json:"text"` + } + require.NoError(t, json.Unmarshal(caller.calls[1].Data.BodyRaw, ¶ms)) + assert.Equal(t, int64(12345), params.ChatID) + assert.NotZero(t, params.DraftID) + assert.Equal(t, " ", params.Text) +} + +func TestBeginStream_CancelClearsExistingDraft(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return &ta.Response{Ok: true, Result: []byte("true")}, nil + }, + } + ch := newTestChannel(t, caller) + ch.tgCfg.Streaming = config.StreamingConfig{Enabled: true} + + streamer, err := ch.BeginStream(context.Background(), "12345") + require.NoError(t, err) + require.NoError(t, streamer.Update(context.Background(), "partial")) + streamer.Cancel(context.Background()) + + require.Len(t, caller.calls, 2) + assert.Contains(t, caller.calls[1].URL, "sendMessageDraft") + + var params struct { + ChatID int64 `json:"chat_id"` + DraftID int `json:"draft_id"` + Text string `json:"text"` + } + require.NoError(t, json.Unmarshal(caller.calls[1].Data.BodyRaw, ¶ms)) + assert.Equal(t, int64(12345), params.ChatID) + assert.NotZero(t, params.DraftID) + assert.Equal(t, " ", params.Text) +} + +func TestBeginStream_FinalizeClearsExistingDraft(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + if strings.Contains(url, "sendMessage") && !strings.Contains(url, "sendMessageDraft") { + return successResponse(t), nil + } + return &ta.Response{Ok: true, Result: []byte("true")}, nil + }, + } + ch := newTestChannel(t, caller) + ch.tgCfg.Streaming = config.StreamingConfig{Enabled: true} + + streamer, err := ch.BeginStream(context.Background(), "12345") + require.NoError(t, err) + require.NoError(t, streamer.Update(context.Background(), "partial")) + require.NoError(t, streamer.Finalize(context.Background(), "final")) + + require.Len(t, caller.calls, 3) + assert.Contains(t, caller.calls[0].URL, "sendMessageDraft") + assert.Contains(t, caller.calls[1].URL, "sendMessage") + assert.Contains(t, caller.calls[2].URL, "sendMessageDraft") + + var params struct { + ChatID int64 `json:"chat_id"` + DraftID int `json:"draft_id"` + Text string `json:"text"` + } + require.NoError(t, json.Unmarshal(caller.calls[2].Data.BodyRaw, ¶ms)) + assert.Equal(t, int64(12345), params.ChatID) + assert.NotZero(t, params.DraftID) + assert.Equal(t, " ", params.Text) +} + func TestBeginStream_FinalizeUsesForumThreadID(t *testing.T) { caller := &stubCaller{ callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { diff --git a/pkg/channels/wecom/wecom.go b/pkg/channels/wecom/wecom.go index a0a23fed..f6252af4 100644 --- a/pkg/channels/wecom/wecom.go +++ b/pkg/channels/wecom/wecom.go @@ -164,6 +164,9 @@ func (c *WeComChannel) Stop(_ context.Context) error { } func (c *WeComChannel) BeginStream(_ context.Context, chatID string) (channels.Streamer, error) { + if c == nil || c.config == nil || !c.config.Streaming.Enabled { + return nil, fmt.Errorf("streaming disabled in config") + } if !c.IsRunning() { return nil, channels.ErrNotRunning } diff --git a/pkg/channels/wecom/wecom_test.go b/pkg/channels/wecom/wecom_test.go index 85a2f6ef..0b8bcc7a 100644 --- a/pkg/channels/wecom/wecom_test.go +++ b/pkg/channels/wecom/wecom_test.go @@ -100,6 +100,7 @@ func TestBeginStream_UpdateAndFinalize(t *testing.T) { t.Parallel() ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.config.Streaming.Enabled = true ch.SetRunning(true) ch.queueTurn("chat-1", wecomTurn{ ReqID: "req-1", @@ -158,6 +159,24 @@ func TestBeginStream_UpdateAndFinalize(t *testing.T) { } } +func TestBeginStream_RequiresStreamingEnabled(t *testing.T) { + t.Parallel() + + ch := newTestWeComChannel(t, bus.NewMessageBus()) + ch.SetRunning(true) + ch.queueTurn("chat-1", wecomTurn{ + ReqID: "req-1", + ChatID: "chat-1", + ChatType: 1, + StreamID: "stream-1", + CreatedAt: time.Now(), + }) + + if _, err := ch.BeginStream(context.Background(), "chat-1"); err == nil { + t.Fatal("BeginStream() error = nil, want disabled streaming error") + } +} + func TestSend_StreamFailureFallsBackToActualChatID(t *testing.T) { t.Parallel() diff --git a/pkg/config/config.go b/pkg/config/config.go index 8ed3da13..d9608d11 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -399,6 +399,7 @@ type AgentDefaults struct { SplitOnMarker bool `json:"split_on_marker" env:"PICOCLAW_AGENTS_DEFAULTS_SPLIT_ON_MARKER"` // split messages on <|[SPLIT]|> marker ContextManager string `json:"context_manager,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_CONTEXT_MANAGER"` ContextManagerConfig json.RawMessage `json:"context_manager_config,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_CONTEXT_MANAGER_CONFIG"` + TurnProfile TurnProfileConfig `json:"turn_profile,omitempty"` MaxLLMRetries int `json:"max_llm_retries,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_LLM_RETRIES"` LLMRetryBackoffSecs int `json:"llm_retry_backoff_secs,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_LLM_RETRY_BACKOFF_SECS"` } @@ -468,9 +469,25 @@ func (p *PlaceholderConfig) GetRandomText() string { } type StreamingConfig struct { - Enabled bool `json:"enabled,omitempty" env:"PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED"` - ThrottleSeconds int `json:"throttle_seconds,omitempty" env:"PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS"` - MinGrowthChars int `json:"min_growth_chars,omitempty" env:"PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS"` + Enabled bool `json:"enabled,omitempty"` + ThrottleSeconds int `json:"throttle_seconds,omitempty"` + MinGrowthChars int `json:"min_growth_chars,omitempty"` +} + +func (c StreamingConfig) IsZero() bool { + return !c.Enabled && c.ThrottleSeconds == 0 && c.MinGrowthChars == 0 +} + +func (c StreamingConfig) WithDefaults(throttleSeconds, minGrowthChars int) StreamingConfig { + if c.Enabled { + if c.ThrottleSeconds == 0 { + c.ThrottleSeconds = throttleSeconds + } + if c.MinGrowthChars == 0 { + c.MinGrowthChars = minGrowthChars + } + } + return c } type WhatsAppSettings struct { @@ -483,7 +500,7 @@ type TelegramSettings struct { Token SecureString `json:"token,omitzero" yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_TELEGRAM_TOKEN"` BaseURL string `json:"base_url" yaml:"-" env:"PICOCLAW_CHANNELS_TELEGRAM_BASE_URL"` Proxy string `json:"proxy" yaml:"-" env:"PICOCLAW_CHANNELS_TELEGRAM_PROXY"` - Streaming StreamingConfig `json:"streaming,omitempty" yaml:"-"` + Streaming StreamingConfig `json:"streaming,omitzero" yaml:"-"` UseMarkdownV2 bool `json:"use_markdown_v2" yaml:"-" env:"PICOCLAW_CHANNELS_TELEGRAM_USE_MARKDOWN_V2"` MediaGroupDelayMS int `json:"media_group_delay_ms" yaml:"-" env:"PICOCLAW_CHANNELS_TELEGRAM_MEDIA_GROUP_DELAY_MS"` } @@ -557,10 +574,11 @@ type WeComGroupConfig struct { } type WeComSettings struct { - BotID string `json:"bot_id" yaml:"-" env:"BOT_ID"` - Secret SecureString `json:"secret,omitzero" yaml:"secret,omitempty" env:"SECRET"` - WebSocketURL string `json:"websocket_url,omitempty" yaml:"-" env:"WEBSOCKET_URL"` - SendThinkingMessage bool `json:"send_thinking_message" yaml:"-" env:"SEND_THINKING_MESSAGE"` + BotID string `json:"bot_id" yaml:"-" env:"BOT_ID"` + Secret SecureString `json:"secret,omitzero" yaml:"secret,omitempty" env:"SECRET"` + WebSocketURL string `json:"websocket_url,omitempty" yaml:"-" env:"WEBSOCKET_URL"` + SendThinkingMessage bool `json:"send_thinking_message" yaml:"-" env:"SEND_THINKING_MESSAGE"` + Streaming StreamingConfig `json:"streaming,omitzero" yaml:"-"` } func (c *WeComSettings) SetSecret(secret string) { @@ -581,13 +599,14 @@ func (c *WeixinSettings) SetToken(token string) { } type PicoSettings struct { - Token SecureString `json:"token,omitzero" yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_PICO_TOKEN"` - AllowTokenQuery bool `json:"allow_token_query,omitempty" yaml:"-"` - AllowOrigins []string `json:"allow_origins,omitempty" yaml:"-"` - PingInterval int `json:"ping_interval,omitempty" yaml:"-"` - ReadTimeout int `json:"read_timeout,omitempty" yaml:"-"` - WriteTimeout int `json:"write_timeout,omitempty" yaml:"-"` - MaxConnections int `json:"max_connections,omitempty" yaml:"-"` + Token SecureString `json:"token,omitzero" yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_PICO_TOKEN"` + AllowTokenQuery bool `json:"allow_token_query,omitempty" yaml:"-"` + AllowOrigins []string `json:"allow_origins,omitempty" yaml:"-"` + Streaming StreamingConfig `json:"streaming,omitzero" yaml:"-"` + PingInterval int `json:"ping_interval,omitempty" yaml:"-"` + ReadTimeout int `json:"read_timeout,omitempty" yaml:"-"` + WriteTimeout int `json:"write_timeout,omitempty" yaml:"-"` + MaxConnections int `json:"max_connections,omitempty" yaml:"-"` } // SetToken sets the Pico token and marks it as dirty for security saving @@ -678,6 +697,14 @@ type VoiceConfig struct { ElevenLabsAPIKey string `json:"elevenlabs_api_key,omitempty" env:"PICOCLAW_VOICE_ELEVENLABS_API_KEY"` } +type ModelStreamingConfig struct { + Enabled bool `json:"enabled,omitempty"` +} + +func (c ModelStreamingConfig) IsZero() bool { + return !c.Enabled +} + // ModelConfig represents a model-centric provider configuration. // It allows adding new providers (especially OpenAI-compatible ones) via configuration only. // The Model field may be either a plain model identifier or a provider-prefixed @@ -702,13 +729,14 @@ type ModelConfig struct { Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers // Optional optimizations - RPM int `json:"rpm,omitempty"` // Requests per minute limit - MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") - RequestTimeout int `json:"request_timeout,omitempty"` - ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive - ToolSchemaTransform string `json:"tool_schema_transform,omitempty"` // Optional tool schema compatibility transform (e.g. "simple") - ExtraBody map[string]any `json:"extra_body,omitempty"` // Additional fields to inject into request body - CustomHeaders map[string]string `json:"custom_headers,omitempty"` // Additional headers to inject into every HTTP request + RPM int `json:"rpm,omitempty"` // Requests per minute limit + MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") + RequestTimeout int `json:"request_timeout,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive + ToolSchemaTransform string `json:"tool_schema_transform,omitempty"` // Optional tool schema compatibility transform (e.g. "simple") + Streaming ModelStreamingConfig `json:"streaming,omitzero"` // Opt-in for provider streaming on this model entry + ExtraBody map[string]any `json:"extra_body,omitempty"` // Additional fields to inject into request body + CustomHeaders map[string]string `json:"custom_headers,omitempty"` // Additional headers to inject into every HTTP request APIKeys SecureStrings `json:"api_keys,omitzero" yaml:"api_keys,omitempty"` // API authentication keys (multiple keys for failover) @@ -1239,7 +1267,7 @@ func LoadConfig(path string) (*Config, error) { return nil, err } - err = makeBackup(path) + err = MakeBackup(path) if err != nil { return nil, err } @@ -1293,7 +1321,7 @@ func LoadConfig(path string) (*Config, error) { return nil, err } - err = makeBackup(path) + err = MakeBackup(path) if err != nil { return nil, err } @@ -1345,7 +1373,7 @@ func LoadConfig(path string) (*Config, error) { return nil, err } - err = makeBackup(path) + err = MakeBackup(path) if err != nil { return nil, err } @@ -1391,6 +1419,9 @@ func LoadConfig(path string) (*Config, error) { if err = InitChannelList(cfg.Channels); err != nil { return nil, err } + if err = cfg.ValidateTurnProfile(); err != nil { + return nil, err + } cfg.Gateway.Host, err = resolveGatewayHostFromEnv(gatewayHostBeforeEnv) if err != nil { return nil, fmt.Errorf("invalid gateway host: %w", err) @@ -1496,7 +1527,7 @@ func applySkillsRegistryEnvCompat(cfg *Config) { cfg.Tools.Skills.Registries.Set("github", githubCfg) } -func makeBackup(path string) error { +func MakeBackup(path string) error { if _, err := os.Stat(path); os.IsNotExist(err) { return nil } @@ -1624,6 +1655,19 @@ func (c *Config) SecurityCopyFrom(path string) error { return loadSecurityConfig(c, securityPath(path)) } +// ResetToDefaults backs up the current config, creates a default config, +// preserves security credentials from the existing config, and saves it. +func ResetToDefaults(configPath string) error { + if err := MakeBackup(configPath); err != nil { + return fmt.Errorf("backup before reset: %w", err) + } + cfg := DefaultConfig() + if err := cfg.SecurityCopyFrom(configPath); err != nil { + logger.WarnF("could not preserve security config", map[string]any{"error": err}) + } + return SaveConfig(configPath, cfg) +} + func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { var expanded []*ModelConfig @@ -1661,6 +1705,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ToolSchemaTransform: m.ToolSchemaTransform, + Streaming: m.Streaming, ExtraBody: m.ExtraBody, CustomHeaders: m.CustomHeaders, UserAgent: m.UserAgent, @@ -1685,6 +1730,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ToolSchemaTransform: m.ToolSchemaTransform, + Streaming: m.Streaming, ExtraBody: m.ExtraBody, CustomHeaders: m.CustomHeaders, UserAgent: m.UserAgent, diff --git a/pkg/config/config_channel.go b/pkg/config/config_channel.go index 52d4b493..63a9b0d8 100644 --- a/pkg/config/config_channel.go +++ b/pkg/config/config_channel.go @@ -3,7 +3,9 @@ package config import ( "encoding/json" "fmt" + "os" "reflect" + "strconv" "strings" "github.com/caarlos0/env/v11" @@ -239,6 +241,7 @@ func (b Channel) MarshalJSON() ([]byte, error) { if err != nil { return nil, err } + raw = preserveExplicitDisabledStreaming(raw, b.Settings) settings = raw } else { settings = b.Settings @@ -252,6 +255,36 @@ func (b Channel) MarshalJSON() ([]byte, error) { return json.Marshal((*Alias)(&out)) } +func preserveExplicitDisabledStreaming(settings, original RawNode) RawNode { + if len(original) == 0 || len(settings) == 0 { + return settings + } + + var originalMap map[string]any + if err := json.Unmarshal(original, &originalMap); err != nil { + return settings + } + originalStreaming, ok := originalMap["streaming"].(map[string]any) + if !ok || originalStreaming["enabled"] != false { + return settings + } + + var settingsMap map[string]any + if err := json.Unmarshal(settings, &settingsMap); err != nil { + return settings + } + if _, exists := settingsMap["streaming"]; exists { + return settings + } + settingsMap["streaming"] = map[string]any{"enabled": false} + + data, err := json.Marshal(settingsMap) + if err != nil { + return settings + } + return data +} + // MarshalYAML implements yaml.ValueMarshaler for Channel. // Outputs only secure fields in the Settings YAML (for .security.yml). // If Decode was called, it serializes from the stored extend (reflecting any @@ -696,6 +729,10 @@ func InitChannelList(channels ChannelsConfig) error { if err := env.Parse(target); err != nil { // Non-fatal: some env vars may not apply } + applyTelegramStreamingEnvCompat(target) + if err := validateChannelStreamingConfig(name, target); err != nil { + return err + } } } @@ -706,3 +743,48 @@ func InitChannelList(channels ChannelsConfig) error { return nil } + +func applyTelegramStreamingEnvCompat(target any) { + settings, ok := target.(*TelegramSettings) + if !ok || settings == nil { + return + } + + if raw, ok := os.LookupEnv("PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED"); ok { + if value, err := strconv.ParseBool(raw); err == nil { + settings.Streaming.Enabled = value + } + } + if raw, ok := os.LookupEnv("PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS"); ok { + if value, err := strconv.Atoi(raw); err == nil { + settings.Streaming.ThrottleSeconds = value + } + } + if raw, ok := os.LookupEnv("PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS"); ok { + if value, err := strconv.Atoi(raw); err == nil { + settings.Streaming.MinGrowthChars = value + } + } +} + +func validateChannelStreamingConfig(channelName string, target any) error { + var streaming StreamingConfig + switch settings := target.(type) { + case *PicoSettings: + streaming = settings.Streaming + case *TelegramSettings: + streaming = settings.Streaming + case *WeComSettings: + streaming = settings.Streaming + default: + return nil + } + + if streaming.ThrottleSeconds < 0 { + return fmt.Errorf("channel %q streaming.throttle_seconds must be >= 0", channelName) + } + if streaming.MinGrowthChars < 0 { + return fmt.Errorf("channel %q streaming.min_growth_chars must be >= 0", channelName) + } + return nil +} diff --git a/pkg/config/config_channel_test.go b/pkg/config/config_channel_test.go index fd3cd824..b34048ff 100644 --- a/pkg/config/config_channel_test.go +++ b/pkg/config/config_channel_test.go @@ -2,6 +2,7 @@ package config import ( "encoding/json" + "reflect" "strings" "testing" @@ -133,6 +134,179 @@ func TestChannel_JSON_Unmarshal(t *testing.T) { assert.Equal(t, "", cfg.Token.String()) } +func TestStreamingConfig_IsChannelGeneric(t *testing.T) { + typ := reflect.TypeOf(StreamingConfig{}) + for i := 0; i < typ.NumField(); i++ { + field := typ.Field(i) + if got := field.Tag.Get("env"); got != "" { + t.Fatalf("StreamingConfig.%s env tag = %q, want no channel-specific env tag", field.Name, got) + } + } +} + +func TestPicoSettings_StreamingConfig(t *testing.T) { + raw := RawNode(`{ + "token": "test-token", + "streaming": { + "enabled": true, + "throttle_seconds": 2, + "min_growth_chars": 80 + } + }`) + ch := &Channel{ + Type: ChannelPico, + Enabled: true, + Settings: raw, + } + ch.SetName("pico") + var picoCfg PicoSettings + if err := ch.Decode(&picoCfg); err != nil { + t.Fatalf("Decode() error = %v", err) + } + assert.True(t, picoCfg.Streaming.Enabled) + assert.Equal(t, 2, picoCfg.Streaming.ThrottleSeconds) + assert.Equal(t, 80, picoCfg.Streaming.MinGrowthChars) +} + +func TestWeComSettings_StreamingConfig(t *testing.T) { + raw := RawNode(`{ + "bot_id": "bot-1", + "streaming": { + "enabled": true, + "throttle_seconds": 4, + "min_growth_chars": 160 + } + }`) + ch := &Channel{ + Type: ChannelWeCom, + Enabled: true, + Settings: raw, + } + ch.SetName("wecom") + var wecomCfg WeComSettings + if err := ch.Decode(&wecomCfg); err != nil { + t.Fatalf("Decode() error = %v", err) + } + assert.True(t, wecomCfg.Streaming.Enabled) + assert.Equal(t, 4, wecomCfg.Streaming.ThrottleSeconds) + assert.Equal(t, 160, wecomCfg.Streaming.MinGrowthChars) +} + +func TestPicoStreamingConfig_Defaults(t *testing.T) { + cfg := StreamingConfig{Enabled: true} + got := cfg.WithDefaults(1, 40) + assert.Equal(t, 1, got.ThrottleSeconds) + assert.Equal(t, 40, got.MinGrowthChars) + + cfg = StreamingConfig{Enabled: true, ThrottleSeconds: 5, MinGrowthChars: 200} + got = cfg.WithDefaults(1, 40) + assert.Equal(t, 5, got.ThrottleSeconds) + assert.Equal(t, 200, got.MinGrowthChars) + + cfg = StreamingConfig{} + got = cfg.WithDefaults(1, 40) + assert.Equal(t, 0, got.ThrottleSeconds) + assert.Equal(t, 0, got.MinGrowthChars) +} + +func TestInitChannelList_TelegramStreamingEnvCompatibility(t *testing.T) { + t.Setenv("PICOCLAW_CHANNELS_TELEGRAM_STREAMING_ENABLED", "true") + t.Setenv("PICOCLAW_CHANNELS_TELEGRAM_STREAMING_THROTTLE_SECONDS", "3") + t.Setenv("PICOCLAW_CHANNELS_TELEGRAM_STREAMING_MIN_GROWTH_CHARS", "120") + + channels := ChannelsConfig{ + "telegram": { + Type: ChannelTelegram, + Enabled: true, + Settings: RawNode(`{"token":"telegram-token"}`), + }, + "pico": { + Type: ChannelPico, + Enabled: true, + Settings: RawNode(`{"token":"pico-token"}`), + }, + } + if err := InitChannelList(channels); err != nil { + t.Fatalf("InitChannelList() error = %v", err) + } + + tgDecoded, err := channels["telegram"].GetDecoded() + if err != nil { + t.Fatalf("telegram GetDecoded() error = %v", err) + } + tgCfg := tgDecoded.(*TelegramSettings) + assert.True(t, tgCfg.Streaming.Enabled) + assert.Equal(t, 3, tgCfg.Streaming.ThrottleSeconds) + assert.Equal(t, 120, tgCfg.Streaming.MinGrowthChars) + + picoDecoded, err := channels["pico"].GetDecoded() + if err != nil { + t.Fatalf("pico GetDecoded() error = %v", err) + } + picoCfg := picoDecoded.(*PicoSettings) + assert.False(t, picoCfg.Streaming.Enabled) + assert.Equal(t, 0, picoCfg.Streaming.ThrottleSeconds) + assert.Equal(t, 0, picoCfg.Streaming.MinGrowthChars) +} + +func TestInitChannelList_RejectsNegativeStreamingDeliveryValues(t *testing.T) { + tests := []struct { + name string + channelType string + settings string + }{ + { + name: "pico throttle", + channelType: ChannelPico, + settings: `{"token":"pico-token","streaming":{"enabled":true,"throttle_seconds":-1}}`, + }, + { + name: "pico growth", + channelType: ChannelPico, + settings: `{"token":"pico-token","streaming":{"enabled":true,"min_growth_chars":-1}}`, + }, + { + name: "telegram throttle", + channelType: ChannelTelegram, + settings: `{"token":"telegram-token","streaming":{"enabled":true,"throttle_seconds":-1}}`, + }, + { + name: "telegram growth", + channelType: ChannelTelegram, + settings: `{"token":"telegram-token","streaming":{"enabled":true,"min_growth_chars":-1}}`, + }, + { + name: "wecom throttle", + channelType: ChannelWeCom, + settings: `{"bot_id":"bot-1","streaming":{"enabled":true,"throttle_seconds":-1}}`, + }, + { + name: "wecom growth", + channelType: ChannelWeCom, + settings: `{"bot_id":"bot-1","streaming":{"enabled":true,"min_growth_chars":-1}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + channels := ChannelsConfig{ + tt.channelType: { + Type: tt.channelType, + Enabled: true, + Settings: RawNode(tt.settings), + }, + } + err := InitChannelList(channels) + if err == nil { + t.Fatal("InitChannelList() error = nil, want validation error") + } + if !strings.Contains(err.Error(), "streaming.") { + t.Fatalf("InitChannelList() error = %v, want streaming field error", err) + } + }) + } +} + // ═══════════════════════════════════════════════════ // JSON marshal: secure fields masked as [NOT_HERE] // ═══════════════════════════════════════════════════ @@ -161,6 +335,22 @@ func TestChannel_JSON_Marshal_SecureMasked(t *testing.T) { assert.Contains(t, string(data), "proxy") } +func TestChannel_JSON_Marshal_OmitsUnconfiguredStreaming(t *testing.T) { + ch := Channel{ + Enabled: true, + Type: ChannelPico, + name: "pico", + Settings: mustParseRawNode(`{"ping_interval":30}`), + } + var cfg PicoSettings + require.NoError(t, ch.Decode(&cfg)) + + data, err := json.MarshalIndent(ch, "", " ") + require.NoError(t, err) + + assert.NotContains(t, string(data), `"streaming"`) +} + // ═══════════════════════════════════════════════════ // YAML unmarshal: security.yml — only secure data // ═══════════════════════════════════════════════════ diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 96417612..e34f2389 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -160,6 +160,130 @@ func TestAgentConfig_FullParse(t *testing.T) { } } +func TestTurnProfileConfig_ParseAndResolve(t *testing.T) { + jsonData := `{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": {"mode": "off"}, + "system_prompt": {"mode": "off"}, + "skills": {"mode": "off"}, + "tools": { + "mode": "custom", + "allow": ["web_search", "web_fetch"] + } + } + } + } + }` + + cfg := DefaultConfig() + if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if err := cfg.ValidateTurnProfile(); err != nil { + t.Fatalf("ValidateTurnProfile() error = %v", err) + } + + profile, ok, err := cfg.Agents.Defaults.ResolveTurnProfile() + if err != nil { + t.Fatalf("ResolveTurnProfile() error = %v", err) + } + if !ok { + t.Fatal("ResolveTurnProfile() ok = false, want true") + } + if profile.HistoryMode != TurnProfileModeOff || + profile.SystemPromptMode != TurnProfileModeOff || + profile.SkillsMode != TurnProfileModeOff || + profile.ToolsMode != TurnProfileModeCustom { + t.Fatalf("resolved clean_web modes = %+v", profile) + } + assert.Equal(t, []string{"web_search", "web_fetch"}, profile.AllowedTools) +} + +func TestTurnProfileConfig_DisabledOrMissingIsNoop(t *testing.T) { + cfg := DefaultConfig() + + profile, ok, err := cfg.Agents.Defaults.ResolveTurnProfile() + if err != nil { + t.Fatalf("ResolveTurnProfile(missing) error = %v", err) + } + if ok { + t.Fatal("ResolveTurnProfile(missing) ok = true, want false") + } + if profile.Enabled { + t.Fatalf("ResolveTurnProfile(missing) profile.Enabled = true, want false") + } + + cfg.Agents.Defaults.TurnProfile = TurnProfileConfig{ + Enabled: false, + History: TurnProfileBlock{ + Mode: TurnProfileModeOff, + }, + } + profile, ok, err = cfg.Agents.Defaults.ResolveTurnProfile() + if err != nil { + t.Fatalf("ResolveTurnProfile(disabled) error = %v", err) + } + if ok || profile.Enabled { + t.Fatalf("disabled profile = (%+v, %v), want no-op", profile, ok) + } + + cfg.Agents.Defaults.TurnProfile = TurnProfileConfig{ + Enabled: false, + History: TurnProfileBlock{ + Mode: TurnProfileModeCustom, + }, + Tools: TurnProfileBlock{ + Mode: TurnProfileMode("sometimes"), + }, + } + if err := cfg.ValidateTurnProfile(); err != nil { + t.Fatalf("ValidateTurnProfile(disabled unsupported modes) error = %v, want nil", err) + } +} + +func TestTurnProfileConfig_ValidationRejectsUnsupportedModes(t *testing.T) { + tests := []struct { + name string + raw string + want string + }{ + { + name: "history custom unsupported", + raw: `{"agents":{"defaults":{"turn_profile":{"enabled":true,"history":{"mode":"custom"}}}}}`, + want: "history.mode", + }, + { + name: "system prompt custom unsupported", + raw: `{"agents":{"defaults":{"turn_profile":{"enabled":true,"system_prompt":{"mode":"custom"}}}}}`, + want: "system_prompt.mode", + }, + { + name: "unknown mode", + raw: `{"agents":{"defaults":{"turn_profile":{"enabled":true,"tools":{"mode":"sometimes"}}}}}`, + want: "unsupported mode", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := DefaultConfig() + if err := json.Unmarshal([]byte(tt.raw), cfg); err != nil { + t.Fatalf("unmarshal: %v", err) + } + err := cfg.ValidateTurnProfile() + if err == nil { + t.Fatal("ValidateTurnProfile() error = nil, want error") + } + if !strings.Contains(err.Error(), tt.want) { + t.Fatalf("ValidateTurnProfile() error = %v, want containing %q", err, tt.want) + } + }) + } +} + func TestDefaultConfig_MCPMaxInlineTextChars(t *testing.T) { cfg := DefaultConfig() if cfg.Tools.MCP.GetMaxInlineTextChars() != DefaultMCPMaxInlineTextChars { @@ -836,6 +960,42 @@ func TestDefaultConfig_Channels(t *testing.T) { } } +func TestDefaultConfig_ChannelStreamingDisabled(t *testing.T) { + cfg := DefaultConfig() + + telegram := cfg.Channels.Get(ChannelTelegram) + if telegram == nil { + t.Fatal("DefaultConfig() missing telegram channel") + } + decoded, err := telegram.GetDecoded() + if err != nil { + t.Fatalf("telegram GetDecoded() error = %v", err) + } + settings, ok := decoded.(*TelegramSettings) + if !ok { + t.Fatalf("telegram settings type = %T, want *TelegramSettings", decoded) + } + if settings.Streaming.Enabled { + t.Fatal("DefaultConfig().telegram.settings.streaming.enabled should be false") + } + + pico := cfg.Channels.Get(ChannelPico) + if pico == nil { + t.Fatal("DefaultConfig() missing pico channel") + } + decoded, err = pico.GetDecoded() + if err != nil { + t.Fatalf("pico GetDecoded() error = %v", err) + } + picoSettings, ok := decoded.(*PicoSettings) + if !ok { + t.Fatalf("pico settings type = %T, want *PicoSettings", decoded) + } + if !picoSettings.Streaming.Enabled { + t.Fatal("DefaultConfig().pico.settings.streaming.enabled should be true") + } +} + func TestValidateSingletonChannels_RejectsMultipleInstances(t *testing.T) { channels := ChannelsConfig{ "pico1": &Channel{Enabled: true, Type: ChannelPico}, @@ -975,6 +1135,49 @@ func TestSaveConfig_PreservesDisabledTelegramPlaceholder(t *testing.T) { } } +func TestSaveConfig_PreservesExplicitDisabledPicoStreaming(t *testing.T) { + tmpDir := t.TempDir() + path := filepath.Join(tmpDir, "config.json") + + cfg := DefaultConfig() + pico := cfg.Channels.Get(ChannelPico) + if pico == nil { + t.Fatal("DefaultConfig() missing pico channel") + } + pico.Settings = RawNode(`{"streaming":{"enabled":false}}`) + + if err := SaveConfig(path, cfg); err != nil { + t.Fatalf("SaveConfig failed: %v", err) + } + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if !strings.Contains(string(data), `"streaming"`) || !strings.Contains(string(data), `"enabled": false`) { + t.Fatalf("saved config should preserve explicit disabled pico streaming, got:\n%s", string(data)) + } + + loaded, err := LoadConfig(path) + if err != nil { + t.Fatalf("LoadConfig failed: %v", err) + } + loadedPico := loaded.Channels.Get(ChannelPico) + if loadedPico == nil { + t.Fatal("loaded config missing pico channel") + } + decoded, err := loadedPico.GetDecoded() + if err != nil { + t.Fatalf("pico GetDecoded() error = %v", err) + } + settings, ok := decoded.(*PicoSettings) + if !ok { + t.Fatalf("pico settings type = %T, want *PicoSettings", decoded) + } + if settings.Streaming.Enabled { + t.Fatal("explicit disabled pico streaming should remain disabled after SaveConfig/LoadConfig round-trip") + } +} + // TestSaveConfig_FiltersVirtualModels verifies that SaveConfig does not write // virtual models (generated by expandMultiKeyModels) to the config file. func TestSaveConfig_FiltersVirtualModels(t *testing.T) { @@ -2602,7 +2805,7 @@ func TestFilterSensitiveData_AllTokenTypes(t *testing.T) { } // --------------------------------------------------------------------------- -// makeBackup tests +// MakeBackup tests // --------------------------------------------------------------------------- // TestMakeBackup_WithDateSuffix verifies backup files include a date suffix. @@ -2613,8 +2816,8 @@ func TestMakeBackup_WithDateSuffix(t *testing.T) { t.Fatalf("WriteFile: %v", err) } - if err := makeBackup(configPath); err != nil { - t.Fatalf("makeBackup: %v", err) + if err := MakeBackup(configPath); err != nil { + t.Fatalf("MakeBackup: %v", err) } entries, err := os.ReadDir(dir) @@ -2653,8 +2856,8 @@ func TestMakeBackup_AlsoBacksSecurityFile(t *testing.T) { os.WriteFile(configPath, []byte(`{"version":2}`), 0o600) os.WriteFile(secPath, []byte(`model_list:\n test:0:\n api_keys:\n - "sk-test"\n`), 0o600) - if err := makeBackup(configPath); err != nil { - t.Fatalf("makeBackup: %v", err) + if err := MakeBackup(configPath); err != nil { + t.Fatalf("MakeBackup: %v", err) } entries, err := os.ReadDir(dir) @@ -2680,14 +2883,14 @@ func TestMakeBackup_AlsoBacksSecurityFile(t *testing.T) { } } -// TestMakeBackup_NonexistentFileSkipsBackup verifies that makeBackup returns nil +// TestMakeBackup_NonexistentFileSkipsBackup verifies that MakeBackup returns nil // when the config file does not exist (no error, no panic). func TestMakeBackup_NonexistentFileSkipsBackup(t *testing.T) { dir := t.TempDir() configPath := filepath.Join(dir, "nonexistent.json") - if err := makeBackup(configPath); err != nil { - t.Fatalf("makeBackup on nonexistent file should return nil, got: %v", err) + if err := MakeBackup(configPath); err != nil { + t.Fatalf("MakeBackup on nonexistent file should return nil, got: %v", err) } } @@ -2698,8 +2901,8 @@ func TestMakeBackup_OnlyConfigNoSecurity(t *testing.T) { configPath := filepath.Join(dir, "config.json") os.WriteFile(configPath, []byte(`{"version":2}`), 0o600) - if err := makeBackup(configPath); err != nil { - t.Fatalf("makeBackup: %v", err) + if err := MakeBackup(configPath); err != nil { + t.Fatalf("MakeBackup: %v", err) } entries, _ := os.ReadDir(dir) @@ -2722,7 +2925,7 @@ func TestMakeBackup_OnlyConfigNoSecurity(t *testing.T) { } // TestMakeBackup_SameDateSuffix verifies that config and security backups -// share the same date suffix (they are created in the same makeBackup call). +// share the same date suffix (they are created in the same MakeBackup call). func TestMakeBackup_SameDateSuffix(t *testing.T) { dir := t.TempDir() configPath := filepath.Join(dir, "config.json") @@ -2731,8 +2934,8 @@ func TestMakeBackup_SameDateSuffix(t *testing.T) { os.WriteFile(configPath, []byte(`{"version":2}`), 0o600) os.WriteFile(secPath, []byte(`key: value`), 0o600) - if err := makeBackup(configPath); err != nil { - t.Fatalf("makeBackup: %v", err) + if err := MakeBackup(configPath); err != nil { + t.Fatalf("MakeBackup: %v", err) } entries, _ := os.ReadDir(dir) diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 995817f1..c411aadf 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -511,7 +511,6 @@ func defaultChannels() ChannelsConfig { "typing": map[string]any{"enabled": true}, "placeholder": map[string]any{"enabled": true, "text": []string{"Thinking... 💭"}}, "settings": map[string]any{ - "streaming": map[string]any{"enabled": true, "throttle_seconds": 3, "min_growth_chars": 200}, "use_markdown_v2": false, "media_group_delay_ms": 500, }, @@ -566,6 +565,7 @@ func defaultChannels() ChannelsConfig { "read_timeout": 60, "write_timeout": 10, "max_connections": 100, + "streaming": map[string]any{"enabled": true}, }, }, "irc": map[string]any{ diff --git a/pkg/config/model_config_test.go b/pkg/config/model_config_test.go index d22eb290..210e963d 100644 --- a/pkg/config/model_config_test.go +++ b/pkg/config/model_config_test.go @@ -7,6 +7,7 @@ package config import ( "encoding/json" + "reflect" "strings" "sync" "testing" @@ -144,6 +145,47 @@ func TestGetModelConfig_Concurrent(t *testing.T) { } } +func TestModelConfig_StreamingConfig(t *testing.T) { + t.Run("loads streaming enabled", func(t *testing.T) { + var cfg ModelConfig + err := json.Unmarshal([]byte(`{ + "model_name": "stream-model", + "model": "openai/gpt-5.4", + "streaming": {"enabled": true} + }`), &cfg) + if err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if !cfg.Streaming.Enabled { + t.Fatal("Streaming.Enabled = false, want true") + } + }) + + t.Run("defaults disabled", func(t *testing.T) { + var cfg ModelConfig + err := json.Unmarshal([]byte(`{ + "model_name": "plain-model", + "model": "openai/gpt-5.4" + }`), &cfg) + if err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if cfg.Streaming.Enabled { + t.Fatal("Streaming.Enabled = true, want false by default") + } + }) + + t.Run("model streaming only has enabled", func(t *testing.T) { + typ := reflect.TypeOf(ModelStreamingConfig{}) + if typ.NumField() != 1 { + t.Fatalf("ModelStreamingConfig field count = %d, want 1", typ.NumField()) + } + if _, ok := typ.FieldByName("Enabled"); !ok { + t.Fatal("ModelStreamingConfig missing Enabled field") + } + }) +} + func TestModelConfig_Validate(t *testing.T) { tests := []struct { name string diff --git a/pkg/config/multikey_test.go b/pkg/config/multikey_test.go index 073cb782..c267eb52 100644 --- a/pkg/config/multikey_test.go +++ b/pkg/config/multikey_test.go @@ -197,6 +197,7 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) { RequestTimeout: 30, ThinkingLevel: "high", ToolSchemaTransform: "simple", + Streaming: ModelStreamingConfig{Enabled: true}, } modelCfg.APIKeys = SimpleSecureStrings("key0", "key1") // Use internal field for multi-key testing models := []*ModelConfig{modelCfg} @@ -229,6 +230,9 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) { if primary.ToolSchemaTransform != "simple" { t.Errorf("expected tool_schema_transform preserved, got %q", primary.ToolSchemaTransform) } + if !primary.Streaming.Enabled { + t.Error("expected streaming config preserved on primary") + } // Check additional entry also preserves fields additional := result[0] @@ -244,6 +248,9 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) { if additional.ToolSchemaTransform != "simple" { t.Errorf("expected additional tool_schema_transform preserved, got %q", additional.ToolSchemaTransform) } + if !additional.Streaming.Enabled { + t.Error("expected streaming config preserved on additional") + } } func TestExpandMultiKeyModels_IsVirtualFlag(t *testing.T) { diff --git a/pkg/config/turn_profile.go b/pkg/config/turn_profile.go new file mode 100644 index 00000000..967a697e --- /dev/null +++ b/pkg/config/turn_profile.go @@ -0,0 +1,133 @@ +package config + +import ( + "fmt" + "strings" +) + +type TurnProfileMode string + +const ( + TurnProfileModeDefault TurnProfileMode = "default" + TurnProfileModeOff TurnProfileMode = "off" + TurnProfileModeCustom TurnProfileMode = "custom" +) + +type TurnProfileConfig struct { + Enabled bool `json:"enabled"` + History TurnProfileBlock `json:"history,omitempty"` + SystemPrompt TurnProfileBlock `json:"system_prompt,omitempty"` + Skills TurnProfileBlock `json:"skills,omitempty"` + Tools TurnProfileBlock `json:"tools,omitempty"` +} + +type TurnProfileBlock struct { + Mode TurnProfileMode `json:"mode,omitempty"` + Allow []string `json:"allow,omitempty"` +} + +type EffectiveTurnProfile struct { + Enabled bool + HistoryMode TurnProfileMode + SystemPromptMode TurnProfileMode + SkillsMode TurnProfileMode + ToolsMode TurnProfileMode + AllowedSkills []string + AllowedTools []string +} + +func (m TurnProfileMode) Effective() TurnProfileMode { + switch TurnProfileMode(strings.ToLower(strings.TrimSpace(string(m)))) { + case "", TurnProfileModeDefault: + return TurnProfileModeDefault + case TurnProfileModeOff: + return TurnProfileModeOff + case TurnProfileModeCustom: + return TurnProfileModeCustom + default: + return TurnProfileMode(strings.ToLower(strings.TrimSpace(string(m)))) + } +} + +func (d *AgentDefaults) ResolveTurnProfile() (EffectiveTurnProfile, bool, error) { + if d == nil { + return EffectiveTurnProfile{}, false, nil + } + profile := d.TurnProfile + if !profile.Enabled { + return EffectiveTurnProfile{}, false, nil + } + if err := validateTurnProfile(profile); err != nil { + return EffectiveTurnProfile{}, false, err + } + return EffectiveTurnProfile{ + Enabled: true, + HistoryMode: profile.History.Mode.Effective(), + SystemPromptMode: profile.SystemPrompt.Mode.Effective(), + SkillsMode: profile.Skills.Mode.Effective(), + ToolsMode: profile.Tools.Mode.Effective(), + AllowedSkills: cleanStringList(profile.Skills.Allow), + AllowedTools: cleanStringList(profile.Tools.Allow), + }, true, nil +} + +func (c *Config) ValidateTurnProfile() error { + if c == nil { + return nil + } + return validateTurnProfile(c.Agents.Defaults.TurnProfile) +} + +func validateTurnProfile(profile TurnProfileConfig) error { + if !profile.Enabled { + return nil + } + if err := validateTurnProfileBlock("history", profile.History, false); err != nil { + return err + } + if err := validateTurnProfileBlock("system_prompt", profile.SystemPrompt, false); err != nil { + return err + } + if err := validateTurnProfileBlock("skills", profile.Skills, true); err != nil { + return err + } + if err := validateTurnProfileBlock("tools", profile.Tools, true); err != nil { + return err + } + return nil +} + +func validateTurnProfileBlock(field string, block TurnProfileBlock, allowCustom bool) error { + mode := block.Mode.Effective() + switch mode { + case TurnProfileModeDefault, TurnProfileModeOff: + return nil + case TurnProfileModeCustom: + if allowCustom { + return nil + } + return fmt.Errorf("turn_profile.%s.mode custom is not supported in this version", field) + default: + return fmt.Errorf("turn_profile.%s.mode has unsupported mode %q", field, block.Mode) + } +} + +func cleanStringList(values []string) []string { + if len(values) == 0 { + return nil + } + out := make([]string, 0, len(values)) + seen := make(map[string]struct{}, len(values)) + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" { + continue + } + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + out = append(out, value) + } + return out +} diff --git a/pkg/mcp/manager_real_server_integration_test.go b/pkg/mcp/manager_real_server_integration_test.go index 742a4fb2..3c002244 100644 --- a/pkg/mcp/manager_real_server_integration_test.go +++ b/pkg/mcp/manager_real_server_integration_test.go @@ -69,7 +69,11 @@ func TestIntegration_RealConfiguredServer(t *testing.T) { t.Fatal("expected at least one discovered tool from real MCP server") } - t.Logf("connected to real MCP server via %s with %d tool(s)", config.EffectiveMCPTransportType(serverCfg), len(tools)) + t.Logf( + "connected to real MCP server via %s with %d tool(s)", + config.EffectiveMCPTransportType(serverCfg), + len(tools), + ) for _, tool := range tools { if tool != nil { t.Logf("discovered tool: %s", tool.Name) diff --git a/pkg/memory/jsonl.go b/pkg/memory/jsonl.go index 49220511..a63598f4 100644 --- a/pkg/memory/jsonl.go +++ b/pkg/memory/jsonl.go @@ -561,6 +561,12 @@ func (s *JSONLStore) addMsg(sessionKey string, msg providers.Message) error { l.Lock() defer l.Unlock() + now := time.Now() + + if msg.CreatedAt == nil { + msg.CreatedAt = &now + } + // Append the message as a single JSON line. line, err := json.Marshal(msg) if err != nil { @@ -598,7 +604,6 @@ func (s *JSONLStore) addMsg(sessionKey string, msg providers.Message) error { if err != nil { return err } - now := time.Now() if meta.Count == 0 && meta.CreatedAt.IsZero() { meta.CreatedAt = now } @@ -726,6 +731,12 @@ func (s *JSONLStore) SetHistory( meta.Count = len(history) meta.UpdatedAt = now + for i := range history { + if history[i].CreatedAt == nil { + history[i].CreatedAt = &now + } + } + // Write meta BEFORE rewriting the JSONL file. If we crash between // the two writes, meta has Skip=0 and the old file is still intact, // so GetHistory reads from line 1 — returning "too many" messages diff --git a/pkg/memory/jsonl_test.go b/pkg/memory/jsonl_test.go index 3a7b9813..0b15f160 100644 --- a/pkg/memory/jsonl_test.go +++ b/pkg/memory/jsonl_test.go @@ -130,6 +130,32 @@ func TestAddFullMessage_WithToolCalls(t *testing.T) { } } +func TestAddFullMessage_PreservesModelName(t *testing.T) { + store := newTestStore(t) + ctx := context.Background() + + msg := providers.Message{ + Role: "assistant", + Content: "done", + ModelName: "gpt-5.4-mini", + } + + if err := store.AddFullMessage(ctx, "model-name", msg); err != nil { + t.Fatalf("AddFullMessage: %v", err) + } + + history, err := store.GetHistory(ctx, "model-name") + if err != nil { + t.Fatalf("GetHistory: %v", err) + } + if len(history) != 1 { + t.Fatalf("expected 1, got %d", len(history)) + } + if history[0].ModelName != "gpt-5.4-mini" { + t.Fatalf("ModelName = %q, want %q", history[0].ModelName, "gpt-5.4-mini") + } +} + func TestAddFullMessage_ToolCallID(t *testing.T) { store := newTestStore(t) ctx := context.Background() @@ -1032,6 +1058,143 @@ func TestMultipleSessions_Isolation(t *testing.T) { } } +func TestStore_SetsCreatedAtWhenNil(t *testing.T) { + type writeOp struct { + name string + fn func(store *JSONLStore, key string) (expectedCount int) + } + + ops := []writeOp{ + { + name: "AddMessage", + fn: func(store *JSONLStore, key string) int { + if err := store.AddMessage(context.Background(), key, "user", "hello"); err != nil { + t.Fatalf("AddMessage: %v", err) + } + return 1 + }, + }, + { + name: "AddFullMessage", + fn: func(store *JSONLStore, key string) int { + if err := store.AddFullMessage(context.Background(), key, providers.Message{ + Role: "user", + Content: "hello from full", + }); err != nil { + t.Fatalf("AddFullMessage: %v", err) + } + return 1 + }, + }, + { + name: "SetHistory", + fn: func(store *JSONLStore, key string) int { + if err := store.SetHistory(context.Background(), key, []providers.Message{ + {Role: "user", Content: "msg1"}, + {Role: "assistant", Content: "msg2"}, + }); err != nil { + t.Fatalf("SetHistory: %v", err) + } + return 2 + }, + }, + } + + for _, op := range ops { + t.Run(op.name, func(t *testing.T) { + store := newTestStore(t) + key := "s1" + + before := time.Now().Add(-time.Second) + expectedCount := op.fn(store, key) + after := time.Now().Add(time.Second) + + history, err := store.GetHistory(context.Background(), key) + if err != nil { + t.Fatalf("GetHistory: %v", err) + } + if len(history) != expectedCount { + t.Fatalf("expected %d messages, got %d", expectedCount, len(history)) + } + for i := range history { + if history[i].CreatedAt == nil || history[i].CreatedAt.IsZero() { + t.Errorf("message %d CreatedAt is zero — not set by %s", i, op.name) + } + if history[i].CreatedAt.Before(before) || history[i].CreatedAt.After(after) { + t.Errorf( + "message %d CreatedAt %v outside expected window [%v, %v]", + i, history[i].CreatedAt, before, after, + ) + } + } + }) + } +} + +func TestStore_PreservesExistingCreatedAt(t *testing.T) { + t1 := time.Date(2026, 1, 1, 10, 0, 0, 0, time.UTC) + t2 := time.Date(2026, 1, 1, 11, 0, 0, 0, time.UTC) + + type writeOp struct { + name string + fn func(store *JSONLStore, key string) + wantTimes []time.Time + } + + ops := []writeOp{ + { + name: "AddFullMessage", + fn: func(store *JSONLStore, key string) { + if err := store.AddFullMessage(context.Background(), key, providers.Message{ + Role: "user", + Content: "custom time", + CreatedAt: &t1, + }); err != nil { + t.Fatalf("AddFullMessage: %v", err) + } + }, + wantTimes: []time.Time{t1}, + }, + { + name: "SetHistory", + fn: func(store *JSONLStore, key string) { + if err := store.SetHistory(context.Background(), key, []providers.Message{ + {Role: "user", Content: "msg1", CreatedAt: &t1}, + {Role: "assistant", Content: "msg2", CreatedAt: &t2}, + }); err != nil { + t.Fatalf("SetHistory: %v", err) + } + }, + wantTimes: []time.Time{t1, t2}, + }, + } + + for _, op := range ops { + t.Run(op.name, func(t *testing.T) { + store := newTestStore(t) + key := "s1" + + op.fn(store, key) + + history, err := store.GetHistory(context.Background(), key) + if err != nil { + t.Fatalf("GetHistory: %v", err) + } + if len(history) != len(op.wantTimes) { + t.Fatalf("expected %d messages, got %d", len(op.wantTimes), len(history)) + } + for i, want := range op.wantTimes { + if history[i].CreatedAt == nil || !history[i].CreatedAt.Equal(want) { + t.Errorf( + "message %d CreatedAt = %v, want %v (should preserve caller-provided time)", + i, history[i].CreatedAt, want, + ) + } + } + }) + } +} + func BenchmarkAddMessage(b *testing.B) { dir := b.TempDir() store, err := NewJSONLStore(dir) diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go index 034c8050..7c03daab 100644 --- a/pkg/providers/factory_provider.go +++ b/pkg/providers/factory_provider.go @@ -18,54 +18,6 @@ import ( "github.com/sipeed/picoclaw/pkg/providers/common" ) -type protocolMeta struct { - defaultAPIBase string - emptyAPIKeyAllowed bool -} - -var protocolMetaByName = map[string]protocolMeta{ - "openai": {defaultAPIBase: "https://api.openai.com/v1"}, - "venice": {defaultAPIBase: "https://api.venice.ai/api/v1"}, - "openrouter": {defaultAPIBase: "https://openrouter.ai/api/v1"}, - "litellm": {defaultAPIBase: "http://localhost:4000/v1"}, - "lmstudio": {defaultAPIBase: "http://localhost:1234/v1", emptyAPIKeyAllowed: true}, - "novita": {defaultAPIBase: "https://api.novita.ai/openai"}, - "groq": {defaultAPIBase: "https://api.groq.com/openai/v1"}, - "zhipu": {defaultAPIBase: "https://open.bigmodel.cn/api/paas/v4"}, - "gemini": {defaultAPIBase: "https://generativelanguage.googleapis.com/v1beta"}, - "nvidia": {defaultAPIBase: "https://integrate.api.nvidia.com/v1"}, - "ollama": {defaultAPIBase: "http://localhost:11434/v1", emptyAPIKeyAllowed: true}, - "moonshot": {defaultAPIBase: "https://api.moonshot.cn/v1"}, - "shengsuanyun": {defaultAPIBase: "https://router.shengsuanyun.com/api/v1"}, - "siliconflow": {defaultAPIBase: "https://api.siliconflow.cn/v1"}, - "deepseek": {defaultAPIBase: "https://api.deepseek.com/v1"}, - "cerebras": {defaultAPIBase: "https://api.cerebras.ai/v1"}, - "vivgrid": {defaultAPIBase: "https://api.vivgrid.com/v1"}, - "volcengine": {defaultAPIBase: "https://ark.cn-beijing.volces.com/api/v3"}, - "qwen": {defaultAPIBase: "https://dashscope.aliyuncs.com/compatible-mode/v1"}, - "qwen-portal": {defaultAPIBase: "https://dashscope.aliyuncs.com/compatible-mode/v1"}, - "qwen-intl": {defaultAPIBase: "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"}, - "qwen-international": {defaultAPIBase: "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"}, - "dashscope-intl": {defaultAPIBase: "https://dashscope-intl.aliyuncs.com/compatible-mode/v1"}, - "qwen-us": {defaultAPIBase: "https://dashscope-us.aliyuncs.com/compatible-mode/v1"}, - "dashscope-us": {defaultAPIBase: "https://dashscope-us.aliyuncs.com/compatible-mode/v1"}, - "coding-plan": {defaultAPIBase: "https://coding-intl.dashscope.aliyuncs.com/v1"}, - "alibaba-coding": {defaultAPIBase: "https://coding-intl.dashscope.aliyuncs.com/v1"}, - "qwen-coding": {defaultAPIBase: "https://coding-intl.dashscope.aliyuncs.com/v1"}, - "coding-plan-anthropic": {defaultAPIBase: "https://coding-intl.dashscope.aliyuncs.com/apps/anthropic"}, - "alibaba-coding-anthropic": {defaultAPIBase: "https://coding-intl.dashscope.aliyuncs.com/apps/anthropic"}, - "zai": {defaultAPIBase: "https://api.z.ai/api/coding/paas/v4"}, - "vllm": {defaultAPIBase: "http://localhost:8000/v1", emptyAPIKeyAllowed: true}, - "mistral": {defaultAPIBase: "https://api.mistral.ai/v1"}, - "avian": {defaultAPIBase: "https://api.avian.io/v1"}, - "minimax": {defaultAPIBase: "https://api.minimaxi.com/v1"}, - "longcat": {defaultAPIBase: "https://api.longcat.chat/openai"}, - "modelscope": {defaultAPIBase: "https://api-inference.modelscope.cn/v1"}, - "mimo": {defaultAPIBase: "https://api.xiaomimimo.com/v1"}, - "anthropic": {defaultAPIBase: "https://api.anthropic.com/v1"}, - "anthropic-messages": {defaultAPIBase: "https://api.anthropic.com/v1"}, -} - // createClaudeAuthProvider creates a Claude provider using OAuth credentials from auth store. func createClaudeAuthProvider() (LLMProvider, error) { cred, err := getCredential("anthropic") @@ -184,7 +136,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err provider.SetProviderName(protocol) return finalizeProviderFromConfig(provider, modelID, cfg) - case "azure", "azure-openai": + case "azure": // Azure OpenAI uses deployment-based URLs, api-key header auth, // and always sends max_completion_tokens. if cfg.APIKey() == "" { @@ -239,11 +191,10 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err } return finalizeProviderFromConfig(provider, modelID, cfg) - case "litellm", "lmstudio", "openrouter", "groq", "zhipu", "nvidia", "venice", + case "litellm", "lmstudio", "gpt4free", "openrouter", "groq", "zhipu", "nvidia", "venice", "ollama", "moonshot", "shengsuanyun", "siliconflow", "deepseek", "cerebras", - "vivgrid", "volcengine", "vllm", "qwen", "qwen-portal", "qwen-intl", "qwen-international", "dashscope-intl", - "qwen-us", "dashscope-us", "mistral", "avian", "longcat", "modelscope", "novita", - "coding-plan", "alibaba-coding", "qwen-coding", "zai", "mimo": + "vivgrid", "volcengine", "vllm", "qwen-portal", "qwen-intl", "qwen-us", "mistral", + "avian", "longcat", "modelscope", "novita", "alibaba-coding", "zai", "mimo": // All other OpenAI-compatible HTTP providers if cfg.APIKey() == "" && cfg.APIBase == "" && !isEmptyAPIKeyAllowed(protocol) { return nil, "", fmt.Errorf("api_key or api_base is required for HTTP-based protocol %q", protocol) @@ -355,7 +306,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err cfg.RequestTimeout, ), modelID, cfg) - case "coding-plan-anthropic", "alibaba-coding-anthropic": + case "alibaba-coding-anthropic": // Alibaba Coding Plan with Anthropic-compatible API apiBase := cfg.APIBase if apiBase == "" { @@ -374,21 +325,21 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err case "antigravity": return finalizeProviderFromConfig(NewAntigravityProvider(), modelID, cfg) - case "claude-cli", "claudecli": + case "claude-cli": workspace := cfg.Workspace if workspace == "" { workspace = "." } return finalizeProviderFromConfig(NewClaudeCliProvider(workspace), modelID, cfg) - case "codex-cli", "codexcli": + case "codex-cli": workspace := cfg.Workspace if workspace == "" { workspace = "." } return finalizeProviderFromConfig(NewCodexCliProvider(workspace), modelID, cfg) - case "github-copilot", "copilot": + case "github-copilot": apiBase := cfg.APIBase if apiBase == "" { apiBase = "localhost:4321" @@ -421,8 +372,8 @@ func finalizeProviderFromConfig( } func isEmptyAPIKeyAllowed(protocol string) bool { - meta, ok := protocolMetaForName(protocol) - return ok && meta.emptyAPIKeyAllowed + option, ok := modelProviderOptionForName(protocol) + return ok && option.EmptyAPIKeyAllowed } // IsEmptyAPIKeyAllowedForProtocol reports whether a protocol allows requests @@ -432,6 +383,16 @@ func IsEmptyAPIKeyAllowedForProtocol(protocol string) bool { return isEmptyAPIKeyAllowed(protocol) } +// IsHTTPAPIProtocol reports whether a provider uses an HTTP API base in the +// model configuration path. This excludes providers such as Bedrock, CLI +// bridges, and OAuth-only managed providers even if they do not require an +// explicit api_key field. +func IsHTTPAPIProtocol(protocol string) bool { + protocol = NormalizeProvider(protocol) + option, ok := modelProviderOptionsByName[protocol] + return ok && option.httpAPI +} + // DefaultAPIBaseForProtocol returns the configured default API base for a protocol. // It returns empty string if the protocol has no default base. func DefaultAPIBaseForProtocol(protocol string) string { @@ -441,19 +402,9 @@ func DefaultAPIBaseForProtocol(protocol string) string { // getDefaultAPIBase returns the default API base URL for a given protocol. func getDefaultAPIBase(protocol string) string { - meta, ok := protocolMetaForName(protocol) + option, ok := modelProviderOptionForName(protocol) if !ok { return "" } - return meta.defaultAPIBase -} - -func protocolMetaForName(protocol string) (protocolMeta, bool) { - if meta, ok := protocolMetaByName[protocol]; ok { - return meta, true - } - if meta, ok := attachedModelProviderMetaByName[protocol]; ok { - return meta.protocolMeta, true - } - return protocolMeta{}, false + return option.DefaultAPIBase } diff --git a/pkg/providers/factory_provider_test.go b/pkg/providers/factory_provider_test.go index 69903711..0b3dd791 100644 --- a/pkg/providers/factory_provider_test.go +++ b/pkg/providers/factory_provider_test.go @@ -171,6 +171,30 @@ func TestCreateProviderFromConfig_UsesExplicitProvider(t *testing.T) { } } +func TestCreateProviderFromConfig_DeepSeekSupportsThinking(t *testing.T) { + cfg := &config.ModelConfig{ + ModelName: "deepseek-v4-flash", + Provider: "deepseek", + Model: "deepseek-v4-flash", + } + cfg.SetAPIKey("test-key") + + provider, modelID, err := CreateProviderFromConfig(cfg) + if err != nil { + t.Fatalf("CreateProviderFromConfig() error = %v", err) + } + if modelID != "deepseek-v4-flash" { + t.Fatalf("modelID = %q, want %q", modelID, "deepseek-v4-flash") + } + tc, ok := provider.(ThinkingCapable) + if !ok { + t.Fatalf("provider %T should implement ThinkingCapable for DeepSeek", provider) + } + if !tc.SupportsThinking() { + t.Fatalf("DeepSeek provider SupportsThinking() = false, want true") + } +} + func TestCreateProviderFromConfig_PreservesExplicitProviderPrefixedModel(t *testing.T) { cfg := &config.ModelConfig{ ModelName: "test-openai", @@ -210,6 +234,7 @@ func TestCreateProviderFromConfig_DefaultAPIBase(t *testing.T) { {"deepseek", "deepseek"}, {"ollama", "ollama"}, {"lmstudio", "lmstudio"}, + {"gpt4free", "gpt4free"}, {"longcat", "longcat"}, {"modelscope", "modelscope"}, {"mimo", "mimo"}, @@ -248,6 +273,15 @@ func TestGetDefaultAPIBase_LMStudio(t *testing.T) { } } +func TestGetDefaultAPIBase_GPT4Free(t *testing.T) { + if got := getDefaultAPIBase("gpt4free"); got != "http://localhost:1337/v1" { + t.Fatalf("getDefaultAPIBase(%q) = %q, want %q", "gpt4free", got, "http://localhost:1337/v1") + } + if got := getDefaultAPIBase("g4f"); got != "http://localhost:1337/v1" { + t.Fatalf("getDefaultAPIBase(%q) = %q, want %q", "g4f", got, "http://localhost:1337/v1") + } +} + func TestGetDefaultAPIBase_Venice(t *testing.T) { if got := getDefaultAPIBase("venice"); got != "https://api.venice.ai/api/v1" { t.Fatalf("getDefaultAPIBase(%q) = %q, want %q", "venice", got, "https://api.venice.ai/api/v1") @@ -330,6 +364,13 @@ func TestCreateProviderFromConfig_LocalProviders(t *testing.T) { apiKey: "", wantModelID: "Qwen/Qwen3-8B", }, + { + name: "GPT4Free without API key", + modelName: "test-gpt4free", + model: "gpt4free/gpt-4o-mini", + apiKey: "", + wantModelID: "gpt-4o-mini", + }, } for _, tt := range tests { @@ -946,7 +987,7 @@ func TestCreateProviderFromConfig_CodingPlanAnthropic(t *testing.T) { if modelID != wantModelID { t.Errorf("modelID = %q, want %q", modelID, wantModelID) } - // coding-plan-anthropic uses Anthropic Messages provider + // alibaba-coding-anthropic uses Anthropic Messages provider // Verify it's the anthropic messages provider by checking interface var _ LLMProvider = provider }) @@ -998,11 +1039,31 @@ func TestModelProviderOptions(t *testing.T) { if option, ok := seen["openai"]; ok && !option.CreateAllowed { t.Fatal("openai should be creatable") } + if option, ok := seen["openai"]; ok && !option.SupportsFetch { + t.Fatal("openai should support upstream model listing") + } else if option.DisplayName != "OpenAI" { + t.Fatalf("openai display_name = %q, want %q", option.DisplayName, "OpenAI") + } else if len(option.CommonModels) == 0 { + t.Fatal("openai common_models should not be empty") + } if option, ok := seen["lmstudio"]; !ok { t.Fatal("lmstudio option missing") } else if !option.EmptyAPIKeyAllowed { t.Fatal("lmstudio should allow empty API keys") } + if option, ok := seen["gpt4free"]; !ok { + t.Fatal("gpt4free option missing") + } else { + if option.DefaultAPIBase != "http://localhost:1337/v1" { + t.Fatalf("gpt4free default_api_base = %q, want %q", option.DefaultAPIBase, "http://localhost:1337/v1") + } + if !option.EmptyAPIKeyAllowed { + t.Fatal("gpt4free should allow empty API keys") + } + if !option.SupportsFetch { + t.Fatal("gpt4free should support upstream model listing") + } + } if option, ok := seen["siliconflow"]; !ok { t.Fatal("siliconflow option missing") } else if option.DefaultAPIBase != "https://api.siliconflow.cn/v1" { @@ -1052,6 +1113,71 @@ func TestModelProviderOptions(t *testing.T) { t.Fatal("github-copilot option missing") } else if option.DefaultAPIBase != "localhost:4321" { t.Fatalf("github-copilot default_api_base = %q, want %q", option.DefaultAPIBase, "localhost:4321") + } else if !option.Local { + t.Fatal("github-copilot should be marked local") + } + if option, ok := seen["qwen-portal"]; !ok { + t.Fatal("qwen-portal option missing") + } else if len(option.Aliases) == 0 || option.Aliases[0] != "qwen" { + t.Fatalf("qwen-portal aliases = %#v, want to include qwen", option.Aliases) + } + + for _, option := range options { + if len(option.CommonModels) > 6 { + t.Fatalf("provider %q exposes %d common_models, want at most 6", option.ID, len(option.CommonModels)) + } + if option.Local && len(option.CommonModels) > 0 { + t.Fatalf("local provider %q should not expose common_models", option.ID) + } + seenModels := make(map[string]struct{}, len(option.CommonModels)) + for _, model := range option.CommonModels { + if strings.TrimSpace(model) == "" { + t.Fatalf("provider %q includes an empty common_model entry", option.ID) + } + if _, exists := seenModels[model]; exists { + t.Fatalf("provider %q includes duplicate common_model %q", option.ID, model) + } + seenModels[model] = struct{}{} + } + } +} + +func TestBuildModelProviderAliasMap(t *testing.T) { + aliases := buildModelProviderAliasMap() + if len(aliases) == 0 { + t.Fatal("buildModelProviderAliasMap() returned empty map") + } + + seenAliases := make(map[string]string, len(aliases)) + for provider, option := range modelProviderOptionsByName { + got, ok := aliases[provider] + if !ok { + t.Fatalf("canonical provider %q missing from alias map", provider) + } + if got != provider { + t.Fatalf("canonical provider %q mapped to %q", provider, got) + } + if existing, ok := seenAliases[provider]; ok { + t.Fatalf("canonical provider key %q collides with provider %q", provider, existing) + } + seenAliases[provider] = provider + for _, alias := range option.Aliases { + normalized := strings.ToLower(strings.TrimSpace(alias)) + if normalized == "" { + t.Fatalf("provider %q includes empty alias", provider) + } + if existing, ok := seenAliases[normalized]; ok && existing != provider { + t.Fatalf("alias %q for provider %q collides with provider %q", alias, provider, existing) + } + seenAliases[normalized] = provider + got, ok := aliases[normalized] + if !ok { + t.Fatalf("alias %q for provider %q missing from alias map", alias, provider) + } + if got != provider { + t.Fatalf("alias %q normalized to %q, want %q", alias, got, provider) + } + } } } diff --git a/pkg/providers/fallback.go b/pkg/providers/fallback.go index 36092105..6529e211 100644 --- a/pkg/providers/fallback.go +++ b/pkg/providers/fallback.go @@ -17,6 +17,7 @@ type FallbackChain struct { type FallbackCandidate struct { Provider string Model string + DisplayName string // optional configured alias/raw model label for persistence/UI RPM int // requests per minute; 0 means unrestricted IdentityKey string // optional stable config identity for cooldown/rate limiting } @@ -32,10 +33,11 @@ func (c FallbackCandidate) StableKey() string { // FallbackResult contains the successful response and metadata about all attempts. type FallbackResult struct { - Response *LLMResponse - Provider string - Model string - Attempts []FallbackAttempt + Response *LLMResponse + Provider string + Model string + IdentityKey string + Attempts []FallbackAttempt } // FallbackAttempt records one attempt in the fallback chain. @@ -85,8 +87,9 @@ func ResolveCandidatesWithLookup( } seen[key] = true candidates = append(candidates, FallbackCandidate{ - Provider: ref.Provider, - Model: ref.Model, + Provider: ref.Provider, + Model: ref.Model, + DisplayName: candidateRaw, }) } @@ -115,6 +118,22 @@ func (fc *FallbackChain) Execute( ctx context.Context, candidates []FallbackCandidate, run func(ctx context.Context, provider, model string) (*LLMResponse, error), +) (*FallbackResult, error) { + return fc.ExecuteCandidate( + ctx, + candidates, + func(ctx context.Context, candidate FallbackCandidate) (*LLMResponse, error) { + return run(ctx, candidate.Provider, candidate.Model) + }, + ) +} + +// ExecuteCandidate runs the fallback chain and passes the complete candidate +// to the caller so model-list identity metadata remains available. +func (fc *FallbackChain) ExecuteCandidate( + ctx context.Context, + candidates []FallbackCandidate, + run func(ctx context.Context, candidate FallbackCandidate) (*LLMResponse, error), ) (*FallbackResult, error) { if len(candidates) == 0 { return nil, fmt.Errorf("fallback: no candidates configured") @@ -178,7 +197,7 @@ func (fc *FallbackChain) Execute( // Execute the run function. start := time.Now() - resp, err := run(ctx, candidate.Provider, candidate.Model) + resp, err := run(ctx, candidate) elapsed := time.Since(start) if err == nil { @@ -187,6 +206,7 @@ func (fc *FallbackChain) Execute( result.Response = resp result.Provider = candidate.Provider result.Model = candidate.Model + result.IdentityKey = candidate.StableKey() return result, nil } @@ -305,6 +325,7 @@ func (fc *FallbackChain) ExecuteImage( result.Response = resp result.Provider = candidate.Provider result.Model = candidate.Model + result.IdentityKey = candidate.StableKey() return result, nil } diff --git a/pkg/providers/httpapi/gemini_provider.go b/pkg/providers/httpapi/gemini_provider.go index 395c555d..9d0b44e1 100644 --- a/pkg/providers/httpapi/gemini_provider.go +++ b/pkg/providers/httpapi/gemini_provider.go @@ -15,8 +15,9 @@ import ( ) const ( - geminiDefaultAPIBase = "https://generativelanguage.googleapis.com/v1beta" - geminiDefaultModel = "gemini-2.0-flash" + geminiDefaultAPIBase = "https://generativelanguage.googleapis.com/v1beta" + geminiDefaultModel = "gemini-2.0-flash" + geminiDefaultStreamingReadIdleLimit = 5 * time.Minute ) type GeminiProvider struct { @@ -114,6 +115,28 @@ func (p *GeminiProvider) ChatStream( model string, options map[string]any, onChunk func(accumulated string), +) (*LLMResponse, error) { + return p.ChatStreamEvents( + ctx, + messages, + tools, + model, + options, + func(chunk StreamChunk) { + if onChunk != nil && strings.TrimSpace(chunk.Content) != "" { + onChunk(chunk.Content) + } + }, + ) +} + +func (p *GeminiProvider) ChatStreamEvents( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(StreamChunk), ) (*LLMResponse, error) { if p.apiBase == "" { return nil, fmt.Errorf("API base not configured") @@ -147,7 +170,43 @@ func (p *GeminiProvider) ChatStream( return nil, common.HandleErrorResponse(resp, p.apiBase) } - return parseGeminiStreamResponse(ctx, resp.Body, onChunk) + return parseGeminiStreamResponse(ctx, + withGeminiStreamingReadIdleTimeout(resp.Body, + geminiDefaultStreamingReadIdleLimit), + onChunk) +} + +func withGeminiStreamingReadIdleTimeout(body io.ReadCloser, timeout time.Duration) io.ReadCloser { + if body == nil || timeout <= 0 { + return body + } + return &geminiStreamingReadIdleTimeoutBody{ + body: body, + timeout: timeout, + } +} + +type geminiStreamingReadIdleTimeoutBody struct { + body io.ReadCloser + timeout time.Duration +} + +func (b *geminiStreamingReadIdleTimeoutBody) Read(p []byte) (int, error) { + timedOut := make(chan struct{}) + timer := time.AfterFunc(b.timeout, func() { + close(timedOut) + _ = b.body.Close() + }) + n, err := b.body.Read(p) + if !timer.Stop() { + <-timedOut + return n, fmt.Errorf("gemini stream idle timeout after %s", b.timeout) + } + return n, err +} + +func (b *geminiStreamingReadIdleTimeoutBody) Close() error { + return b.body.Close() } func (p *GeminiProvider) applyHeaders(req *http.Request) { @@ -458,7 +517,7 @@ func parseGeminiResponse(resp *geminiGenerateContentResponse) *LLMResponse { func parseGeminiStreamResponse( ctx context.Context, reader io.Reader, - onChunk func(accumulated string), + onChunk func(StreamChunk), ) (*LLMResponse, error) { var contentBuilder strings.Builder var reasoningBuilder strings.Builder @@ -498,10 +557,13 @@ func parseGeminiStreamResponse( if part.Text != "" { if part.Thought { reasoningBuilder.WriteString(part.Text) + if onChunk != nil { + onChunk(StreamChunk{ReasoningContent: reasoningBuilder.String()}) + } } else { contentBuilder.WriteString(part.Text) if onChunk != nil { - onChunk(contentBuilder.String()) + onChunk(StreamChunk{Content: contentBuilder.String()}) } } } diff --git a/pkg/providers/httpapi/gemini_provider_test.go b/pkg/providers/httpapi/gemini_provider_test.go index b455357c..05080b81 100644 --- a/pkg/providers/httpapi/gemini_provider_test.go +++ b/pkg/providers/httpapi/gemini_provider_test.go @@ -3,10 +3,13 @@ package httpapi import ( "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" "strings" + "sync" "testing" + "time" ) func TestGeminiProvider_ChatSeparatesThoughtAndToolCall(t *testing.T) { @@ -212,6 +215,109 @@ func TestGeminiProvider_ChatStreamParsesThoughtTextAndToolCalls(t *testing.T) { } } +func TestGeminiProvider_ChatStreamEventsStreamsThoughtBeforeContent(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, ":streamGenerateContent") { + t.Fatalf("path = %s, expected streamGenerateContent endpoint", r.URL.Path) + } + w.Header().Set("Content-Type", "text/event-stream") + flusher, ok := w.(http.Flusher) + if !ok { + t.Fatal("response writer is not flushable") + } + + chunk := map[string]any{ + "candidates": []any{map[string]any{ + "content": map[string]any{ + "parts": []any{ + map[string]any{"text": "think", "thought": true}, + map[string]any{"text": "answer"}, + }, + }, + "finishReason": "STOP", + }}, + } + raw, err := json.Marshal(chunk) + if err != nil { + t.Fatalf("marshal chunk: %v", err) + } + _, _ = fmt.Fprintf(w, "data: %s\n\n", raw) + _, _ = fmt.Fprint(w, "data: [DONE]\n\n") + flusher.Flush() + })) + defer server.Close() + + provider := NewGeminiProvider("test-key", server.URL, "", "", 0, nil, nil) + events := make([]string, 0) + resp, err := provider.ChatStreamEvents( + t.Context(), + []Message{{Role: "user", Content: "hello"}}, + nil, + "gemini-2.5-flash", + nil, + func(chunk StreamChunk) { + if chunk.ReasoningContent != "" { + events = append(events, "reasoning:"+chunk.ReasoningContent) + } + if chunk.Content != "" { + events = append(events, "content:"+chunk.Content) + } + }, + ) + if err != nil { + t.Fatalf("ChatStreamEvents() error = %v", err) + } + if resp.Content != "answer" { + t.Fatalf("Content = %q, want %q", resp.Content, "answer") + } + if resp.ReasoningContent != "think" { + t.Fatalf("ReasoningContent = %q, want %q", resp.ReasoningContent, "think") + } + want := []string{"reasoning:think", "content:answer"} + if len(events) != len(want) { + t.Fatalf("events = %#v, want %#v", events, want) + } + for i := range want { + if events[i] != want[i] { + t.Fatalf("events = %#v, want %#v", events, want) + } + } +} + +type geminiBlockingReadCloser struct { + closeOnce sync.Once + closed chan struct{} +} + +func newGeminiBlockingReadCloser() *geminiBlockingReadCloser { + return &geminiBlockingReadCloser{closed: make(chan struct{})} +} + +func (r *geminiBlockingReadCloser) Read([]byte) (int, error) { + <-r.closed + return 0, io.ErrClosedPipe +} + +func (r *geminiBlockingReadCloser) Close() error { + r.closeOnce.Do(func() { + close(r.closed) + }) + return nil +} + +func TestGeminiStreamingReadIdleTimeoutClosesStalledBody(t *testing.T) { + body := newGeminiBlockingReadCloser() + wrapped := withGeminiStreamingReadIdleTimeout(body, 10*time.Millisecond) + + _, err := wrapped.Read(make([]byte, 1)) + if err == nil { + t.Fatal("expected stalled stream read to return an error") + } + if !strings.Contains(err.Error(), "gemini stream idle timeout") { + t.Fatalf("error = %v, want gemini stream idle timeout", err) + } +} + func TestGeminiProvider_ChatStreamSkipsEmptyDataFrames(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") diff --git a/pkg/providers/httpapi/http_provider.go b/pkg/providers/httpapi/http_provider.go index 90f389cc..4d6fbec9 100644 --- a/pkg/providers/httpapi/http_provider.go +++ b/pkg/providers/httpapi/http_provider.go @@ -70,6 +70,17 @@ func (p *HTTPProvider) ChatStream( return p.delegate.ChatStream(ctx, messages, tools, model, options, onChunk) } +func (p *HTTPProvider) ChatStreamEvents( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(StreamChunk), +) (*LLMResponse, error) { + return p.delegate.ChatStreamEvents(ctx, messages, tools, model, options, onChunk) +} + func (p *HTTPProvider) GetDefaultModel() string { return "" } @@ -78,6 +89,13 @@ func (p *HTTPProvider) SupportsNativeSearch() bool { return p.delegate.SupportsNativeSearch() } +func (p *HTTPProvider) SupportsThinking() bool { + if p == nil || p.delegate == nil { + return false + } + return p.delegate.SupportsThinking() +} + func (p *HTTPProvider) SetProviderName(providerName string) { if p == nil || p.delegate == nil { return diff --git a/pkg/providers/httpapi/types.go b/pkg/providers/httpapi/types.go index c8bcdc0d..2b4fcef7 100644 --- a/pkg/providers/httpapi/types.go +++ b/pkg/providers/httpapi/types.go @@ -10,6 +10,7 @@ type ( ToolCall = protocoltypes.ToolCall FunctionCall = protocoltypes.FunctionCall LLMResponse = protocoltypes.LLMResponse + StreamChunk = protocoltypes.StreamChunk UsageInfo = protocoltypes.UsageInfo Message = protocoltypes.Message ToolDefinition = protocoltypes.ToolDefinition @@ -41,3 +42,14 @@ type StreamingProvider interface { onChunk func(accumulated string), ) (*LLMResponse, error) } + +type StreamingEventProvider interface { + ChatStreamEvents( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(StreamChunk), + ) (*LLMResponse, error) +} diff --git a/pkg/providers/model_ref.go b/pkg/providers/model_ref.go index 48e3fb4c..ac7cae78 100644 --- a/pkg/providers/model_ref.go +++ b/pkg/providers/model_ref.go @@ -29,46 +29,14 @@ func ParseModelRef(raw string, defaultProvider string) *ModelRef { // NormalizeProvider normalizes provider identifiers to canonical form. func NormalizeProvider(provider string) string { - p := strings.ToLower(strings.TrimSpace(provider)) - - switch p { - case "z.ai", "z-ai": - return "zai" - case "opencode-zen": - return "opencode" - case "qwen": - return "qwen-portal" - case "kimi-code": - return "kimi-coding" - case "gpt": - return "openai" - case "claude": - return "anthropic" - case "glm": - return "zhipu" - case "google": - return "gemini" - case "google-antigravity": - return "antigravity" - case "alibaba-coding", "qwen-coding": - return "coding-plan" - case "alibaba-coding-anthropic": - return "coding-plan-anthropic" - case "qwen-international", "dashscope-intl": - return "qwen-intl" - case "dashscope-us": - return "qwen-us" - case "azure-openai": - return "azure" - case "claudecli": - return "claude-cli" - case "codexcli": - return "codex-cli" - case "copilot": - return "github-copilot" + normalized := strings.ToLower(strings.TrimSpace(provider)) + if normalized == "" { + return "" } - - return p + if canonical, ok := normalizedModelProviderAliasesByName[normalized]; ok { + return canonical + } + return normalized } // ModelKey returns a canonical "provider/model" key for deduplication. diff --git a/pkg/providers/model_ref_test.go b/pkg/providers/model_ref_test.go index 9a164bf4..57438a44 100644 --- a/pkg/providers/model_ref_test.go +++ b/pkg/providers/model_ref_test.go @@ -65,9 +65,7 @@ func TestNormalizeProvider(t *testing.T) { {"z.ai", "zai"}, {"z-ai", "zai"}, {"Z.AI", "zai"}, - {"opencode-zen", "opencode"}, {"qwen", "qwen-portal"}, - {"kimi-code", "kimi-coding"}, {"gpt", "openai"}, {"claude", "anthropic"}, {"glm", "zhipu"}, @@ -78,10 +76,13 @@ func TestNormalizeProvider(t *testing.T) { {"claudecli", "claude-cli"}, {"codexcli", "codex-cli"}, {"copilot", "github-copilot"}, + {"g4f", "gpt4free"}, // Alibaba Coding Plan aliases - {"alibaba-coding", "coding-plan"}, - {"qwen-coding", "coding-plan"}, - {"alibaba-coding-anthropic", "coding-plan-anthropic"}, + {"alibaba-coding", "alibaba-coding"}, + {"coding-plan", "alibaba-coding"}, + {"qwen-coding", "alibaba-coding"}, + {"alibaba-coding-anthropic", "alibaba-coding-anthropic"}, + {"coding-plan-anthropic", "alibaba-coding-anthropic"}, // Qwen international aliases {"qwen-international", "qwen-intl"}, {"dashscope-intl", "qwen-intl"}, diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index c154544c..ff30d49c 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers/common" "github.com/sipeed/picoclaw/pkg/providers/messageutil" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" @@ -23,6 +24,7 @@ type ( ToolCall = protocoltypes.ToolCall FunctionCall = protocoltypes.FunctionCall LLMResponse = protocoltypes.LLMResponse + StreamChunk = protocoltypes.StreamChunk UsageInfo = protocoltypes.UsageInfo Message = protocoltypes.Message ToolDefinition = protocoltypes.ToolDefinition @@ -45,7 +47,10 @@ type Provider struct { type Option func(*Provider) -const defaultRequestTimeout = common.DefaultRequestTimeout +const ( + defaultRequestTimeout = common.DefaultRequestTimeout + defaultStreamingReadIdleTimeout = 5 * time.Minute +) var stripModelPrefixProviders = map[string]struct{}{ "litellm": {}, @@ -189,6 +194,8 @@ func (p *Provider) buildRequestBody( } } + p.applyThinkingControl(requestBody, model, options) + // Merge extra body fields configured per-provider/model. // These are injected last so they take precedence over defaults. maps.Copy(requestBody, p.extraBody) @@ -196,6 +203,112 @@ func (p *Provider) buildRequestBody( return requestBody } +func (p *Provider) applyThinkingControl(requestBody map[string]any, model string, options map[string]any) { + level, ok := normalizedThinkingLevel(options) + if !ok { + return + } + + if p.SupportsThinking() { + p.applyDeepSeekThinkingControl(requestBody, level) + return + } + + if level != "off" { + return + } + + switch p.thinkingControlKind(model) { + case "thinking_type": + requestBody["thinking"] = map[string]any{"type": "disabled"} + case "enable_thinking": + requestBody["enable_thinking"] = false + } +} + +func (p *Provider) applyDeepSeekThinkingControl(requestBody map[string]any, level string) { + switch level { + case "off": + requestBody["thinking"] = map[string]any{"type": "disabled"} + case "low", "medium", "high": + requestBody["thinking"] = map[string]any{"type": "enabled"} + requestBody["reasoning_effort"] = "high" + case "xhigh": + requestBody["thinking"] = map[string]any{"type": "enabled"} + requestBody["reasoning_effort"] = "max" + case "adaptive": + logger.WarnCF("provider.openai_compat", + `DeepSeek does not support thinking_level="adaptive"; using provider default thinking behavior`, + map[string]any{ + "provider": p.providerName, + "api_base": p.apiBase, + "thinking_level": level, + }, + ) + } +} + +func normalizedThinkingLevel(options map[string]any) (string, bool) { + raw, ok := options["thinking_level"].(string) + if !ok { + return "", false + } + switch strings.ToLower(strings.TrimSpace(raw)) { + case "off", "low", "medium", "high", "xhigh", "adaptive": + return strings.ToLower(strings.TrimSpace(raw)), true + default: + return "", false + } +} + +func (p *Provider) thinkingControlKind(model string) string { + providerName := strings.ToLower(strings.TrimSpace(p.providerName)) + lowerModel := strings.ToLower(strings.TrimSpace(model)) + + switch providerName { + case "volcengine": + return "thinking_type" + case "zhipu", "zai": + return "thinking_type" + case "qwen", "qwen-portal", "qwen-intl", "qwen-international", "dashscope-intl", "qwen-us", "dashscope-us": + return "enable_thinking" + case "modelscope": + if strings.Contains(lowerModel, "qwen") { + return "enable_thinking" + } + } + + if providerName == "openai" || providerName == "" { + if isVolcengineHost(p.apiBase) || strings.Contains(lowerModel, "doubao") { + return "thinking_type" + } + if isDashScopeHost(p.apiBase) || strings.Contains(lowerModel, "qwen") { + return "enable_thinking" + } + } + + return "" +} + +func isVolcengineHost(apiBase string) bool { + host := normalizedHostname(apiBase) + return host == "volcengine.com" || strings.HasSuffix(host, ".volcengine.com") || + host == "volces.com" || strings.HasSuffix(host, ".volces.com") +} + +func isDashScopeHost(apiBase string) bool { + host := normalizedHostname(apiBase) + return host == "dashscope.aliyuncs.com" || strings.HasSuffix(host, ".dashscope.aliyuncs.com") +} + +func normalizedHostname(rawURL string) string { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil { + return "" + } + return strings.ToLower(strings.TrimSpace(parsed.Hostname())) +} + func (p *Provider) applyCustomHeaders(req *http.Request) { for k, v := range p.customHeaders { if strings.TrimSpace(k) == "" { @@ -209,6 +322,10 @@ func (p *Provider) SetProviderName(providerName string) { p.providerName = strings.ToLower(strings.TrimSpace(providerName)) } +func (p *Provider) SupportsThinking() bool { + return strings.EqualFold(strings.TrimSpace(p.providerName), "deepseek") || isDeepSeekHost(p.apiBase) +} + func (p *Provider) prepareMessagesForRequest(messages []Message) []Message { if len(messages) == 0 { return nil @@ -380,6 +497,28 @@ func (p *Provider) ChatStream( model string, options map[string]any, onChunk func(accumulated string), +) (*LLMResponse, error) { + return p.ChatStreamEvents( + ctx, + messages, + tools, + model, + options, + func(chunk StreamChunk) { + if onChunk != nil && strings.TrimSpace(chunk.Content) != "" { + onChunk(chunk.Content) + } + }, + ) +} + +func (p *Provider) ChatStreamEvents( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(StreamChunk), ) (*LLMResponse, error) { if p.apiBase == "" { return nil, fmt.Errorf("API base not configured") @@ -422,14 +561,47 @@ func (p *Provider) ChatStream( return nil, common.HandleErrorResponse(resp, p.apiBase) } - return parseStreamResponse(ctx, resp.Body, onChunk) + return parseStreamResponse(ctx, withStreamingReadIdleTimeout(resp.Body, defaultStreamingReadIdleTimeout), onChunk) +} + +func withStreamingReadIdleTimeout(body io.ReadCloser, timeout time.Duration) io.ReadCloser { + if body == nil || timeout <= 0 { + return body + } + return &streamingReadIdleTimeoutBody{ + body: body, + timeout: timeout, + } +} + +type streamingReadIdleTimeoutBody struct { + body io.ReadCloser + timeout time.Duration +} + +func (b *streamingReadIdleTimeoutBody) Read(p []byte) (int, error) { + timedOut := make(chan struct{}) + timer := time.AfterFunc(b.timeout, func() { + close(timedOut) + _ = b.body.Close() + }) + n, err := b.body.Read(p) + if !timer.Stop() { + <-timedOut + return n, fmt.Errorf("stream idle timeout after %s", b.timeout) + } + return n, err +} + +func (b *streamingReadIdleTimeoutBody) Close() error { + return b.body.Close() } // parseStreamResponse parses an OpenAI-compatible SSE stream. func parseStreamResponse( ctx context.Context, reader io.Reader, - onChunk func(accumulated string), + onChunk func(StreamChunk), ) (*LLMResponse, error) { var textContent strings.Builder var reasoningContent strings.Builder @@ -489,22 +661,29 @@ func parseStreamResponse( choice := chunk.Choices[0] - // Accumulate text content - if choice.Delta.Content != "" { - textContent.WriteString(choice.Delta.Content) - if onChunk != nil { - onChunk(textContent.String()) - } - } if choice.Delta.ReasoningContent != "" { reasoningContent.WriteString(choice.Delta.ReasoningContent) + if onChunk != nil { + onChunk(StreamChunk{ReasoningContent: reasoningContent.String()}) + } } if choice.Delta.Reasoning != "" { reasoning.WriteString(choice.Delta.Reasoning) + if onChunk != nil { + onChunk(StreamChunk{ReasoningContent: reasoning.String()}) + } } if len(choice.Delta.ReasoningDetails) > 0 { reasoningDetails = append(reasoningDetails, choice.Delta.ReasoningDetails...) } + // Accumulate text content after reasoning so UIs can show thought first + // when a provider sends both fields in the same event. + if choice.Delta.Content != "" { + textContent.WriteString(choice.Delta.Content) + if onChunk != nil { + onChunk(StreamChunk{Content: textContent.String()}) + } + } // Accumulate tool call deltas for _, tc := range choice.Delta.ToolCalls { diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index 5caf3039..bc7e12de 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -8,10 +8,13 @@ import ( "net/http" "net/http/httptest" "net/url" + "os" "strings" + "sync" "testing" "time" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers/common" "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) @@ -61,6 +64,209 @@ func TestProviderChat_UsesMaxCompletionTokensForGLM(t *testing.T) { } } +func TestBuildRequestBody_DisablesDoubaoThinkingWhenThinkingLevelOff(t *testing.T) { + p := NewProvider("key", "https://ark.cn-beijing.volces.com/api/v3", "") + p.SetProviderName("openai") + + body := p.buildRequestBody( + []Message{{Role: "user", Content: "hi"}}, + nil, + "doubao-seed-1-6-flash-250828", + map[string]any{"thinking_level": "off"}, + ) + + thinking, ok := body["thinking"].(map[string]any) + if !ok { + t.Fatalf("thinking = %#v, want map", body["thinking"]) + } + if got := thinking["type"]; got != "disabled" { + t.Fatalf("thinking.type = %#v, want %q", got, "disabled") + } +} + +func TestBuildRequestBody_DisablesModelDependentQwenThinkingWhenThinkingLevelOff(t *testing.T) { + p := NewProvider("key", "https://api-inference.modelscope.cn/v1", "") + p.SetProviderName("modelscope") + + body := p.buildRequestBody( + []Message{{Role: "user", Content: "hi"}}, + nil, + "qwen3-coder-plus", + map[string]any{"thinking_level": "off"}, + ) + + if got := body["enable_thinking"]; got != false { + t.Fatalf("enable_thinking = %#v, want false", got) + } +} + +func TestBuildRequestBody_PreservesDoubaoRequestWhenThinkingLevelIsNotOff(t *testing.T) { + p := NewProvider("key", "https://ark.cn-beijing.volces.com/api/v3", "") + p.SetProviderName("openai") + + for _, level := range []string{"low", "adaptive", "unexpected"} { + t.Run(level, func(t *testing.T) { + body := p.buildRequestBody( + []Message{{Role: "user", Content: "hi"}}, + nil, + "doubao-seed-1-6-flash-250828", + map[string]any{"thinking_level": level}, + ) + + if _, ok := body["thinking"]; ok { + t.Fatalf( + "thinking should be omitted for %q to preserve existing behavior, got %#v", + level, + body["thinking"], + ) + } + if _, ok := body["enable_thinking"]; ok { + t.Fatalf("enable_thinking should be omitted for %q, got %#v", level, body["enable_thinking"]) + } + }) + } +} + +func TestBuildRequestBody_MapsDeepSeekThinkingLevels(t *testing.T) { + p := NewProvider("key", "https://api.deepseek.com/v1", "") + p.SetProviderName("deepseek") + + tests := []struct { + name string + level string + wantThinkingType string + wantEffort any + }{ + {name: "off", level: "off", wantThinkingType: "disabled"}, + {name: "low", level: "low", wantThinkingType: "enabled", wantEffort: "high"}, + {name: "medium", level: "medium", wantThinkingType: "enabled", wantEffort: "high"}, + {name: "high", level: "high", wantThinkingType: "enabled", wantEffort: "high"}, + {name: "xhigh", level: "xhigh", wantThinkingType: "enabled", wantEffort: "max"}, + {name: "adaptive", level: "adaptive"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := p.buildRequestBody( + []Message{{Role: "user", Content: "hi"}}, + nil, + "deepseek-v4-pro", + map[string]any{"thinking_level": tt.level}, + ) + + if tt.wantThinkingType == "" { + if _, ok := body["thinking"]; ok { + t.Fatalf("thinking should be omitted for %q, got %#v", tt.level, body["thinking"]) + } + } else { + thinking, ok := body["thinking"].(map[string]any) + if !ok { + t.Fatalf("thinking = %#v, want map", body["thinking"]) + } + if got := thinking["type"]; got != tt.wantThinkingType { + t.Fatalf("thinking.type = %#v, want %q", got, tt.wantThinkingType) + } + } + + if tt.wantEffort == nil { + if _, ok := body["reasoning_effort"]; ok { + t.Fatalf("reasoning_effort should be omitted for %q, got %#v", tt.level, body["reasoning_effort"]) + } + } else if got := body["reasoning_effort"]; got != tt.wantEffort { + t.Fatalf("reasoning_effort = %#v, want %#v", got, tt.wantEffort) + } + }) + } +} + +func TestBuildRequestBody_MapsDeepSeekThinkingLevelsByHost(t *testing.T) { + p := NewProvider("key", "https://api.deepseek.com/v1", "") + + body := p.buildRequestBody( + []Message{{Role: "user", Content: "hi"}}, + nil, + "deepseek-v4-flash", + map[string]any{"thinking_level": "xhigh"}, + ) + + thinking, ok := body["thinking"].(map[string]any) + if !ok { + t.Fatalf("thinking = %#v, want map", body["thinking"]) + } + if got := thinking["type"]; got != "enabled" { + t.Fatalf("thinking.type = %#v, want enabled", got) + } + if got := body["reasoning_effort"]; got != "max" { + t.Fatalf("reasoning_effort = %#v, want max", got) + } +} + +func TestBuildRequestBody_DeepSeekExtraBodyStillOverridesThinkingFields(t *testing.T) { + extraBody := map[string]any{ + "thinking": map[string]any{"type": "disabled"}, + "reasoning_effort": "max", + } + p := NewProvider("key", "https://api.deepseek.com/v1", "", WithExtraBody(extraBody)) + p.SetProviderName("deepseek") + + body := p.buildRequestBody( + []Message{{Role: "user", Content: "hi"}}, + nil, + "deepseek-v4-pro", + map[string]any{"thinking_level": "high"}, + ) + + thinking, ok := body["thinking"].(map[string]any) + if !ok { + t.Fatalf("thinking = %#v, want map", body["thinking"]) + } + if got := thinking["type"]; got != "disabled" { + t.Fatalf("thinking.type = %#v, want disabled from extra_body override", got) + } + if got := body["reasoning_effort"]; got != "max" { + t.Fatalf("reasoning_effort = %#v, want max from extra_body override", got) + } +} + +func TestBuildRequestBody_WarnsForUnsupportedDeepSeekAdaptiveThinkingLevel(t *testing.T) { + logFile := t.TempDir() + "/deepseek-adaptive-warning.log" + prevLevel := logger.GetLevel() + logger.SetLevel(logger.WARN) + if err := logger.EnableFileLogging(logFile); err != nil { + t.Fatalf("EnableFileLogging() error = %v", err) + } + defer func() { + logger.DisableFileLogging() + logger.SetLevel(prevLevel) + }() + + p := NewProvider("key", "https://api.deepseek.com/v1", "") + p.SetProviderName("deepseek") + + body := p.buildRequestBody( + []Message{{Role: "user", Content: "hi"}}, + nil, + "deepseek-v4-pro", + map[string]any{"thinking_level": "adaptive"}, + ) + + if _, ok := body["thinking"]; ok { + t.Fatalf("thinking should be omitted for adaptive, got %#v", body["thinking"]) + } + if _, ok := body["reasoning_effort"]; ok { + t.Fatalf("reasoning_effort should be omitted for adaptive, got %#v", body["reasoning_effort"]) + } + + data, err := os.ReadFile(logFile) + if err != nil { + t.Fatalf("ReadFile(%q) error = %v", logFile, err) + } + logs := string(data) + if !strings.Contains(logs, `thinking_level=\"adaptive\"`) { + t.Fatalf("warning log = %q, want adaptive warning message", logs) + } +} + func TestProviderChat_ParsesToolCalls(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { resp := map[string]any{ @@ -1318,6 +1524,53 @@ func TestProviderChatStream_ParsesReasoningContent(t *testing.T) { } } +func TestProviderChatStreamEvents_EmitsReasoningBeforeContentFromSameEvent(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte( + "data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"think\",\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n", + )) + _, _ = w.Write([]byte("data: [DONE]\n\n")) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + events := make([]string, 0) + out, err := p.ChatStreamEvents( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "deepseek-v4-flash", + nil, + func(chunk StreamChunk) { + if chunk.ReasoningContent != "" { + events = append(events, "reasoning:"+chunk.ReasoningContent) + } + if chunk.Content != "" { + events = append(events, "content:"+chunk.Content) + } + }, + ) + if err != nil { + t.Fatalf("ChatStreamEvents() error = %v", err) + } + if out.Content != "answer" { + t.Fatalf("Content = %q, want %q", out.Content, "answer") + } + if out.ReasoningContent != "think" { + t.Fatalf("ReasoningContent = %q, want %q", out.ReasoningContent, "think") + } + want := []string{"reasoning:think", "content:answer"} + if len(events) != len(want) { + t.Fatalf("events = %#v, want %#v", events, want) + } + for i := range want { + if events[i] != want[i] { + t.Fatalf("events = %#v, want %#v", events, want) + } + } +} + func TestProviderChatStream_ParsesMultilineSSEEvent(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") @@ -1463,6 +1716,40 @@ func (r *errAfterDataReadCloser) Close() error { return nil } +type blockingReadCloser struct { + closeOnce sync.Once + closed chan struct{} +} + +func newBlockingReadCloser() *blockingReadCloser { + return &blockingReadCloser{closed: make(chan struct{})} +} + +func (r *blockingReadCloser) Read([]byte) (int, error) { + <-r.closed + return 0, io.ErrClosedPipe +} + +func (r *blockingReadCloser) Close() error { + r.closeOnce.Do(func() { + close(r.closed) + }) + return nil +} + +func TestStreamingReadIdleTimeoutClosesStalledBody(t *testing.T) { + body := newBlockingReadCloser() + wrapped := withStreamingReadIdleTimeout(body, 10*time.Millisecond) + + _, err := wrapped.Read(make([]byte, 1)) + if err == nil { + t.Fatal("expected stalled stream read to return an error") + } + if !strings.Contains(err.Error(), "stream idle timeout") { + t.Fatalf("error = %v, want stream idle timeout", err) + } +} + func TestProvider_FunctionalOptionMaxTokensField(t *testing.T) { p := NewProvider("key", "https://example.com/v1", "", WithMaxTokensField("max_completion_tokens")) if p.maxTokensField != "max_completion_tokens" { diff --git a/pkg/providers/protocoltypes/types.go b/pkg/providers/protocoltypes/types.go index bab4433e..b60c4670 100644 --- a/pkg/providers/protocoltypes/types.go +++ b/pkg/providers/protocoltypes/types.go @@ -1,5 +1,7 @@ package protocoltypes +import "time" + type ToolCall struct { ID string `json:"id"` Type string `json:"type,omitempty"` @@ -35,6 +37,11 @@ type LLMResponse struct { ReasoningDetails []ReasoningDetail `json:"reasoning_details"` } +type StreamChunk struct { + Content string + ReasoningContent string +} + type ReasoningDetail struct { Format string `json:"format"` Index int `json:"index"` @@ -81,6 +88,8 @@ type Attachment struct { type Message struct { Role string `json:"role"` Content string `json:"content"` + ModelName string `json:"model_name,omitempty"` + CreatedAt *time.Time `json:"created_at,omitempty"` Media []string `json:"media,omitempty"` Attachments []Attachment `json:"attachments,omitempty"` ReasoningContent string `json:"reasoning_content,omitempty"` diff --git a/pkg/providers/provider_catalog.go b/pkg/providers/provider_catalog.go index a9178cb8..b2056563 100644 --- a/pkg/providers/provider_catalog.go +++ b/pkg/providers/provider_catalog.go @@ -5,97 +5,10 @@ import ( "strings" ) -// ModelProviderOption describes a canonical provider entry exposed to the Web UI. -type ModelProviderOption struct { - ID string `json:"id"` - DefaultAPIBase string `json:"default_api_base"` - EmptyAPIKeyAllowed bool `json:"empty_api_key_allowed"` - CreateAllowed bool `json:"create_allowed"` - DefaultModelAllowed bool `json:"default_model_allowed"` - DefaultAuthMethod string `json:"default_auth_method,omitempty"` - AuthMethodLocked bool `json:"auth_method_locked,omitempty"` -} - -type attachedModelProviderMeta struct { - protocolMeta - createAllowed bool - defaultModelAllowed bool - defaultAuthMethod string - authMethodLocked bool -} - -// attachedModelProviderMetaByName augments protocolMetaByName for provider -// families that are implemented in CreateProviderFromConfig but intentionally -// kept out of the core HTTP metadata map because they have special auth/runtime -// semantics. -var attachedModelProviderMetaByName = map[string]attachedModelProviderMeta{ - "azure": {createAllowed: true, defaultModelAllowed: true}, - "anthropic": { - protocolMeta: protocolMeta{defaultAPIBase: "https://api.anthropic.com/v1"}, - createAllowed: true, - defaultModelAllowed: true, - }, - "anthropic-messages": { - protocolMeta: protocolMeta{defaultAPIBase: "https://api.anthropic.com/v1"}, - createAllowed: true, - defaultModelAllowed: true, - }, - "bedrock": {createAllowed: true, defaultModelAllowed: true}, - "antigravity": { - createAllowed: true, - defaultModelAllowed: true, - defaultAuthMethod: "oauth", - authMethodLocked: true, - }, - "claude-cli": {createAllowed: true, defaultModelAllowed: true}, - "codex-cli": {createAllowed: true, defaultModelAllowed: true}, - "github-copilot": { - protocolMeta: protocolMeta{defaultAPIBase: "localhost:4321"}, - createAllowed: true, - defaultModelAllowed: true, - }, - // ElevenLabs is intentionally exposed only as an ASR-capable provider. It - // belongs in the shared model catalog because ASR is configured via - // model_list, but it must not be selectable as the default chat model. - "elevenlabs": { - protocolMeta: protocolMeta{defaultAPIBase: "https://api.elevenlabs.io"}, - createAllowed: true, - defaultModelAllowed: false, - }, -} - // ModelProviderOptions returns the canonical provider catalog exposed to the Web UI. func ModelProviderOptions() []ModelProviderOption { - optionsByID := make(map[string]ModelProviderOption, len(protocolMetaByName)+len(attachedModelProviderMetaByName)) - for provider := range protocolMetaByName { - if NormalizeProvider(provider) != provider { - continue - } - optionsByID[provider] = ModelProviderOption{ - ID: provider, - DefaultAPIBase: DefaultAPIBaseForProtocol(provider), - EmptyAPIKeyAllowed: IsEmptyAPIKeyAllowedForProtocol(provider), - CreateAllowed: true, - DefaultModelAllowed: true, - } - } - for provider, meta := range attachedModelProviderMetaByName { - if NormalizeProvider(provider) != provider { - continue - } - optionsByID[provider] = ModelProviderOption{ - ID: provider, - DefaultAPIBase: meta.defaultAPIBase, - EmptyAPIKeyAllowed: meta.emptyAPIKeyAllowed, - CreateAllowed: meta.createAllowed, - DefaultModelAllowed: meta.defaultModelAllowed, - DefaultAuthMethod: meta.defaultAuthMethod, - AuthMethodLocked: meta.authMethodLocked, - } - } - - options := make([]ModelProviderOption, 0, len(optionsByID)) - for _, option := range optionsByID { + options := make([]ModelProviderOption, 0, len(modelProviderOptionsByName)) + for _, option := range modelProviderOptionsByName { options = append(options, option) } sort.Slice(options, func(i, j int) bool { @@ -107,44 +20,30 @@ func ModelProviderOptions() []ModelProviderOption { // IsSupportedModelProvider reports whether provider resolves to a provider ID // returned by ModelProviderOptions. func IsSupportedModelProvider(provider string) bool { - normalized := NormalizeProvider(provider) - if normalized == "" { - return false - } - if _, ok := protocolMetaByName[normalized]; ok { - return true - } - _, ok := attachedModelProviderMetaByName[normalized] + _, ok := modelProviderOptionForName(provider) return ok } +// IsModelProviderFetchable reports whether provider supports upstream /models +// listing through the launcher fetch endpoint. +func IsModelProviderFetchable(provider string) bool { + option, ok := modelProviderOptionForName(provider) + return ok && option.SupportsFetch +} + // IsCreatableModelProvider reports whether provider can be selected for a new // model entry from the Web UI. func IsCreatableModelProvider(provider string) bool { - normalized := NormalizeProvider(provider) - if normalized == "" { - return false - } - if _, ok := protocolMetaByName[normalized]; ok { - return true - } - meta, ok := attachedModelProviderMetaByName[normalized] - return ok && meta.createAllowed + option, ok := modelProviderOptionForName(provider) + return ok && option.CreateAllowed } // IsDefaultModelProvider reports whether provider can be used as the default // chat model. Some providers such as ASR-only entries are intentionally // exposed in model_list management but cannot drive the gateway default model. func IsDefaultModelProvider(provider string) bool { - normalized := NormalizeProvider(provider) - if normalized == "" { - return false - } - if _, ok := protocolMetaByName[normalized]; ok { - return true - } - meta, ok := attachedModelProviderMetaByName[normalized] - return ok && meta.defaultModelAllowed + option, ok := modelProviderOptionForName(provider) + return ok && option.DefaultModelAllowed } // SplitModelProviderAndID separates a legacy "provider/model" string into its diff --git a/pkg/providers/provider_metadata.go b/pkg/providers/provider_metadata.go new file mode 100644 index 00000000..7b9c9e0e --- /dev/null +++ b/pkg/providers/provider_metadata.go @@ -0,0 +1,595 @@ +package providers + +import "strings" + +// ModelProviderOption describes a canonical provider entry exposed to the Web UI. +// It also serves as the backend-owned source of truth for shared provider metadata. +type ModelProviderOption struct { + ID string `json:"id"` + DisplayName string `json:"display_name,omitempty"` + IconSlug string `json:"icon_slug,omitempty"` + Domain string `json:"domain,omitempty"` + DefaultAPIBase string `json:"default_api_base"` + EmptyAPIKeyAllowed bool `json:"empty_api_key_allowed"` + CreateAllowed bool `json:"create_allowed"` + DefaultModelAllowed bool `json:"default_model_allowed"` + SupportsFetch bool `json:"supports_fetch,omitempty"` + DefaultAuthMethod string `json:"default_auth_method,omitempty"` + AuthMethodLocked bool `json:"auth_method_locked,omitempty"` + Local bool `json:"local,omitempty"` + Priority float64 `json:"priority,omitempty"` + CommonModels []string `json:"common_models,omitempty"` + Aliases []string `json:"aliases,omitempty"` + + httpAPI bool `json:"-"` +} + +var modelProviderOptionsByName = map[string]ModelProviderOption{ + "openai": { + ID: "openai", + DisplayName: "OpenAI", + IconSlug: "openai", + Domain: "openai.com", + DefaultAPIBase: "https://api.openai.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 100, + CommonModels: []string{"gpt-5.4", "gpt-5.4-mini", "gpt-5.5"}, + Aliases: []string{"gpt"}, + httpAPI: true, + }, + "anthropic": { + ID: "anthropic", + DisplayName: "Anthropic", + IconSlug: "anthropic", + Domain: "anthropic.com", + DefaultAPIBase: "https://api.anthropic.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 95, + CommonModels: []string{"claude-opus-4-7", "claude-sonnet-4-6", "claude-haiku-4-5"}, + Aliases: []string{"claude"}, + httpAPI: true, + }, + "anthropic-messages": { + ID: "anthropic-messages", + DisplayName: "Anthropic Messages", + IconSlug: "anthropic", + Domain: "anthropic.com", + DefaultAPIBase: "https://api.anthropic.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 94, + CommonModels: []string{"claude-opus-4-7", "claude-sonnet-4-6", "claude-haiku-4-5"}, + httpAPI: true, + }, + "gemini": { + ID: "gemini", + DisplayName: "Google Gemini", + IconSlug: "googlegemini", + Domain: "gemini.google.com", + DefaultAPIBase: "https://generativelanguage.googleapis.com/v1beta", + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 90, + CommonModels: []string{"gemini-3.1-pro-preview", "gemini-3-flash-preview", "gemini-3.1-flash-lite"}, + Aliases: []string{"google"}, + httpAPI: true, + }, + "deepseek": { + ID: "deepseek", + DisplayName: "DeepSeek", + IconSlug: "deepseek", + Domain: "deepseek.com", + DefaultAPIBase: "https://api.deepseek.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 85, + CommonModels: []string{"deepseek-v4-flash", "deepseek-v4-pro"}, + httpAPI: true, + }, + "openrouter": { + ID: "openrouter", + DisplayName: "OpenRouter", + IconSlug: "openrouter", + Domain: "openrouter.ai", + DefaultAPIBase: "https://openrouter.ai/api/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 80, + CommonModels: []string{ + "openai/gpt-5.4", + "anthropic/claude-opus-4.7", + "google/gemini-3.1-pro-preview", + "qwen/qwen3-coder-next", + }, + httpAPI: true, + }, + "qwen-portal": { + ID: "qwen-portal", + DisplayName: "Qwen", + IconSlug: "alibabacloud", + Domain: "qwenlm.ai", + DefaultAPIBase: "https://dashscope.aliyuncs.com/compatible-mode/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 75, + CommonModels: []string{"qwen3.6-max-preview", "qwen3.6-plus", "qwen3.6-flash", "qwen3-coder-next"}, + Aliases: []string{"qwen"}, + httpAPI: true, + }, + "qwen-intl": { + ID: "qwen-intl", + DisplayName: "Qwen International", + IconSlug: "alibabacloud", + Domain: "alibabacloud.com", + DefaultAPIBase: "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 74, + CommonModels: []string{"qwen3.6-max-preview", "qwen3.6-plus", "qwen3.6-flash", "qwen3-coder-next"}, + Aliases: []string{"qwen-international", "dashscope-intl"}, + httpAPI: true, + }, + "moonshot": { + ID: "moonshot", + DisplayName: "Moonshot", + Domain: "moonshot.ai", + DefaultAPIBase: "https://api.moonshot.cn/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 70, + CommonModels: []string{ + "kimi-k2.5", + "kimi-k2-thinking", + "kimi-k2-thinking-turbo", + "kimi-k2-turbo-preview", + }, + httpAPI: true, + }, + "volcengine": { + ID: "volcengine", + DisplayName: "Volcengine", + IconSlug: "bytedance", + Domain: "volcengine.com", + DefaultAPIBase: "https://ark.cn-beijing.volces.com/api/v3", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 69, + CommonModels: []string{ + "doubao-seed-1-6-251015", + "doubao-seed-1-6-flash-250828", + "doubao-seed-1-6-thinking", + }, + httpAPI: true, + }, + "zhipu": { + ID: "zhipu", + DisplayName: "Zhipu AI", + IconSlug: "zhipu", + Domain: "zhipuai.cn", + DefaultAPIBase: "https://open.bigmodel.cn/api/paas/v4", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 68, + CommonModels: []string{"glm-5", "glm-4.7", "glm-4.5-air", "glm-4-flash-250414"}, + Aliases: []string{"glm"}, + httpAPI: true, + }, + "groq": { + ID: "groq", + DisplayName: "Groq", + IconSlug: "groq", + Domain: "groq.com", + DefaultAPIBase: "https://api.groq.com/openai/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 65, + CommonModels: []string{ + "openai/gpt-oss-120b", + "openai/gpt-oss-20b", + "llama-3.3-70b-versatile", + "qwen/qwen3-32b", + }, + httpAPI: true, + }, + "mistral": { + ID: "mistral", + DisplayName: "Mistral AI", + IconSlug: "mistralai", + Domain: "mistral.ai", + DefaultAPIBase: "https://api.mistral.ai/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 64, + CommonModels: []string{ + "mistral-large-latest", + "mistral-medium-3-5", + "mistral-small-latest", + "devstral-latest", + }, + httpAPI: true, + }, + "nvidia": { + ID: "nvidia", + DisplayName: "NVIDIA", + IconSlug: "nvidia", + Domain: "nvidia.com", + DefaultAPIBase: "https://integrate.api.nvidia.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 63, + CommonModels: []string{ + "openai/gpt-oss-120b", + "openai/gpt-oss-20b", + "qwen/qwen3-coder-480b-a35b-instruct", + "qwen/qwen3-next-80b-a3b-thinking", + }, + httpAPI: true, + }, + "cerebras": { + ID: "cerebras", + DisplayName: "Cerebras", + IconSlug: "cerebras", + Domain: "cerebras.ai", + DefaultAPIBase: "https://api.cerebras.ai/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 62, + CommonModels: []string{"gpt-oss-120b", "zai-glm-4.7"}, + httpAPI: true, + }, + "azure": { + ID: "azure", + DisplayName: "Azure OpenAI", + IconSlug: "microsoftazure", + Domain: "azure.com", + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 61, + Aliases: []string{"azure-openai"}, + httpAPI: true, + }, + "bedrock": { + ID: "bedrock", + DisplayName: "AWS Bedrock", + IconSlug: "amazonwebservices", + Domain: "aws.amazon.com", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 60, + }, + "github-copilot": { + ID: "github-copilot", + DisplayName: "GitHub Copilot", + IconSlug: "githubcopilot", + Domain: "github.com", + DefaultAPIBase: "localhost:4321", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + Local: true, + Priority: 55, + Aliases: []string{"copilot"}, + }, + "antigravity": { + ID: "antigravity", + DisplayName: "Google Code Assist", + Domain: "antigravity.google", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + DefaultAuthMethod: "oauth", + AuthMethodLocked: true, + Priority: 54, + Aliases: []string{"google-antigravity"}, + }, + "claude-cli": { + ID: "claude-cli", + DisplayName: "Claude CLI", + IconSlug: "anthropic", + Domain: "anthropic.com", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + Local: true, + Priority: 53, + Aliases: []string{"claudecli"}, + }, + "codex-cli": { + ID: "codex-cli", + DisplayName: "Codex CLI", + IconSlug: "openai", + Domain: "openai.com", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + Local: true, + Priority: 52, + Aliases: []string{"codexcli"}, + }, + "ollama": { + ID: "ollama", + DisplayName: "Ollama", + IconSlug: "ollama", + Domain: "ollama.com", + DefaultAPIBase: "http://localhost:11434/v1", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Local: true, + Priority: 50, + httpAPI: true, + }, + "vllm": { + ID: "vllm", + DisplayName: "VLLM", + Domain: "vllm.ai", + DefaultAPIBase: "http://localhost:8000/v1", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Local: true, + Priority: 49, + httpAPI: true, + }, + "lmstudio": { + ID: "lmstudio", + DisplayName: "LM Studio", + Domain: "lmstudio.ai", + DefaultAPIBase: "http://localhost:1234/v1", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Local: true, + Priority: 48, + httpAPI: true, + }, + "gpt4free": { + ID: "gpt4free", + DisplayName: "GPT4Free", + Domain: "g4f.dev", + DefaultAPIBase: "http://localhost:1337/v1", + EmptyAPIKeyAllowed: true, + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Local: true, + Priority: 47.5, + Aliases: []string{"g4f"}, + httpAPI: true, + }, + "elevenlabs": { + ID: "elevenlabs", + DisplayName: "ElevenLabs ASR", + IconSlug: "elevenlabs", + Domain: "elevenlabs.io", + DefaultAPIBase: "https://api.elevenlabs.io", + CreateAllowed: true, + DefaultModelAllowed: false, + Priority: 47, + httpAPI: true, + }, + "venice": { + ID: "venice", + DisplayName: "Venice AI", + IconSlug: "venice", + Domain: "venice.ai", + DefaultAPIBase: "https://api.venice.ai/api/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 45, + httpAPI: true, + }, + "shengsuanyun": { + ID: "shengsuanyun", + DisplayName: "ShengsuanYun", + Domain: "shengsuanyun.com", + DefaultAPIBase: "https://router.shengsuanyun.com/api/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 44, + httpAPI: true, + }, + "siliconflow": { + ID: "siliconflow", + DisplayName: "SiliconFlow", + Domain: "siliconflow.cn", + DefaultAPIBase: "https://api.siliconflow.cn/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 43.5, + httpAPI: true, + }, + "vivgrid": { + ID: "vivgrid", + DisplayName: "Vivgrid", + Domain: "vivgrid.com", + DefaultAPIBase: "https://api.vivgrid.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 43, + httpAPI: true, + }, + "minimax": { + ID: "minimax", + DisplayName: "MiniMax", + Domain: "minimaxi.com", + DefaultAPIBase: "https://api.minimaxi.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 42, + httpAPI: true, + }, + "longcat": { + ID: "longcat", + DisplayName: "LongCat", + Domain: "longcat.chat", + DefaultAPIBase: "https://api.longcat.chat/openai", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 41, + httpAPI: true, + }, + "modelscope": { + ID: "modelscope", + DisplayName: "ModelScope", + Domain: "modelscope.cn", + DefaultAPIBase: "https://api-inference.modelscope.cn/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 40, + httpAPI: true, + }, + "mimo": { + ID: "mimo", + DisplayName: "Xiaomi MiMo", + IconSlug: "xiaomi", + Domain: "xiaomi.com", + DefaultAPIBase: "https://api.xiaomimimo.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 39, + httpAPI: true, + }, + "avian": { + ID: "avian", + DisplayName: "Avian", + Domain: "avian.io", + DefaultAPIBase: "https://api.avian.io/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 38, + httpAPI: true, + }, + "zai": { + ID: "zai", + DisplayName: "Z.ai", + Domain: "z.ai", + DefaultAPIBase: "https://api.z.ai/api/coding/paas/v4", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 37, + CommonModels: []string{"glm-5", "glm-4.7", "glm-4.5-air", "glm-4-flash-250414"}, + Aliases: []string{"z.ai", "z-ai"}, + httpAPI: true, + }, + "alibaba-coding": { + ID: "alibaba-coding", + DisplayName: "Alibaba Coding Plan", + IconSlug: "alibabacloud", + Domain: "alibabacloud.com", + DefaultAPIBase: "https://coding-intl.dashscope.aliyuncs.com/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 36.5, + CommonModels: []string{"qwen3.6-plus", "kimi-k2.5", "glm-5", "MiniMax-M2.5"}, + Aliases: []string{"coding-plan", "qwen-coding"}, + httpAPI: true, + }, + "alibaba-coding-anthropic": { + ID: "alibaba-coding-anthropic", + DisplayName: "Alibaba Coding Plan (Anthropic)", + IconSlug: "alibabacloud", + Domain: "alibabacloud.com", + DefaultAPIBase: "https://coding-intl.dashscope.aliyuncs.com/apps/anthropic", + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 36.25, + CommonModels: []string{"qwen3.6-plus", "kimi-k2.5", "glm-5", "MiniMax-M2.5"}, + Aliases: []string{"coding-plan-anthropic"}, + httpAPI: true, + }, + "novita": { + ID: "novita", + DisplayName: "Novita AI", + Domain: "novita.ai", + DefaultAPIBase: "https://api.novita.ai/openai", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 36, + httpAPI: true, + }, + "litellm": { + ID: "litellm", + DisplayName: "LiteLLM", + Domain: "litellm.ai", + DefaultAPIBase: "http://localhost:4000/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + SupportsFetch: true, + Priority: 35, + httpAPI: true, + }, + "qwen-us": { + ID: "qwen-us", + DisplayName: "Qwen US", + IconSlug: "alibabacloud", + Domain: "alibabacloud.com", + DefaultAPIBase: "https://dashscope-us.aliyuncs.com/compatible-mode/v1", + CreateAllowed: true, + DefaultModelAllowed: true, + Priority: 34, + CommonModels: []string{"qwen3.6-max-preview", "qwen3.6-plus", "qwen3.6-flash", "qwen3-coder-next"}, + Aliases: []string{"dashscope-us"}, + httpAPI: true, + }, +} + +var normalizedModelProviderAliasesByName = buildModelProviderAliasMap() + +func buildModelProviderAliasMap() map[string]string { + totalAliases := 0 + for _, option := range modelProviderOptionsByName { + totalAliases += len(option.Aliases) + } + + aliases := make(map[string]string, len(modelProviderOptionsByName)+totalAliases) + for provider, option := range modelProviderOptionsByName { + aliases[provider] = provider + for _, alias := range option.Aliases { + normalized := strings.ToLower(strings.TrimSpace(alias)) + if normalized == "" { + continue + } + aliases[normalized] = provider + } + } + return aliases +} + +func modelProviderOptionForName(provider string) (ModelProviderOption, bool) { + normalized := NormalizeProvider(provider) + if normalized == "" { + return ModelProviderOption{}, false + } + option, ok := modelProviderOptionsByName[normalized] + return option, ok +} diff --git a/pkg/providers/tool_schema_transform.go b/pkg/providers/tool_schema_transform.go index 6b6cab7a..e312a99a 100644 --- a/pkg/providers/tool_schema_transform.go +++ b/pkg/providers/tool_schema_transform.go @@ -67,6 +67,29 @@ func (p *toolSchemaStreamingProvider) ChatStream( return streaming.ChatStream(ctx, messages, transformed, model, options, onChunk) } +func (p *toolSchemaStreamingProvider) ChatStreamEvents( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(StreamChunk), +) (*LLMResponse, error) { + streaming, ok := p.delegate.(StreamingEventProvider) + if !ok { + return p.ChatStream(ctx, messages, tools, model, options, func(accumulated string) { + if onChunk != nil { + onChunk(StreamChunk{Content: accumulated}) + } + }) + } + transformed, err := common.TransformToolDefinitions(tools, p.transform) + if err != nil { + return nil, err + } + return streaming.ChatStreamEvents(ctx, messages, transformed, model, options, onChunk) +} + func (p *toolSchemaTransformProvider) SupportsThinking() bool { tc, ok := p.delegate.(ThinkingCapable) return ok && tc.SupportsThinking() diff --git a/pkg/providers/types.go b/pkg/providers/types.go index 23406bc4..c6724f31 100644 --- a/pkg/providers/types.go +++ b/pkg/providers/types.go @@ -11,6 +11,7 @@ type ( ToolCall = protocoltypes.ToolCall FunctionCall = protocoltypes.FunctionCall LLMResponse = protocoltypes.LLMResponse + StreamChunk = protocoltypes.StreamChunk UsageInfo = protocoltypes.UsageInfo Message = protocoltypes.Message ToolDefinition = protocoltypes.ToolDefinition @@ -52,6 +53,17 @@ type StreamingProvider interface { ) (*LLMResponse, error) } +type StreamingEventProvider interface { + ChatStreamEvents( + ctx context.Context, + messages []Message, + tools []ToolDefinition, + model string, + options map[string]any, + onChunk func(StreamChunk), + ) (*LLMResponse, error) +} + // ThinkingCapable is an optional interface for providers that support // extended thinking (e.g. Anthropic). Used by the agent loop to warn // when thinking_level is configured but the active provider cannot use it. diff --git a/pkg/seahorse/schema.go b/pkg/seahorse/schema.go index 5b67fe9e..97e638c9 100644 --- a/pkg/seahorse/schema.go +++ b/pkg/seahorse/schema.go @@ -46,6 +46,7 @@ func runSchema(db *sql.DB) error { conversation_id INTEGER NOT NULL REFERENCES conversations(conversation_id), role TEXT NOT NULL, content TEXT NOT NULL DEFAULT '', + model_name TEXT NOT NULL DEFAULT '', reasoning_content TEXT NOT NULL DEFAULT '', token_count INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT (datetime('now')) @@ -162,6 +163,9 @@ func runSchema(db *sql.DB) error { if err := ensureMessagesReasoningContentColumn(db); err != nil { return err } + if err := ensureMessagesModelNameColumn(db); err != nil { + return err + } return nil } @@ -180,6 +184,21 @@ func ensureMessagesReasoningContentColumn(db *sql.DB) error { return nil } +func ensureMessagesModelNameColumn(db *sql.DB) error { + hasColumn, err := tableHasColumn(db, "messages", "model_name") + if err != nil { + return fmt.Errorf("check messages.model_name: %w", err) + } + if hasColumn { + return nil + } + + if _, err := db.Exec(`ALTER TABLE messages ADD COLUMN model_name TEXT NOT NULL DEFAULT ''`); err != nil { + return fmt.Errorf("add messages.model_name: %w", err) + } + return nil +} + func tableHasColumn(db *sql.DB, tableName, columnName string) (bool, error) { rows, err := db.Query(fmt.Sprintf(`PRAGMA table_info(%s)`, tableName)) if err != nil { diff --git a/pkg/seahorse/schema_test.go b/pkg/seahorse/schema_test.go index 943b742b..4618eeff 100644 --- a/pkg/seahorse/schema_test.go +++ b/pkg/seahorse/schema_test.go @@ -138,6 +138,37 @@ func TestRunSchemaAddsMessagesReasoningContentColumn(t *testing.T) { } } +func TestRunSchemaAddsMessagesModelNameColumn(t *testing.T) { + db := openTestDB(t) + + _, err := db.Exec(`CREATE TABLE messages ( + message_id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL DEFAULT '', + reasoning_content TEXT NOT NULL DEFAULT '', + token_count INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL DEFAULT (datetime('now')) + )`) + if err != nil { + t.Fatalf("create legacy messages table: %v", err) + } + + err = runSchema(db) + if err != nil { + t.Fatalf("runSchema: %v", err) + } + + var count int + err = db.QueryRow(`SELECT count(*) FROM pragma_table_info('messages') WHERE name = 'model_name'`).Scan(&count) + if err != nil { + t.Fatalf("query pragma_table_info: %v", err) + } + if count != 1 { + t.Fatalf("model_name column count = %d, want 1", count) + } +} + func TestMigrationConversationUnique(t *testing.T) { db := openTestDB(t) if err := runSchema(db); err != nil { diff --git a/pkg/seahorse/short_engine.go b/pkg/seahorse/short_engine.go index 0a817561..d275da51 100644 --- a/pkg/seahorse/short_engine.go +++ b/pkg/seahorse/short_engine.go @@ -258,6 +258,7 @@ func (e *Engine) Ingest(ctx context.Context, sessionKey string, messages []Messa conv.ConversationID, msg.Role, msg.Parts, + msg.ModelName, msg.ReasoningContent, msg.TokenCount, ) @@ -267,6 +268,7 @@ func (e *Engine) Ingest(ctx context.Context, sessionKey string, messages []Messa conv.ConversationID, msg.Role, msg.Content, + msg.ModelName, msg.ReasoningContent, msg.TokenCount, ) @@ -431,6 +433,31 @@ func (e *Engine) Bootstrap(ctx context.Context, sessionKey string, messages []Me return fmt.Errorf("bootstrap: get messages: %w", err) } + // Migration repair path: old SeaHorse rows may be missing reasoning_content + // even though the canonical JSONL history already has it. Backfill those + // rows in place so we do not treat this as edited history and leave stale + // summaries/context behind after a partial raw-message rebuild. + repairedReasoning, err := e.repairBootstrapReasoningContent(ctx, dbMsgs, messages) + if err != nil { + return fmt.Errorf("bootstrap: repair reasoning_content: %w", err) + } + repairedModelName, err := e.repairBootstrapModelName(ctx, dbMsgs, messages) + if err != nil { + return fmt.Errorf("bootstrap: repair model_name: %w", err) + } + if (repairedReasoning || repairedModelName) && len(dbMsgs) == len(messages) { + matched := true + for i := range messages { + if !messageMatches(dbMsgs[i], messages[i]) { + matched = false + break + } + } + if matched { + return nil + } + } + // Fast path: DB has same count and exact match → no-op if len(dbMsgs) == len(messages) { matched := true @@ -445,16 +472,6 @@ func (e *Engine) Bootstrap(ctx context.Context, sessionKey string, messages []Me } } - // Migration repair path: old SeaHorse rows may be missing reasoning_content - // even though the canonical JSONL history already has it. Backfill those - // rows in place so we do not treat this as edited history and leave stale - // summaries/context behind after a partial raw-message rebuild. - if repaired, err := e.repairBootstrapReasoningContent(ctx, dbMsgs, messages); err != nil { - return fmt.Errorf("bootstrap: repair reasoning_content: %w", err) - } else if repaired && len(dbMsgs) == len(messages) { - return nil - } - // Find longest matching prefix from the start anchor := -1 compareLen := min(len(dbMsgs), len(messages)) @@ -465,14 +482,16 @@ func (e *Engine) Bootstrap(ctx context.Context, sessionKey string, messages []Me } else { // Mismatch detected - log details and rebuild logger.InfoCF("seahorse", "bootstrap: mismatch detected", map[string]any{ - "conv_id": conv.ConversationID, - "index": i, - "db_role": dbMsgs[i].Role, - "db_content": truncate(dbMsgs[i].Content, 50), - "db_parts": len(dbMsgs[i].Parts), - "msg_role": messages[i].Role, - "msg_content": truncate(messages[i].Content, 50), - "msg_parts": len(messages[i].Parts), + "conv_id": conv.ConversationID, + "index": i, + "db_role": dbMsgs[i].Role, + "db_content": truncate(dbMsgs[i].Content, 50), + "db_parts": len(dbMsgs[i].Parts), + "db_model_name": dbMsgs[i].ModelName, + "msg_role": messages[i].Role, + "msg_content": truncate(messages[i].Content, 50), + "msg_parts": len(messages[i].Parts), + "msg_model_name": messages[i].ModelName, }) break } @@ -559,7 +578,7 @@ func (e *Engine) repairBootstrapReasoningContent(ctx context.Context, dbMsgs, me } for i := range overlap { - if !messageMatchesIgnoringReasoning(dbMsgs[i], messages[i]) { + if !messageMatchesIgnoringReasoningAndModelName(dbMsgs[i], messages[i]) { return false, nil } if dbMsgs[i].ReasoningContent == messages[i].ReasoningContent { @@ -596,6 +615,57 @@ func (e *Engine) repairBootstrapReasoningContent(ctx context.Context, dbMsgs, me return true, nil } +func (e *Engine) repairBootstrapModelName(ctx context.Context, dbMsgs, messages []Message) (bool, error) { + if len(dbMsgs) == 0 || len(messages) == 0 { + return false, nil + } + + overlap := min(len(messages), len(dbMsgs)) + + var updates []struct { + index int + messageID int64 + modelName string + } + + for i := range overlap { + if !messageMatchesIgnoringReasoningAndModelName(dbMsgs[i], messages[i]) { + return false, nil + } + if dbMsgs[i].ModelName == messages[i].ModelName { + continue + } + if messages[i].ModelName == "" { + return false, nil + } + updates = append(updates, struct { + index int + messageID int64 + modelName string + }{ + index: i, + messageID: dbMsgs[i].ID, + modelName: messages[i].ModelName, + }) + } + + if len(updates) == 0 { + return false, nil + } + + for _, update := range updates { + if err := e.store.UpdateMessageModelName(ctx, update.messageID, update.modelName); err != nil { + return false, err + } + dbMsgs[update.index].ModelName = update.modelName + } + + logger.InfoCF("seahorse", "bootstrap: repaired missing model_name", map[string]any{ + "messages": len(updates), + }) + return true, nil +} + // truncate shortens a string for logging. func truncate(s string, maxLen int) string { if len(s) <= maxLen { @@ -610,13 +680,20 @@ func truncate(s string, maxLen int) string { // For messages with Parts (tool_use, tool_result), compare Parts instead of Content // because structured messages are matched by their parts payload. func messageMatches(a, b Message) bool { - if a.Role != b.Role || a.ReasoningContent != b.ReasoningContent { + if a.Role != b.Role || a.ReasoningContent != b.ReasoningContent || a.ModelName != b.ModelName { return false } return messageMatchesIgnoringReasoning(a, b) } func messageMatchesIgnoringReasoning(a, b Message) bool { + if a.ModelName != b.ModelName { + return false + } + return messageMatchesIgnoringReasoningAndModelName(a, b) +} + +func messageMatchesIgnoringReasoningAndModelName(a, b Message) bool { if a.Role != b.Role { return false } diff --git a/pkg/seahorse/short_engine_test.go b/pkg/seahorse/short_engine_test.go index 2a5c6c5d..337416f6 100644 --- a/pkg/seahorse/short_engine_test.go +++ b/pkg/seahorse/short_engine_test.go @@ -25,6 +25,43 @@ func newTestEngine(t *testing.T) *Engine { } } +func prepareBootstrapRepairConversation( + t *testing.T, + eng *Engine, + ctx context.Context, + sessionKey string, +) (*Conversation, []Message) { + t.Helper() + + conv, err := eng.store.GetOrCreateConversation(ctx, sessionKey) + if err != nil { + t.Fatalf("GetOrCreateConversation: %v", err) + } + + userMsg, err := eng.store.AddMessage(ctx, conv.ConversationID, "user", "hello", 3) + if err != nil { + t.Fatalf("AddMessage user: %v", err) + } + + assistantMsg, err := eng.store.AddMessage(ctx, conv.ConversationID, "assistant", "world", 3) + if err != nil { + t.Fatalf("AddMessage assistant: %v", err) + } + + if err := eng.store.AppendContextMessages( + ctx, + conv.ConversationID, + []int64{userMsg.ID, assistantMsg.ID}, + ); err != nil { + t.Fatalf("AppendContextMessages: %v", err) + } + + return conv, []Message{ + {Role: "user", Content: "hello", TokenCount: 3}, + {Role: "assistant", Content: "world", TokenCount: 3}, + } +} + // --- compileSessionPattern --- func TestCompileSessionPattern(t *testing.T) { @@ -328,6 +365,7 @@ func TestEngineIngestPreservesReasoningContent(t *testing.T) { { Role: "assistant", Content: "world", + ModelName: "gpt-5.4-mini", ReasoningContent: "let me think this through", TokenCount: 4, }, @@ -353,6 +391,9 @@ func TestEngineIngestPreservesReasoningContent(t *testing.T) { "let me think this through", ) } + if stored[0].ModelName != "gpt-5.4-mini" { + t.Errorf("stored[0].ModelName = %q, want %q", stored[0].ModelName, "gpt-5.4-mini") + } result, err := eng.Assemble(ctx, "agent:reasoning", AssembleInput{Budget: 1000}) if err != nil { @@ -368,6 +409,140 @@ func TestEngineIngestPreservesReasoningContent(t *testing.T) { "let me think this through", ) } + if result.Messages[0].ModelName != "gpt-5.4-mini" { + t.Errorf("assembled model_name = %q, want %q", result.Messages[0].ModelName, "gpt-5.4-mini") + } +} + +func TestBootstrapRepairsMissingModelName(t *testing.T) { + eng := newTestEngine(t) + ctx := context.Background() + sessionKey := "agent:repair-model-name" + conv, msgs := prepareBootstrapRepairConversation(t, eng, ctx, sessionKey) + msgs[1].ModelName = "gpt-5.4" + + err := eng.Bootstrap(ctx, sessionKey, msgs) + if err != nil { + t.Fatalf("Bootstrap: %v", err) + } + + stored, err := eng.store.GetMessages(ctx, conv.ConversationID, 10, 0) + if err != nil { + t.Fatalf("GetMessages: %v", err) + } + if len(stored) != 2 { + t.Fatalf("stored messages = %d, want 2", len(stored)) + } + if stored[1].ModelName != "gpt-5.4" { + t.Fatalf("stored[1].ModelName = %q, want %q", stored[1].ModelName, "gpt-5.4") + } +} + +func TestBootstrapRepairsReasoningContentAndModelNameTogether(t *testing.T) { + eng := newTestEngine(t) + ctx := context.Background() + sessionKey := "agent:repair-both-fields" + + conv, err := eng.store.GetOrCreateConversation(ctx, sessionKey) + if err != nil { + t.Fatalf("GetOrCreateConversation: %v", err) + } + + userMsg, err := eng.store.AddMessage(ctx, conv.ConversationID, "user", "hello", 3) + if err != nil { + t.Fatalf("AddMessage user: %v", err) + } + + assistantMsg, err := eng.store.AddMessage(ctx, conv.ConversationID, "assistant", "world", 3) + if err != nil { + t.Fatalf("AddMessage assistant: %v", err) + } + + err = eng.store.AppendContextMessages(ctx, conv.ConversationID, []int64{userMsg.ID, assistantMsg.ID}) + if err != nil { + t.Fatalf("AppendContextMessages: %v", err) + } + + err = eng.Bootstrap(ctx, sessionKey, []Message{ + {Role: "user", Content: "hello", TokenCount: 3}, + { + Role: "assistant", + Content: "world", + ModelName: "gpt-5.4", + ReasoningContent: "let me think this through", + TokenCount: 3, + }, + }) + if err != nil { + t.Fatalf("Bootstrap: %v", err) + } + + stored, err := eng.store.GetMessages(ctx, conv.ConversationID, 10, 0) + if err != nil { + t.Fatalf("GetMessages: %v", err) + } + if len(stored) != 2 { + t.Fatalf("stored messages = %d, want 2", len(stored)) + } + if stored[1].ReasoningContent != "let me think this through" { + t.Fatalf("stored[1].ReasoningContent = %q, want %q", stored[1].ReasoningContent, "let me think this through") + } + if stored[1].ModelName != "gpt-5.4" { + t.Fatalf("stored[1].ModelName = %q, want %q", stored[1].ModelName, "gpt-5.4") + } +} + +func TestBootstrapRepairsIncorrectNonEmptyModelName(t *testing.T) { + eng := newTestEngine(t) + ctx := context.Background() + sessionKey := "agent:repair-wrong-model-name" + + conv, err := eng.store.GetOrCreateConversation(ctx, sessionKey) + if err != nil { + t.Fatalf("GetOrCreateConversation: %v", err) + } + + userMsg, err := eng.store.AddMessage(ctx, conv.ConversationID, "user", "hello", 3) + if err != nil { + t.Fatalf("AddMessage user: %v", err) + } + + assistantMsg, err := eng.store.AddMessageWithReasoning( + ctx, + conv.ConversationID, + "assistant", + "world", + "wrong-model", + "", + 3, + ) + if err != nil { + t.Fatalf("AddMessageWithReasoning assistant: %v", err) + } + + err = eng.store.AppendContextMessages(ctx, conv.ConversationID, []int64{userMsg.ID, assistantMsg.ID}) + if err != nil { + t.Fatalf("AppendContextMessages: %v", err) + } + + err = eng.Bootstrap(ctx, sessionKey, []Message{ + {Role: "user", Content: "hello", TokenCount: 3}, + {Role: "assistant", Content: "world", ModelName: "gpt-5.4", TokenCount: 3}, + }) + if err != nil { + t.Fatalf("Bootstrap: %v", err) + } + + stored, err := eng.store.GetMessages(ctx, conv.ConversationID, 10, 0) + if err != nil { + t.Fatalf("GetMessages: %v", err) + } + if len(stored) != 2 { + t.Fatalf("stored messages = %d, want 2", len(stored)) + } + if stored[1].ModelName != "gpt-5.4" { + t.Fatalf("stored[1].ModelName = %q, want %q", stored[1].ModelName, "gpt-5.4") + } } func TestEngineIngestWithPartsPreservesReasoningContent(t *testing.T) { @@ -620,35 +795,10 @@ func TestBootstrapRepairsMissingReasoningContent(t *testing.T) { eng := newTestEngine(t) ctx := context.Background() sessionKey := "agent:repair-reasoning" + conv, msgs := prepareBootstrapRepairConversation(t, eng, ctx, sessionKey) + msgs[1].ReasoningContent = "let me think this through" - conv, err := eng.store.GetOrCreateConversation(ctx, sessionKey) - if err != nil { - t.Fatalf("GetOrCreateConversation: %v", err) - } - - userMsg, err := eng.store.AddMessage(ctx, conv.ConversationID, "user", "hello", 3) - if err != nil { - t.Fatalf("AddMessage user: %v", err) - } - - assistantMsg, err := eng.store.AddMessage(ctx, conv.ConversationID, "assistant", "world", 3) - if err != nil { - t.Fatalf("AddMessage assistant: %v", err) - } - - err = eng.store.AppendContextMessages( - ctx, - conv.ConversationID, - []int64{userMsg.ID, assistantMsg.ID}, - ) - if err != nil { - t.Fatalf("AppendContextMessages: %v", err) - } - - err = eng.Bootstrap(ctx, sessionKey, []Message{ - {Role: "user", Content: "hello", TokenCount: 3}, - {Role: "assistant", Content: "world", ReasoningContent: "let me think this through", TokenCount: 3}, - }) + err := eng.Bootstrap(ctx, sessionKey, msgs) if err != nil { t.Fatalf("Bootstrap: %v", err) } diff --git a/pkg/seahorse/store.go b/pkg/seahorse/store.go index 0edbbd12..b5e32e89 100644 --- a/pkg/seahorse/store.go +++ b/pkg/seahorse/store.go @@ -162,19 +162,25 @@ func (s *Store) getMessageTimeRange(ctx context.Context, convID int64) (time.Tim // AddMessage appends a message to a conversation. func (s *Store) AddMessage(ctx context.Context, convID int64, role, content string, tokenCount int) (*Message, error) { - return s.AddMessageWithReasoning(ctx, convID, role, content, "", tokenCount) + return s.AddMessageWithReasoning(ctx, convID, role, content, "", "", tokenCount) } // AddMessageWithReasoning appends a message with reasoning content to a conversation. func (s *Store) AddMessageWithReasoning( ctx context.Context, convID int64, - role, content, reasoningContent string, + role, content, modelName, reasoningContent string, tokenCount int, ) (*Message, error) { - result, err := s.db.ExecContext(ctx, - "INSERT INTO messages (conversation_id, role, content, reasoning_content, token_count) VALUES (?, ?, ?, ?, ?)", - convID, role, content, reasoningContent, tokenCount, + result, err := s.db.ExecContext( + ctx, + "INSERT INTO messages (conversation_id, role, content, model_name, reasoning_content, token_count) VALUES (?, ?, ?, ?, ?, ?)", + convID, + role, + content, + modelName, + reasoningContent, + tokenCount, ) if err != nil { return nil, fmt.Errorf("add message: %w", err) @@ -185,6 +191,7 @@ func (s *Store) AddMessageWithReasoning( ConversationID: convID, Role: role, Content: content, + ModelName: modelName, ReasoningContent: reasoningContent, TokenCount: tokenCount, }, nil @@ -224,7 +231,7 @@ func (s *Store) AddMessageWithParts( parts []MessagePart, tokenCount int, ) (*Message, error) { - return s.AddMessageWithPartsAndReasoning(ctx, convID, role, parts, "", tokenCount) + return s.AddMessageWithPartsAndReasoning(ctx, convID, role, parts, "", "", tokenCount) } // AddMessageWithPartsAndReasoning adds a message with structured parts and reasoning content. @@ -233,6 +240,7 @@ func (s *Store) AddMessageWithPartsAndReasoning( convID int64, role string, parts []MessagePart, + modelName string, reasoningContent string, tokenCount int, ) (*Message, error) { @@ -245,9 +253,15 @@ func (s *Store) AddMessageWithPartsAndReasoning( // Derive readable content from Parts for FTS5 indexing and summary formatting readableContent := partsToReadableContent(parts) - result, err := tx.ExecContext(ctx, - "INSERT INTO messages (conversation_id, role, content, reasoning_content, token_count) VALUES (?, ?, ?, ?, ?)", - convID, role, readableContent, reasoningContent, tokenCount, + result, err := tx.ExecContext( + ctx, + "INSERT INTO messages (conversation_id, role, content, model_name, reasoning_content, token_count) VALUES (?, ?, ?, ?, ?, ?)", + convID, + role, + readableContent, + modelName, + reasoningContent, + tokenCount, ) if err != nil { return nil, fmt.Errorf("add message: %w", err) @@ -282,6 +296,7 @@ func (s *Store) AddMessageWithPartsAndReasoning( ID: msgID, ConversationID: convID, Role: role, + ModelName: modelName, ReasoningContent: reasoningContent, TokenCount: tokenCount, Parts: make([]MessagePart, len(parts)), @@ -295,7 +310,7 @@ func (s *Store) AddMessageWithPartsAndReasoning( // GetMessages retrieves messages for a conversation. func (s *Store) GetMessages(ctx context.Context, convID int64, limit int, beforeID int64) ([]Message, error) { - query := "SELECT message_id, conversation_id, role, content, reasoning_content, token_count, created_at FROM messages WHERE conversation_id = ?" + query := "SELECT message_id, conversation_id, role, content, model_name, reasoning_content, token_count, created_at FROM messages WHERE conversation_id = ?" args := []any{convID} if beforeID > 0 { query += " AND message_id < ?" @@ -322,6 +337,7 @@ func (s *Store) GetMessages(ctx context.Context, convID int64, limit int, before &msg.ConversationID, &msg.Role, &msg.Content, + &msg.ModelName, &msg.ReasoningContent, &msg.TokenCount, &createdAt, @@ -362,9 +378,9 @@ func (s *Store) GetMessageByID(ctx context.Context, messageID int64) (*Message, var createdAt string err := s.db.QueryRowContext( ctx, - "SELECT message_id, conversation_id, role, content, reasoning_content, token_count, created_at FROM messages WHERE message_id = ?", + "SELECT message_id, conversation_id, role, content, model_name, reasoning_content, token_count, created_at FROM messages WHERE message_id = ?", messageID, - ).Scan(&msg.ID, &msg.ConversationID, &msg.Role, &msg.Content, &msg.ReasoningContent, &msg.TokenCount, &createdAt) + ).Scan(&msg.ID, &msg.ConversationID, &msg.Role, &msg.Content, &msg.ModelName, &msg.ReasoningContent, &msg.TokenCount, &createdAt) if err == sql.ErrNoRows { return nil, fmt.Errorf("message %d not found", messageID) } @@ -398,6 +414,27 @@ func (s *Store) UpdateMessageReasoningContent(ctx context.Context, messageID int return nil } +func (s *Store) UpdateMessageModelName(ctx context.Context, messageID int64, modelName string) error { + result, err := s.db.ExecContext( + ctx, + "UPDATE messages SET model_name = ? WHERE message_id = ?", + modelName, + messageID, + ) + if err != nil { + return fmt.Errorf("update message model_name: %w", err) + } + + rowsAffected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("update message model_name rows affected: %w", err) + } + if rowsAffected == 0 { + return fmt.Errorf("message %d not found", messageID) + } + return nil +} + func (s *Store) loadMessageParts(ctx context.Context, msgID int64) ([]MessagePart, error) { rows, err := s.db.QueryContext(ctx, `SELECT part_id, message_id, type, text, name, arguments, tool_call_id, media_uri, mime_type @@ -581,8 +618,9 @@ func (s *Store) LinkSummaryToMessages(ctx context.Context, summaryID string, mes // GetSummarySourceMessages retrieves source messages for a summary. func (s *Store) GetSummarySourceMessages(ctx context.Context, summaryID string) ([]Message, error) { - rows, err := s.db.QueryContext(ctx, - `SELECT m.message_id, m.conversation_id, m.role, m.content, m.reasoning_content, m.token_count, m.created_at + rows, err := s.db.QueryContext( + ctx, + `SELECT m.message_id, m.conversation_id, m.role, m.content, m.model_name, m.reasoning_content, m.token_count, m.created_at FROM summary_messages sm JOIN messages m ON m.message_id = sm.message_id WHERE sm.summary_id = ? @@ -603,6 +641,7 @@ func (s *Store) GetSummarySourceMessages(ctx context.Context, summaryID string) &msg.ConversationID, &msg.Role, &msg.Content, + &msg.ModelName, &msg.ReasoningContent, &msg.TokenCount, &createdAt, diff --git a/pkg/seahorse/store_test.go b/pkg/seahorse/store_test.go index 67bed1c1..4ed2bb3b 100644 --- a/pkg/seahorse/store_test.go +++ b/pkg/seahorse/store_test.go @@ -210,6 +210,7 @@ func TestStoreAddAndGetMessagesWithReasoningContent(t *testing.T) { conv.ConversationID, "assistant", "hello world", + "gpt-5.4-mini", "let me think", 5, ) @@ -219,6 +220,9 @@ func TestStoreAddAndGetMessagesWithReasoningContent(t *testing.T) { if msg.ReasoningContent != "let me think" { t.Fatalf("ReasoningContent = %q, want %q", msg.ReasoningContent, "let me think") } + if msg.ModelName != "gpt-5.4-mini" { + t.Fatalf("ModelName = %q, want %q", msg.ModelName, "gpt-5.4-mini") + } msgs, err := s.GetMessages(ctx, conv.ConversationID, 10, 0) if err != nil { @@ -230,6 +234,9 @@ func TestStoreAddAndGetMessagesWithReasoningContent(t *testing.T) { if msgs[0].ReasoningContent != "let me think" { t.Errorf("ReasoningContent = %q, want %q", msgs[0].ReasoningContent, "let me think") } + if msgs[0].ModelName != "gpt-5.4-mini" { + t.Errorf("ModelName = %q, want %q", msgs[0].ModelName, "gpt-5.4-mini") + } found, err := s.GetMessageByID(ctx, msg.ID) if err != nil { @@ -238,6 +245,9 @@ func TestStoreAddAndGetMessagesWithReasoningContent(t *testing.T) { if found.ReasoningContent != "let me think" { t.Errorf("GetMessageByID ReasoningContent = %q, want %q", found.ReasoningContent, "let me think") } + if found.ModelName != "gpt-5.4-mini" { + t.Errorf("GetMessageByID ModelName = %q, want %q", found.ModelName, "gpt-5.4-mini") + } } func TestStoreAddMessageWithParts(t *testing.T) { @@ -288,6 +298,7 @@ func TestStoreAddMessageWithPartsAndReasoningContent(t *testing.T) { conv.ConversationID, "assistant", parts, + "gpt-5.4", "need to inspect the file first", 10, ) @@ -309,6 +320,9 @@ func TestStoreAddMessageWithPartsAndReasoningContent(t *testing.T) { "need to inspect the file first", ) } + if msgs[0].ModelName != "gpt-5.4" { + t.Errorf("ModelName = %q, want %q", msgs[0].ModelName, "gpt-5.4") + } } func TestStoreGetMessageCount(t *testing.T) { diff --git a/pkg/seahorse/types.go b/pkg/seahorse/types.go index 2bc7f931..af323d2b 100644 --- a/pkg/seahorse/types.go +++ b/pkg/seahorse/types.go @@ -22,6 +22,7 @@ type Message struct { ConversationID int64 `json:"conversationId"` Role string `json:"role"` Content string `json:"content"` + ModelName string `json:"modelName,omitempty"` ReasoningContent string `json:"reasoningContent,omitempty"` TokenCount int `json:"tokenCount"` CreatedAt time.Time `json:"createdAt"` @@ -135,6 +136,7 @@ func EstimateMessageTokens(msg Message) int { pm := providers.Message{ Role: msg.Role, Content: msg.Content, + ModelName: msg.ModelName, ReasoningContent: msg.ReasoningContent, } diff --git a/pkg/session/jsonl_backend_test.go b/pkg/session/jsonl_backend_test.go index 0b79ad84..6f67109d 100644 --- a/pkg/session/jsonl_backend_test.go +++ b/pkg/session/jsonl_backend_test.go @@ -66,6 +66,25 @@ func TestJSONLBackend_AddFullMessage(t *testing.T) { } } +func TestJSONLBackend_AddFullMessage_PreservesModelName(t *testing.T) { + b := newBackend(t) + + msg := providers.Message{ + Role: "assistant", + Content: "done", + ModelName: "gpt-5.4-mini", + } + b.AddFullMessage("s1", msg) + + history := b.GetHistory("s1") + if len(history) != 1 { + t.Fatalf("got %d, want 1", len(history)) + } + if history[0].ModelName != "gpt-5.4-mini" { + t.Fatalf("ModelName = %q, want %q", history[0].ModelName, "gpt-5.4-mini") + } +} + func TestJSONLBackend_Summary(t *testing.T) { b := newBackend(t) diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 1d6fa310..7ca03744 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -87,8 +87,13 @@ func (sm *SessionManager) AddFullMessage(sessionKey string, msg providers.Messag sm.sessions[sessionKey] = session } + now := time.Now() + if msg.CreatedAt == nil { + msg.CreatedAt = &now + } + session.Messages = append(session.Messages, msg) - session.Updated = time.Now() + session.Updated = now } func (sm *SessionManager) GetHistory(key string) []providers.Message { @@ -300,7 +305,13 @@ func (sm *SessionManager) SetHistory(key string, history []providers.Message) { // from the caller's slice. msgs := make([]providers.Message, len(history)) copy(msgs, history) + now := time.Now() + for i := range msgs { + if msgs[i].CreatedAt == nil { + msgs[i].CreatedAt = &now + } + } session.Messages = msgs - session.Updated = time.Now() + session.Updated = now } } diff --git a/web/backend/api/channels.go b/web/backend/api/channels.go index e77b11f8..b98c0557 100644 --- a/web/backend/api/channels.go +++ b/web/backend/api/channels.go @@ -3,6 +3,7 @@ package api import ( "encoding/json" "net/http" + "reflect" "github.com/sipeed/picoclaw/pkg/config" ) @@ -170,6 +171,42 @@ func addChannelCommonConfig(settings map[string]any, bc *config.Channel) { if bc.Placeholder.Enabled || len(bc.Placeholder.Text) > 0 { settings["placeholder"] = bc.Placeholder } + if _, exists := settings["streaming"]; !exists { + if streaming, ok := channelStreamingConfig(bc); ok { + if !streaming.IsZero() { + settings["streaming"] = streaming + } + } + } +} + +func channelStreamingConfig(bc *config.Channel) (config.StreamingConfig, bool) { + if bc == nil { + return config.StreamingConfig{}, false + } + + decoded, err := bc.GetDecoded() + if err != nil || decoded == nil { + return config.StreamingConfig{}, false + } + + value := reflect.ValueOf(decoded) + if value.Kind() == reflect.Pointer { + if value.IsNil() { + return config.StreamingConfig{}, false + } + value = value.Elem() + } + if value.Kind() != reflect.Struct { + return config.StreamingConfig{}, false + } + + field := value.FieldByName("Streaming") + if !field.IsValid() || !field.CanInterface() { + return config.StreamingConfig{}, false + } + streaming, ok := field.Interface().(config.StreamingConfig) + return streaming, ok } func detectConfiguredSecrets(settings config.RawNode, channelName string) []string { diff --git a/web/backend/api/channels_test.go b/web/backend/api/channels_test.go index 0208af8e..a70c9ba4 100644 --- a/web/backend/api/channels_test.go +++ b/web/backend/api/channels_test.go @@ -147,6 +147,96 @@ func TestHandleGetChannelConfig_ReturnsCommonFieldsWhenSettingsEmpty(t *testing. } } +func TestHandleGetChannelConfig_OmitsUnconfiguredStreaming(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + req := httptest.NewRequest(http.MethodGet, "/api/channels/telegram/config", nil) + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf( + "GET /api/channels/telegram/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) + } + + var resp struct { + Config map[string]any `json:"config"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("json.Unmarshal() error = %v", err) + } + if _, ok := resp.Config["streaming"]; ok { + t.Fatalf("config.streaming = %#v, want omitted when not configured", resp.Config["streaming"]) + } +} + +func TestHandleGetChannelConfig_ReturnsConfiguredStreaming(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + pico := cfg.Channels.Get(config.ChannelPico) + if pico == nil { + t.Fatal("missing pico channel") + } + pico.Settings = config.RawNode(`{"streaming":{"enabled":true,"throttle_seconds":2,"min_growth_chars":80}}`) + if err := config.InitChannelList(cfg.Channels); err != nil { + t.Fatalf("InitChannelList() error = %v", err) + } + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + req := httptest.NewRequest(http.MethodGet, "/api/channels/pico/config", nil) + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf( + "GET /api/channels/pico/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) + } + + var resp struct { + Config map[string]any `json:"config"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("json.Unmarshal() error = %v", err) + } + streaming, ok := resp.Config["streaming"].(map[string]any) + if !ok { + t.Fatalf("config.streaming = %#v, want object", resp.Config["streaming"]) + } + if got := streaming["enabled"]; got != true { + t.Fatalf("config.streaming.enabled = %#v, want true", got) + } + if got := streaming["throttle_seconds"]; got != float64(2) { + t.Fatalf("config.streaming.throttle_seconds = %#v, want 2", got) + } + if got := streaming["min_growth_chars"]; got != float64(80) { + t.Fatalf("config.streaming.min_growth_chars = %#v, want 80", got) + } +} + func TestHandleGetChannelConfig_ReturnsDefaultShapeForMissingChannel(t *testing.T) { configPath, cleanup := setupOAuthTestEnv(t) defer cleanup() diff --git a/web/backend/api/config.go b/web/backend/api/config.go index afcd3f74..b3795804 100644 --- a/web/backend/api/config.go +++ b/web/backend/api/config.go @@ -18,6 +18,7 @@ func (h *Handler) registerConfigRoutes(mux *http.ServeMux) { mux.HandleFunc("GET /api/config", h.handleGetConfig) mux.HandleFunc("PUT /api/config", h.handleUpdateConfig) mux.HandleFunc("PATCH /api/config", h.handlePatchConfig) + mux.HandleFunc("POST /api/config/reset", h.handleResetConfig) mux.HandleFunc("POST /api/config/test-command-patterns", h.handleTestCommandPatterns) } @@ -211,6 +212,32 @@ func (h *Handler) handlePatchConfig(w http.ResponseWriter, r *http.Request) { json.NewEncoder(w).Encode(map[string]string{"status": "ok"}) } +// handleResetConfig resets the configuration to factory defaults. +// API keys and security credentials are preserved. +// +// POST /api/config/reset +func (h *Handler) handleResetConfig(w http.ResponseWriter, r *http.Request) { + if err := config.ResetToDefaults(h.configPath); err != nil { + http.Error(w, fmt.Sprintf("Failed to reset config: %v", err), http.StatusInternalServerError) + return + } + + h.applyRuntimeLogLevel() + logger.Infof("configuration reset to factory defaults") + + // Restart gateway if running + status := h.gatewayStatusData() + gatewayStatus, _ := status["gateway_status"].(string) + if gatewayStatus == "running" { + if _, err := h.RestartGateway(); err != nil { + logger.ErrorF("failed to restart gateway after config reset", map[string]any{"error": err.Error()}) + } + } + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]string{"status": "ok"}) +} + // handleTestCommandPatterns tests a command against whitelist and blacklist patterns. // // POST /api/config/test-command-patterns @@ -285,11 +312,28 @@ func validateConfig(cfg *config.Config) []string { errs = append(errs, err.Error()) } + if err := cfg.ValidateTurnProfile(); err != nil { + errs = append(errs, err.Error()) + } + // Gateway port range if cfg.Gateway.Port != 0 && (cfg.Gateway.Port < 1 || cfg.Gateway.Port > 65535) { errs = append(errs, fmt.Sprintf("gateway.port %d is out of valid range (1-65535)", cfg.Gateway.Port)) } + for name, bc := range cfg.Channels { + streaming, ok := channelStreamingConfig(bc) + if !ok { + continue + } + if streaming.ThrottleSeconds < 0 { + errs = append(errs, fmt.Sprintf("channel %q streaming.throttle_seconds must be >= 0", name)) + } + if streaming.MinGrowthChars < 0 { + errs = append(errs, fmt.Sprintf("channel %q streaming.min_growth_chars must be >= 0", name)) + } + } + // Pico channel: token required when enabled { bc := cfg.Channels.GetByType(config.ChannelPico) diff --git a/web/backend/api/config_test.go b/web/backend/api/config_test.go index 8377c2ec..8a23b7a3 100644 --- a/web/backend/api/config_test.go +++ b/web/backend/api/config_test.go @@ -13,6 +13,95 @@ import ( "github.com/sipeed/picoclaw/pkg/logger" ) +func TestHandlePatchConfig_PreservesTurnProfile(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.Agents.Defaults.TurnProfile = config.TurnProfileConfig{ + Enabled: true, + History: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + SystemPrompt: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Skills: config.TurnProfileBlock{Mode: config.TurnProfileModeOff}, + Tools: config.TurnProfileBlock{ + Mode: config.TurnProfileModeCustom, + Allow: []string{"web_search", "web_fetch"}, + }, + } + if saveErr := config.SaveConfig(configPath, cfg); saveErr != nil { + t.Fatalf("SaveConfig() error = %v", saveErr) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + req := httptest.NewRequest(http.MethodPatch, "/api/config", bytes.NewBufferString(`{ + "agents": { + "defaults": { + "max_tokens": 1234 + } + } + }`)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + updated, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig(updated) error = %v", err) + } + profile := updated.Agents.Defaults.TurnProfile + if profile.Tools.Mode != config.TurnProfileModeCustom || + strings.Join(profile.Tools.Allow, ",") != "web_search,web_fetch" { + t.Fatalf("profile tools = %#v, want custom web_search/web_fetch", profile.Tools) + } +} + +func TestHandlePatchConfig_RejectsInvalidTurnProfile(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + req := httptest.NewRequest(http.MethodPatch, "/api/config", bytes.NewBufferString(`{ + "agents": { + "defaults": { + "turn_profile": { + "enabled": true, + "history": { "mode": "custom" } + } + } + } + }`)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, req) + if rec.Code != http.StatusBadRequest { + t.Fatalf( + "status = %d, want %d, body=%s", + rec.Code, + http.StatusBadRequest, + rec.Body.String(), + ) + } + if !strings.Contains(rec.Body.String(), "history.mode custom is not supported") { + t.Fatalf("body=%s, want turn profile validation error", rec.Body.String()) + } + + if _, err := config.LoadConfig(configPath); err != nil { + t.Fatalf("LoadConfig() after rejected patch error = %v", err) + } +} + func assertGatewayLogLevelApplied(t *testing.T, method, body string, want logger.LogLevel) { t.Helper() @@ -35,7 +124,13 @@ func assertGatewayLogLevelApplied(t *testing.T, method, body string, want logger rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("%s /api/config status = %d, want %d, body=%s", method, rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "%s /api/config status = %d, want %d, body=%s", + method, + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } if got := logger.GetLevel(); got != want { t.Fatalf("logger.GetLevel() = %v, want %v", got, want) @@ -141,10 +236,18 @@ func TestHandlePatchConfig_RejectsInvalidExecRegexPatterns(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusBadRequest { - t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusBadRequest, rec.Body.String()) + t.Fatalf( + "status = %d, want %d, body=%s", + rec.Code, + http.StatusBadRequest, + rec.Body.String(), + ) } if !bytes.Contains(rec.Body.Bytes(), []byte("custom_deny_patterns")) { - t.Fatalf("expected validation error mentioning custom_deny_patterns, body=%s", rec.Body.String()) + t.Fatalf( + "expected validation error mentioning custom_deny_patterns, body=%s", + rec.Body.String(), + ) } } @@ -200,7 +303,12 @@ func TestHandlePatchConfig_SavesChannelListSettingsPatch(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } cfg, err := config.LoadConfig(configPath) @@ -264,7 +372,12 @@ func TestHandlePatchConfig_NormalizesStringChannelArrayFields(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } cfg, err := config.LoadConfig(configPath) @@ -277,7 +390,10 @@ func TestHandlePatchConfig_NormalizesStringChannelArrayFields(t *testing.T) { picoChannel.AllowFrom[0] != "ou_a" || picoChannel.AllowFrom[1] != "ou_b" || picoChannel.AllowFrom[2] != "ou_c" { - t.Fatalf("pico allow_from = %#v, want [\"ou_a\", \"ou_b\", \"ou_c\"]", picoChannel.AllowFrom) + t.Fatalf( + "pico allow_from = %#v, want [\"ou_a\", \"ou_b\", \"ou_c\"]", + picoChannel.AllowFrom, + ) } if len(picoChannel.GroupTrigger.Prefixes) != 3 || picoChannel.GroupTrigger.Prefixes[0] != "/" || @@ -346,7 +462,12 @@ func TestHandlePatchConfig_NormalizesSingleNumericAllowFrom(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } cfg, err := config.LoadConfig(configPath) @@ -420,7 +541,11 @@ func TestHandlePatchConfig_RejectsInvalidChannelArrayFields(t *testing.T) { mux := http.NewServeMux() h.RegisterRoutes(mux) - req := httptest.NewRequest(http.MethodPatch, "/api/config", bytes.NewBufferString(tt.body)) + req := httptest.NewRequest( + http.MethodPatch, + "/api/config", + bytes.NewBufferString(tt.body), + ) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() @@ -439,13 +564,57 @@ func TestHandlePatchConfig_RejectsInvalidChannelArrayFields(t *testing.T) { t.Fatalf("LoadConfig() error = %v", err) } telegramChannel := cfg.Channels[config.ChannelTelegram] - if len(telegramChannel.AllowFrom) != 1 || telegramChannel.AllowFrom[0] != "existing-user" { - t.Fatalf("telegram allow_from = %#v, want unchanged [\"existing-user\"]", telegramChannel.AllowFrom) + if len(telegramChannel.AllowFrom) != 1 || + telegramChannel.AllowFrom[0] != "existing-user" { + t.Fatalf( + "telegram allow_from = %#v, want unchanged [\"existing-user\"]", + telegramChannel.AllowFrom, + ) } }) } } +func TestHandlePatchConfig_RejectsNegativeStreamingDeliveryValues(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + req := httptest.NewRequest(http.MethodPatch, "/api/config", bytes.NewBufferString(`{ + "channel_list": { + "pico": { + "settings": { + "streaming": { + "enabled": true, + "throttle_seconds": -1 + } + } + } + } + }`)) + req.Header.Set("Content-Type", "application/json") + + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, req) + if rec.Code != http.StatusBadRequest { + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusBadRequest, + rec.Body.String(), + ) + } + if !strings.Contains(rec.Body.String(), "streaming.throttle_seconds") { + t.Fatalf( + "response body = %q, want streaming.throttle_seconds validation error", + rec.Body.String(), + ) + } +} + func TestHandlePatchConfig_ClearingAllowFromDoesNotLeaveEmptyStringItem(t *testing.T) { configPath, cleanup := setupOAuthTestEnv(t) defer cleanup() @@ -488,7 +657,12 @@ func TestHandlePatchConfig_ClearingAllowFromDoesNotLeaveEmptyStringItem(t *testi rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } cfg, err = config.LoadConfig(configPath) @@ -505,7 +679,10 @@ func TestHandlePatchConfig_ClearingAllowFromDoesNotLeaveEmptyStringItem(t *testi t.Fatalf("ReadFile(configPath) error = %v", err) } if strings.Contains(string(configData), `"allow_from": [""]`) { - t.Fatalf("config file should not contain empty-string allow_from item: %s", string(configData)) + t.Fatalf( + "config file should not contain empty-string allow_from item: %s", + string(configData), + ) } } @@ -543,7 +720,12 @@ func TestHandlePatchConfig_CreatesMissingChannelWithTypeAndSecret(t *testing.T) rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } cfg, err = config.LoadConfig(configPath) @@ -664,7 +846,12 @@ func TestHandleUpdateConfig_SucceedsWhenPicoTokenInSecurityOnly(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PUT /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PUT /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } } @@ -687,7 +874,12 @@ func TestHandlePatchConfig_SucceedsWhenPicoTokenInSecurityOnly(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } } @@ -746,7 +938,12 @@ func TestHandlePatchConfig_PreservesDebugFlagOverride(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } if got := logger.GetLevel(); got != logger.DEBUG { t.Fatalf("logger.GetLevel() = %v, want %v", got, logger.DEBUG) @@ -776,7 +973,12 @@ func TestHandlePatchConfig_SavesDiscordTokenFromPayload(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } cfg, err := config.LoadConfig(configPath) @@ -820,7 +1022,12 @@ func TestHandlePatchConfig_DoesNotPersistShadowRegistryAuthTokenField(t *testing rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusOK { - t.Fatalf("PATCH /api/config status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + t.Fatalf( + "PATCH /api/config status = %d, want %d, body=%s", + rec.Code, + http.StatusOK, + rec.Body.String(), + ) } cfg, err := config.LoadConfig(configPath) @@ -843,7 +1050,10 @@ func TestHandlePatchConfig_DoesNotPersistShadowRegistryAuthTokenField(t *testing t.Fatalf("ReadFile(configPath) error = %v", err) } if strings.Contains(string(rawConfig), "_auth_token") { - t.Fatalf("config.json should not persist _auth_token shadow field, got:\n%s", string(rawConfig)) + t.Fatalf( + "config.json should not persist _auth_token shadow field, got:\n%s", + string(rawConfig), + ) } } @@ -879,7 +1089,11 @@ func testCommandPatterns(t *testing.T, configPath string, body string) *httptest h := NewHandler(configPath) mux := http.NewServeMux() h.RegisterRoutes(mux) - req := httptest.NewRequest(http.MethodPost, "/api/config/test-command-patterns", bytes.NewBufferString(body)) + req := httptest.NewRequest( + http.MethodPost, + "/api/config/test-command-patterns", + bytes.NewBufferString(body), + ) req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) @@ -922,7 +1136,10 @@ func TestHandleTestCommandPatterns_MatchesBlacklistNotWhitelist(t *testing.T) { t.Fatalf("expected blocked=true, body=%s", rec.Body.String()) } if bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { - t.Fatalf("expected allowed=false when blacklist matches but not whitelist, body=%s", rec.Body.String()) + t.Fatalf( + "expected allowed=false when blacklist matches but not whitelist, body=%s", + rec.Body.String(), + ) } } @@ -996,7 +1213,10 @@ func TestHandleTestCommandPatterns_InvalidRegexSkipped(t *testing.T) { t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) } if !bytes.Contains(rec.Body.Bytes(), []byte(`"allowed":true`)) { - t.Fatalf("expected allowed=true, invalid pattern skipped and valid one matched, body=%s", rec.Body.String()) + t.Fatalf( + "expected allowed=true, invalid pattern skipped and valid one matched, body=%s", + rec.Body.String(), + ) } } @@ -1036,7 +1256,12 @@ func TestHandleTestCommandPatterns_InvalidJSON(t *testing.T) { rec := httptest.NewRecorder() mux.ServeHTTP(rec, req) if rec.Code != http.StatusBadRequest { - t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusBadRequest, rec.Body.String()) + t.Fatalf( + "status = %d, want %d, body=%s", + rec.Code, + http.StatusBadRequest, + rec.Body.String(), + ) } } diff --git a/web/backend/api/gateway.go b/web/backend/api/gateway.go index 45f7e691..af6c6560 100644 --- a/web/backend/api/gateway.go +++ b/web/backend/api/gateway.go @@ -25,6 +25,7 @@ import ( "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/netbind" ppid "github.com/sipeed/picoclaw/pkg/pid" + "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/web/backend/utils" ) @@ -413,6 +414,10 @@ func computeConfigSignature(cfg *config.Config) string { if defaultModel != "" { parts = append(parts, "model:"+defaultModel) } + modelStreamingSignatures := computeModelStreamingSignatures(cfg) + if len(modelStreamingSignatures) > 0 { + parts = append(parts, "model_streaming:"+strings.Join(modelStreamingSignatures, ",")) + } toolSignatures := []string{} if cfg.Tools.ReadFile.Enabled { toolSignatures = append(toolSignatures, "read_file") @@ -491,6 +496,152 @@ func computeConfigSignature(cfg *config.Config) string { return strings.Join(parts, ";") } +func computeModelStreamingSignatures(cfg *config.Config) []string { + if cfg == nil { + return nil + } + defaultProvider := strings.TrimSpace(cfg.Agents.Defaults.Provider) + if defaultProvider == "" { + defaultProvider = "openai" + } + names := []string{strings.TrimSpace(cfg.Agents.Defaults.GetModelName())} + names = append(names, cfg.Agents.Defaults.ModelFallbacks...) + if cfg.Agents.Defaults.Routing != nil { + names = append(names, cfg.Agents.Defaults.Routing.LightModel) + } + for _, agent := range cfg.Agents.List { + if agent.Model == nil { + continue + } + names = append(names, agent.Model.Primary) + names = append(names, agent.Model.Fallbacks...) + } + + seenNames := make(map[string]bool) + seenEntries := make(map[string]bool) + signatures := make([]string, 0, len(names)) + for _, name := range names { + name = strings.TrimSpace(name) + if name == "" || seenNames[name] { + continue + } + seenNames[name] = true + for _, match := range modelConfigsMatchingSignatureRef(cfg.ModelList, name, defaultProvider) { + mc := match.model + entry := strings.Join([]string{ + name, + strconv.Itoa(match.index), + strings.TrimSpace(mc.Provider), + strings.TrimSpace(mc.Model), + strconv.FormatBool(mc.Streaming.Enabled), + }, ":") + if seenEntries[entry] { + continue + } + seenEntries[entry] = true + signatures = append(signatures, entry) + } + } + sort.Strings(signatures) + return signatures +} + +type signatureModelConfigMatch struct { + index int + model *config.ModelConfig +} + +func modelConfigsMatchingSignatureRef( + modelList []*config.ModelConfig, + raw string, + defaultProvider string, +) []signatureModelConfigMatch { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil + } + matches := make([]signatureModelConfigMatch, 0, 1) + for i, mc := range modelList { + if mc == nil || strings.TrimSpace(mc.ModelName) != raw { + continue + } + matches = append(matches, signatureModelConfigMatch{index: i, model: mc}) + } + if len(matches) > 0 { + return matches + } + for i, mc := range modelList { + if mc == nil || strings.TrimSpace(mc.Model) != raw { + continue + } + return []signatureModelConfigMatch{{index: i, model: mc}} + } + for i, mc := range modelList { + if modelConfigMatchesBareRef(mc, raw, defaultProvider) { + return []signatureModelConfigMatch{{index: i, model: mc}} + } + } + + rawRef := providers.ParseModelRef(raw, "") + rawHasProvider := rawRef != nil && hasUnambiguousProviderPrefix(raw) && + strings.TrimSpace(rawRef.Provider) != "" && strings.TrimSpace(rawRef.Model) != "" + if rawHasProvider { + for i, mc := range modelList { + if modelConfigMatchesProviderRef(mc, raw) { + return []signatureModelConfigMatch{{index: i, model: mc}} + } + } + } + return nil +} + +func hasUnambiguousProviderPrefix(raw string) bool { + provider, _, found := strings.Cut(strings.TrimSpace(raw), "/") + if !found { + return false + } + provider = strings.ToLower(strings.TrimSpace(provider)) + if provider == "" { + return false + } + normalizedProvider := providers.NormalizeProvider(provider) + if !providers.IsSupportedModelProvider(normalizedProvider) { + return false + } + return true +} + +func modelConfigMatchesProviderRef(mc *config.ModelConfig, raw string) bool { + if mc == nil { + return false + } + raw = strings.TrimSpace(raw) + if raw == "" { + return false + } + rawRef := providers.ParseModelRef(raw, "") + if rawRef == nil || strings.TrimSpace(rawRef.Provider) == "" || strings.TrimSpace(rawRef.Model) == "" { + return false + } + protocol, modelID := providers.ExtractProtocol(mc) + return providers.ModelKey(protocol, modelID) == providers.ModelKey(rawRef.Provider, rawRef.Model) +} + +func modelConfigMatchesBareRef(mc *config.ModelConfig, raw string, defaultProvider string) bool { + if mc == nil { + return false + } + raw = strings.TrimSpace(raw) + if raw == "" { + return false + } + protocol, modelID := providers.ExtractProtocol(mc) + if strings.TrimSpace(modelID) != raw { + return false + } + return providers.NormalizeProvider(protocol) == providers.NormalizeProvider(defaultProvider) +} + func computeChannelSignatures(channels config.ChannelsConfig) []string { if len(channels) == 0 { return nil diff --git a/web/backend/api/gateway_test.go b/web/backend/api/gateway_test.go index f383089a..3e118cd3 100644 --- a/web/backend/api/gateway_test.go +++ b/web/backend/api/gateway_test.go @@ -1267,6 +1267,528 @@ func TestGatewayStatusRequiresRestartAfterChannelChange(t *testing.T) { } } +func TestGatewayStatusRequiresRestartAfterDefaultModelStreamingChange(t *testing.T) { + resetGatewayTestState(t) + + configPath := filepath.Join(t.TempDir(), "config.json") + cfg := config.DefaultConfig() + cfg.Agents.Defaults.ModelName = cfg.ModelList[0].ModelName + cfg.ModelList[0].SetAPIKey("test-key") + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: false} + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + process, err := os.FindProcess(os.Getpid()) + if err != nil { + t.Fatalf("FindProcess() error = %v", err) + } + + bootSignature := computeConfigSignature(cfg) + gateway.mu.Lock() + gateway.cmd = &exec.Cmd{Process: process} + gateway.bootDefaultModel = cfg.ModelList[0].ModelName + gateway.bootConfigSignature = bootSignature + setGatewayRuntimeStatusLocked("running") + gateway.mu.Unlock() + + updatedCfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + updatedCfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + if err := config.SaveConfig(configPath, updatedCfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + gatewayHealthGet = func(string, time.Duration) (*http.Response, error) { + return mockGatewayHealthResponse(http.StatusOK, os.Getpid()), nil + } + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/gateway/status", nil) + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK) + } + + var body map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + + if got := body["gateway_status"]; got != "running" { + t.Fatalf("gateway_status = %#v, want %q", got, "running") + } + if got := body["gateway_restart_required"]; got != true { + t.Fatalf("gateway_restart_required = %#v, want true", got) + } +} + +func TestConfigSignatureIncludesModelStreamingForDefaultModelRef(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList[0].ModelName = "friendly-alias" + cfg.ModelList[0].Provider = "" + cfg.ModelList[0].Model = "openai/gpt-4o-ref" + cfg.Agents.Defaults.ModelName = "openai/gpt-4o-ref" + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: false} + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal("config signature should change when streaming changes for a default model referenced by model ref") + } +} + +func TestConfigSignatureIncludesModelStreamingForLoadBalancedAliasEntries(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "lb-alias", + Model: "openai/gpt-4o-primary", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + { + ModelName: "lb-alias", + Model: "openai/gpt-4o-secondary", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.ModelName = "lb-alias" + + before := computeConfigSignature(cfg) + + cfg.ModelList[1].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal("config signature should change when streaming changes for a load-balanced alias entry") + } +} + +func TestConfigSignatureIncludesSlashModelIDForDefaultProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "nvidia-model", + Provider: "nvidia", + Model: "z-ai/glm-5.1", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "nvidia" + cfg.Agents.Defaults.ModelName = "z-ai/glm-5.1" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change when streaming changes for a slash-containing model id on the default provider", + ) + } +} + +func TestConfigSignatureIncludesSupportedPrefixSlashModelIDForDefaultProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "openrouter-openai", + Provider: "openrouter", + Model: "openai/gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "openrouter" + cfg.Agents.Defaults.ModelName = "openai/gpt-4o" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change when streaming changes for a supported-prefix slash model id on the default provider", + ) + } +} + +func TestConfigSignatureIncludesLegacyDefaultProviderPrefixedSlashModelID(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "legacy-openrouter-openai", + Model: "openrouter/openai/gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "openrouter" + cfg.Agents.Defaults.ModelName = "openai/gpt-4o" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change when streaming changes for a legacy default-provider prefixed slash model id", + ) + } +} + +func TestConfigSignatureIncludesSlashModelIDWithoutProviderFieldForDefaultProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "nvidia-model", + Model: "z-ai/glm-5.1", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "nvidia" + cfg.Agents.Defaults.ModelName = "z-ai/glm-5.1" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change when streaming changes for a default-provider slash model id without provider field", + ) + } +} + +func TestConfigSignatureIncludesUnknownSlashPrefixModelIDWithoutProviderFieldForDefaultProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "nvidia-meta", + Model: "meta/llama-3.1-8b", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "nvidia" + cfg.Agents.Defaults.ModelName = "meta/llama-3.1-8b" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change when streaming changes for unknown-prefix default-provider slash model id", + ) + } +} + +func TestConfigSignatureDashAliasSlashModelIDMatchesProviderAlias(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "zai-model", + Provider: "zai", + Model: "glm-5.1", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "nvidia" + cfg.Agents.Defaults.ModelName = "z-ai/glm-5.1" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal("config signature should change when a dash-alias slash ref matches a provider alias") + } +} + +func TestConfigSignatureDashAliasSlashModelIDMatchesProviderAliasWithOpenAIDefault(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "zai-model", + Provider: "zai", + Model: "glm-5.1", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "openai" + cfg.Agents.Defaults.ModelName = "z-ai/glm-5.1" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change when a dash-alias slash ref matches a provider alias with OpenAI default", + ) + } +} + +func TestConfigSignatureProviderAliasRefIgnoresDefaultProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "openai-gpt", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "nvidia" + cfg.Agents.Defaults.ModelName = "gpt/gpt-4o" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal("config signature should change for a provider alias ref even when default provider differs") + } +} + +func TestConfigSignatureExplicitProviderRefIgnoresDefaultProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "openai-gpt", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "nvidia" + cfg.Agents.Defaults.ModelName = "openai/gpt-4o" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal("config signature should change for an explicit provider ref even when default provider differs") + } +} + +func TestConfigSignatureExactModelNameTakesPrecedenceOverResolvedRefs(t *testing.T) { + tests := []struct { + name string + defaultProvider string + defaultModelName string + models []*config.ModelConfig + shadowedEntryIndex int + exactModelNameIndex int + shadowedChangeMessage string + exactChangeMessage string + }{ + { + name: "slash model name shadows explicit provider ref", + defaultProvider: "nvidia", + defaultModelName: "openai/gpt-4o", + models: []*config.ModelConfig{ + { + ModelName: "openai/gpt-4o", + Provider: "nvidia", + Model: "openai/gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + { + ModelName: "openai-gpt", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + }, + shadowedEntryIndex: 1, + exactModelNameIndex: 0, + shadowedChangeMessage: "config signature should not change when an exact slash model_name shadows an explicit provider ref", + exactChangeMessage: "config signature should change when the exact slash model_name entry changes", + }, + { + name: "bare model name shadows default provider model id", + defaultProvider: "openai", + defaultModelName: "gpt-4o", + models: []*config.ModelConfig{ + { + ModelName: "gpt-4o", + Provider: "anthropic", + Model: "claude-sonnet", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + { + ModelName: "openai-gpt", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + }, + shadowedEntryIndex: 1, + exactModelNameIndex: 0, + shadowedChangeMessage: "config signature should not change when an exact bare model_name shadows a default-provider model id", + exactChangeMessage: "config signature should change when the exact bare model_name entry changes", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = tt.models + cfg.Agents.Defaults.Provider = tt.defaultProvider + cfg.Agents.Defaults.ModelName = tt.defaultModelName + + before := computeConfigSignature(cfg) + + cfg.ModelList[tt.shadowedEntryIndex].Streaming = config.ModelStreamingConfig{Enabled: true} + afterShadowedChange := computeConfigSignature(cfg) + + if before != afterShadowedChange { + t.Fatal(tt.shadowedChangeMessage) + } + + cfg.ModelList[tt.exactModelNameIndex].Streaming = config.ModelStreamingConfig{Enabled: true} + afterExactModelNameChange := computeConfigSignature(cfg) + + if before == afterExactModelNameChange { + t.Fatal(tt.exactChangeMessage) + } + }) + } +} + +func TestConfigSignatureIncludesLoadBalancedDuplicateEntryIndex(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "lb-alias", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + { + ModelName: "lb-alias", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: true}, + }, + } + cfg.Agents.Defaults.ModelName = "lb-alias" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming.Enabled = true + cfg.ModelList[1].Streaming.Enabled = false + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal("config signature should change when duplicate load-balanced entries swap streaming state") + } +} + +func TestConfigSignatureProviderDotAliasRefIgnoresDefaultProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "zai-model", + Provider: "zai", + Model: "glm-5.1", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "nvidia" + cfg.Agents.Defaults.ModelName = "z.ai/glm-5.1" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change for an explicit dot-alias provider ref even when default provider differs", + ) + } +} + +func TestConfigSignatureIncludesDefaultProviderPrefixedRefWithSplitConfig(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "openai-split", + Provider: "openai", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "openai" + cfg.Agents.Defaults.ModelName = "openai/gpt-4o" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + after := computeConfigSignature(cfg) + + if before == after { + t.Fatal( + "config signature should change when streaming changes for default-provider prefixed ref with split config", + ) + } +} + +func TestConfigSignatureBareModelRefUsesExactModelBeforeDefaultProviderModelID(t *testing.T) { + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "azure-alias", + Provider: "azure", + Model: "gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + { + ModelName: "openai-alias", + Model: "openai/gpt-4o", + Streaming: config.ModelStreamingConfig{Enabled: false}, + }, + } + cfg.Agents.Defaults.Provider = "openai" + cfg.Agents.Defaults.ModelName = "gpt-4o" + + before := computeConfigSignature(cfg) + + cfg.ModelList[0].Streaming = config.ModelStreamingConfig{Enabled: true} + afterExactModelChange := computeConfigSignature(cfg) + + if before == afterExactModelChange { + t.Fatal("config signature should change when the exact bare model entry changes streaming") + } + + cfg.ModelList[1].Streaming = config.ModelStreamingConfig{Enabled: true} + afterDefaultProviderModelChange := computeConfigSignature(cfg) + + if afterExactModelChange != afterDefaultProviderModelChange { + t.Fatal("config signature should not change when a shadowed default-provider model id changes streaming") + } +} + func TestGatewayStatusRequiresRestartAfterWebSearchConfigChange(t *testing.T) { resetGatewayTestState(t) diff --git a/web/backend/api/model_catalog.go b/web/backend/api/model_catalog.go index ce50deaf..edce03cf 100644 --- a/web/backend/api/model_catalog.go +++ b/web/backend/api/model_catalog.go @@ -12,6 +12,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/fileutil" + "github.com/sipeed/picoclaw/pkg/providers" ) // CatalogModel represents a single model entry in a saved catalog. @@ -42,7 +43,7 @@ func catalogFilePath() string { // generateCatalogKey creates a deterministic key for a provider+base+key combination. func generateCatalogKey(provider, apiBase, apiKey string) string { - provider = strings.ToLower(strings.TrimSpace(provider)) + provider = providers.NormalizeProvider(provider) apiBase = strings.TrimRight(strings.TrimSpace(apiBase), "/") hash := sha256.Sum256([]byte(apiKey)) return fmt.Sprintf("%s|%s|%x", provider, apiBase, hash[:6]) @@ -104,9 +105,10 @@ func SaveCatalog(provider, apiBase, apiKey string, models []CatalogModel) error return err } key := generateCatalogKey(provider, apiBase, apiKey) + provider = providers.NormalizeProvider(provider) store.Entries[key] = &CatalogEntry{ ID: key, - Provider: strings.ToLower(strings.TrimSpace(provider)), + Provider: provider, APIBase: strings.TrimRight(strings.TrimSpace(apiBase), "/"), APIKeyMask: maskAPIKeyValue(apiKey), Models: models, diff --git a/web/backend/api/model_status.go b/web/backend/api/model_status.go index 6cfda501..03666ae9 100644 --- a/web/backend/api/model_status.go +++ b/web/backend/api/model_status.go @@ -126,7 +126,7 @@ func hasStoredOAuthCredential(m *config.ModelConfig) (bool, bool) { func providerUsesImplicitOAuth(protocol string) bool { switch protocol { - case "antigravity", "google-antigravity": + case "antigravity": return true default: return false @@ -168,11 +168,11 @@ func requiresRuntimeProbe(m *config.ModelConfig) bool { protocol := modelProtocol(m) switch protocol { - case "claude-cli", "claudecli", "codex-cli", "codexcli", "github-copilot", "copilot": + case "claude-cli", "codex-cli", "github-copilot": return true } - if providers.IsEmptyAPIKeyAllowedForProtocol(protocol) { + if providers.IsHTTPAPIProtocol(protocol) && providers.IsEmptyAPIKeyAllowedForProtocol(protocol) { apiBase := strings.TrimSpace(m.APIBase) return apiBase == "" || hasLocalAPIBase(apiBase) } @@ -218,13 +218,13 @@ func runLocalModelProbe(m *config.ModelConfig) bool { switch protocol { case "ollama": return probeOllamaModelFunc(apiBase, modelID) - case "vllm", "lmstudio": + case "vllm", "lmstudio", "gpt4free": return probeOpenAICompatibleModelFunc(apiBase, modelID, m.APIKey()) - case "github-copilot", "copilot": + case "github-copilot": return probeTCPServiceFunc(apiBase) - case "claude-cli", "claudecli": + case "claude-cli": return probeCommandAvailableFunc("claude") - case "codex-cli", "codexcli": + case "codex-cli": return probeCommandAvailableFunc("codex") default: if hasLocalAPIBase(apiBase) { @@ -442,7 +442,7 @@ func modelProbeAPIBase(m *config.ModelConfig) string { } switch protocol { - case "github-copilot", "copilot": + case "github-copilot": return "localhost:4321" default: return "" @@ -477,7 +477,7 @@ func oauthProviderForModel(m *config.ModelConfig) (string, bool) { return oauthProviderOpenAI, true case "anthropic": return oauthProviderAnthropic, true - case "antigravity", "google-antigravity": + case "antigravity": return oauthProviderGoogleAntigravity, true default: return "", false @@ -563,16 +563,29 @@ func probeOpenAICompatibleModel(apiBase, modelID, apiKey string) bool { return false } - var resp struct { + var respEnvelope struct { Data []struct { ID string `json:"id"` } `json:"data"` } - if err := getJSON(strings.TrimRight(strings.TrimSpace(apiBase), "/")+"/models", &resp, apiKey); err != nil { + fetchURL := strings.TrimRight(strings.TrimSpace(apiBase), "/") + "/models" + if err := getJSON(fetchURL, &respEnvelope, apiKey); err == nil { + for _, model := range respEnvelope.Data { + if strings.EqualFold(strings.TrimSpace(model.ID), modelID) { + return true + } + } return false } - for _, model := range resp.Data { + var respArray []struct { + ID string `json:"id"` + } + if err := getJSON(fetchURL, &respArray, apiKey); err != nil { + return false + } + + for _, model := range respArray { if strings.EqualFold(strings.TrimSpace(model.ID), modelID) { return true } diff --git a/web/backend/api/model_status_test.go b/web/backend/api/model_status_test.go index d5463a85..7a8155e4 100644 --- a/web/backend/api/model_status_test.go +++ b/web/backend/api/model_status_test.go @@ -1,6 +1,7 @@ package api import ( + "fmt" "net/http" "net/http/httptest" "sync" @@ -39,6 +40,21 @@ func TestProbeLocalModelAvailability_OpenAICompatibleIncludesAPIKey(t *testing.T } } +func TestProbeOpenAICompatibleModel_SupportsBareArrayModels(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Fatalf("path = %q, want %q", r.URL.Path, "/v1/models") + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `[{"id":"gpt-4o-mini"},{"id":"deepseek-r1"}]`) + })) + defer srv.Close() + + if !probeOpenAICompatibleModel(srv.URL+"/v1", "gpt-4o-mini", "") { + t.Fatal("probeOpenAICompatibleModel() = false, want true for bare array /models response") + } +} + func TestRequiresRuntimeProbe_LMStudio(t *testing.T) { if !requiresRuntimeProbe(&config.ModelConfig{ Model: "lmstudio/openai/gpt-oss-20b", @@ -54,6 +70,21 @@ func TestRequiresRuntimeProbe_LMStudio(t *testing.T) { } } +func TestRequiresRuntimeProbe_GPT4Free(t *testing.T) { + if !requiresRuntimeProbe(&config.ModelConfig{ + Model: "gpt4free/gpt-4o-mini", + }) { + t.Fatal("requiresRuntimeProbe(gpt4free with default base) = false, want true") + } + + if requiresRuntimeProbe(&config.ModelConfig{ + Model: "gpt4free/gpt-4o-mini", + APIBase: "https://g4f.space/v1", + }) { + t.Fatal("requiresRuntimeProbe(gpt4free with remote base) = true, want false") + } +} + func TestModelProbeAPIBase_LMStudioDefault(t *testing.T) { got := modelProbeAPIBase(&config.ModelConfig{Model: "lmstudio/openai/gpt-oss-20b"}) if got != "http://localhost:1234/v1" { @@ -61,6 +92,13 @@ func TestModelProbeAPIBase_LMStudioDefault(t *testing.T) { } } +func TestModelProbeAPIBase_GPT4FreeDefault(t *testing.T) { + got := modelProbeAPIBase(&config.ModelConfig{Model: "gpt4free/gpt-4o-mini"}) + if got != "http://localhost:1337/v1" { + t.Fatalf("modelProbeAPIBase(gpt4free) = %q, want %q", got, "http://localhost:1337/v1") + } +} + func TestProbeLocalModelAvailability_LMStudioUsesOpenAICompatibleProbe(t *testing.T) { originalProbe := probeOpenAICompatibleModelFunc defer func() { probeOpenAICompatibleModelFunc = originalProbe }() @@ -89,6 +127,34 @@ func TestProbeLocalModelAvailability_LMStudioUsesOpenAICompatibleProbe(t *testin } } +func TestProbeLocalModelAvailability_GPT4FreeUsesOpenAICompatibleProbe(t *testing.T) { + originalProbe := probeOpenAICompatibleModelFunc + defer func() { probeOpenAICompatibleModelFunc = originalProbe }() + + called := false + probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool { + called = true + if apiBase != "http://localhost:1337/v1" { + t.Fatalf("apiBase = %q, want %q", apiBase, "http://localhost:1337/v1") + } + if modelID != "gpt-4o-mini" { + t.Fatalf("modelID = %q, want %q", modelID, "gpt-4o-mini") + } + if apiKey != "" { + t.Fatalf("apiKey = %q, want empty", apiKey) + } + return true + } + + model := &config.ModelConfig{Model: "gpt4free/gpt-4o-mini"} + if !probeLocalModelAvailability(model) { + t.Fatal("probeLocalModelAvailability(gpt4free) = false, want true") + } + if !called { + t.Fatal("probeOpenAICompatibleModelFunc was not called for gpt4free") + } +} + func TestModelProbeCacheKey_DifferentAPIKeysProduceDifferentKeys(t *testing.T) { base := &config.ModelConfig{ ModelName: "local-vllm", diff --git a/web/backend/api/models.go b/web/backend/api/models.go index dd14c54a..be594e3e 100644 --- a/web/backend/api/models.go +++ b/web/backend/api/models.go @@ -18,29 +18,16 @@ import ( "github.com/sipeed/picoclaw/pkg/providers" ) -// fetchableProviders lists providers that support OpenAI-compatible /models listing. -var fetchableProviders = map[string]bool{ - "openai": true, "deepseek": true, "openrouter": true, - "qwen-portal": true, "qwen-intl": true, "moonshot": true, - "volcengine": true, "zhipu": true, "groq": true, - "mistral": true, "nvidia": true, "cerebras": true, - "venice": true, "shengsuanyun": true, "vivgrid": true, - "siliconflow": true, - "minimax": true, "longcat": true, "modelscope": true, - "mimo": true, "avian": true, "zai": true, "novita": true, - "litellm": true, "vllm": true, "lmstudio": true, "ollama": true, -} - // registerModelRoutes binds model list management endpoints to the ServeMux. func (h *Handler) registerModelRoutes(mux *http.ServeMux) { mux.HandleFunc("GET /api/models", h.handleListModels) + mux.HandleFunc("POST /api/models/fetch", h.handleFetchModels) + mux.HandleFunc("GET /api/models/catalog", h.handleListCatalogs) + mux.HandleFunc("DELETE /api/models/catalog/{id}", h.handleDeleteCatalog) mux.HandleFunc("POST /api/models", h.handleAddModel) mux.HandleFunc("POST /api/models/default", h.handleSetDefaultModel) mux.HandleFunc("PUT /api/models/{index}", h.handleUpdateModel) mux.HandleFunc("DELETE /api/models/{index}", h.handleDeleteModel) - mux.HandleFunc("POST /api/models/fetch", h.handleFetchModels) - mux.HandleFunc("GET /api/models/catalog", h.handleListCatalogs) - mux.HandleFunc("DELETE /api/models/catalog/{id}", h.handleDeleteCatalog) mux.HandleFunc("POST /api/models/{index}/test", h.handleTestModel) mux.HandleFunc("POST /api/models/test-inline", h.handleTestInlineModel) } @@ -57,15 +44,16 @@ type modelResponse struct { Proxy string `json:"proxy,omitempty"` AuthMethod string `json:"auth_method,omitempty"` // Advanced fields - ConnectMode string `json:"connect_mode,omitempty"` - Workspace string `json:"workspace,omitempty"` - RPM int `json:"rpm,omitempty"` - MaxTokensField string `json:"max_tokens_field,omitempty"` - RequestTimeout int `json:"request_timeout,omitempty"` - ThinkingLevel string `json:"thinking_level,omitempty"` - ToolSchemaTransform string `json:"tool_schema_transform,omitempty"` - ExtraBody map[string]any `json:"extra_body,omitempty"` - CustomHeaders map[string]string `json:"custom_headers,omitempty"` + ConnectMode string `json:"connect_mode,omitempty"` + Workspace string `json:"workspace,omitempty"` + RPM int `json:"rpm,omitempty"` + MaxTokensField string `json:"max_tokens_field,omitempty"` + RequestTimeout int `json:"request_timeout,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` + ToolSchemaTransform string `json:"tool_schema_transform,omitempty"` + Streaming config.ModelStreamingConfig `json:"streaming,omitempty"` + ExtraBody map[string]any `json:"extra_body,omitempty"` + CustomHeaders map[string]string `json:"custom_headers,omitempty"` // Meta Enabled bool `json:"enabled"` Available bool `json:"available"` @@ -293,6 +281,7 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ToolSchemaTransform: m.ToolSchemaTransform, + Streaming: m.Streaming, ExtraBody: m.ExtraBody, CustomHeaders: m.CustomHeaders, Enabled: m.Enabled, @@ -441,6 +430,9 @@ func (h *Handler) handleUpdateModel(w http.ResponseWriter, r *http.Request) { if _, ok := rawFields["tool_schema_transform"]; !ok { mc.ToolSchemaTransform = cfg.ModelList[idx].ToolSchemaTransform } + if _, ok := rawFields["streaming"]; !ok { + mc.Streaming = cfg.ModelList[idx].Streaming + } // Preserve the existing Provider when the caller omits it. This keeps the // update API backward-compatible for clients that haven't started sending // the new field yet, while still allowing explicit clearing via "". @@ -648,9 +640,10 @@ func (h *Handler) handleFetchModels(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() var req struct { - Provider string `json:"provider"` - APIKey string `json:"api_key"` - APIBase string `json:"api_base"` + Provider string `json:"provider"` + APIKey string `json:"api_key"` + APIBase string `json:"api_base"` + ModelIndex *int `json:"model_index,omitempty"` } if err = json.Unmarshal(body, &req); err != nil { http.Error(w, fmt.Sprintf("Invalid JSON: %v", err), http.StatusBadRequest) @@ -662,12 +655,20 @@ func (h *Handler) handleFetchModels(w http.ResponseWriter, r *http.Request) { return } - if !fetchableProviders[strings.ToLower(req.Provider)] { + if !providers.IsModelProviderFetchable(req.Provider) { http.Error(w, fmt.Sprintf("provider %q does not support model listing", req.Provider), http.StatusBadRequest) return } + apiKey := strings.TrimSpace(req.APIKey) apiBase := strings.TrimSpace(req.APIBase) + + if apiKey == "" && req.ModelIndex != nil { + if stored := h.lookupStoredAPIKey(*req.ModelIndex, req.Provider, apiBase); stored != "" { + apiKey = stored + } + } + if apiBase == "" { apiBase = providers.DefaultAPIBaseForProtocol(req.Provider) } @@ -679,7 +680,7 @@ func (h *Handler) handleFetchModels(w http.ResponseWriter, r *http.Request) { ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second) defer cancel() - models, err := fetchUpstreamModels(ctx, req.Provider, apiBase, req.APIKey) + models, err := fetchUpstreamModels(ctx, req.Provider, apiBase, apiKey) if err != nil { http.Error(w, fmt.Sprintf("Failed to fetch models: %v", err), http.StatusBadGateway) return @@ -690,7 +691,7 @@ func (h *Handler) handleFetchModels(w http.ResponseWriter, r *http.Request) { for i, m := range models { catalogModels[i] = CatalogModel{ID: m.ID, OwnedBy: m.OwnedBy} } - if saveErr := SaveCatalog(req.Provider, apiBase, req.APIKey, catalogModels); saveErr != nil { + if saveErr := SaveCatalog(req.Provider, apiBase, apiKey, catalogModels); saveErr != nil { // Log but don't fail the request — saving catalog is non-critical logger.Warnf("Failed to save model catalog: %v", saveErr) } @@ -702,6 +703,30 @@ func (h *Handler) handleFetchModels(w http.ResponseWriter, r *http.Request) { }) } +func (h *Handler) lookupStoredAPIKey(index int, reqProvider, reqAPIBase string) string { + cfg, err := config.LoadConfig(h.configPath) + if err != nil || index < 0 || index >= len(cfg.ModelList) { + return "" + } + stored := cfg.ModelList[index] + storedProvider, _ := providers.ExtractProtocol(stored) + if providers.NormalizeProvider(reqProvider) != providers.NormalizeProvider(storedProvider) { + return "" + } + effectiveReqBase := strings.TrimSpace(reqAPIBase) + if effectiveReqBase == "" { + effectiveReqBase = providers.DefaultAPIBaseForProtocol(reqProvider) + } + effectiveStoredBase := strings.TrimSpace(stored.APIBase) + if effectiveStoredBase == "" { + effectiveStoredBase = providers.DefaultAPIBaseForProtocol(storedProvider) + } + if normalizeAPIBaseForCompare(effectiveReqBase) != normalizeAPIBaseForCompare(effectiveStoredBase) { + return "" + } + return stored.APIKey() +} + type upstreamModel struct { ID string `json:"id"` OwnedBy string `json:"owned_by,omitempty"` @@ -1005,13 +1030,13 @@ func probeModelConnectivity(m *config.ModelConfig) bool { switch protocol { case "ollama": return probeOllamaModel(apiBase, modelID) - case "vllm", "lmstudio": + case "vllm", "lmstudio", "gpt4free": return probeOpenAICompatibleModel(apiBase, modelID, m.APIKey()) - case "github-copilot", "copilot": + case "github-copilot": return probeTCPService(apiBase) - case "claude-cli", "claudecli": + case "claude-cli": return probeCommandAvailable("claude") - case "codex-cli", "codexcli": + case "codex-cli": return probeCommandAvailable("codex") default: // For remote providers (OpenAI, Anthropic, Gemini, DeepSeek, etc.), diff --git a/web/backend/api/models_test.go b/web/backend/api/models_test.go index ad6b3028..f734d214 100644 --- a/web/backend/api/models_test.go +++ b/web/backend/api/models_test.go @@ -6,6 +6,8 @@ import ( "fmt" "net/http" "net/http/httptest" + "os" + "path/filepath" "strings" "sync" "testing" @@ -762,6 +764,51 @@ func TestHandleAddModel_PersistsProvider(t *testing.T) { } } +func TestHandleListModels_ReturnsStreamingConfig(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.ModelList = []*config.ModelConfig{{ + ModelName: "streaming-model", + Provider: "openai", + Model: "gpt-4o-mini", + APIKeys: config.SimpleSecureStrings("sk-existing"), + Streaming: config.ModelStreamingConfig{Enabled: true}, + }} + if err = config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/models", nil) + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + var resp struct { + Models []modelResponse `json:"models"` + } + if err = json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("Unmarshal() error = %v", err) + } + if len(resp.Models) != 1 { + t.Fatalf("len(models) = %d, want 1", len(resp.Models)) + } + if !resp.Models[0].Streaming.Enabled { + t.Fatal("streaming.enabled = false, want true") + } +} + func TestHandleAddModel_RejectsUnsupportedProvider(t *testing.T) { configPath, cleanup := setupOAuthTestEnv(t) defer cleanup() @@ -1214,6 +1261,71 @@ func TestHandleUpdateModel_ToolSchemaTransformPreserveAndClear(t *testing.T) { } } +func TestHandleUpdateModel_StreamingPreserveAndChange(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.ModelList = []*config.ModelConfig{{ + ModelName: "editable", + Provider: "openai", + Model: "gpt-4o-mini", + APIKeys: config.SimpleSecureStrings("sk-existing"), + Streaming: config.ModelStreamingConfig{Enabled: true}, + }} + if err = config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + recPreserve := httptest.NewRecorder() + reqPreserve := httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{ + "model_name":"editable", + "provider":"openai", + "model":"gpt-4o-mini" + }`)) + reqPreserve.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(recPreserve, reqPreserve) + if recPreserve.Code != http.StatusOK { + t.Fatalf("preserve status = %d, want %d, body=%s", recPreserve.Code, http.StatusOK, recPreserve.Body.String()) + } + + afterPreserve, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() after preserve error = %v", err) + } + if !afterPreserve.ModelList[0].Streaming.Enabled { + t.Fatal("preserved streaming.enabled = false, want true") + } + + recChange := httptest.NewRecorder() + reqChange := httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{ + "model_name":"editable", + "provider":"openai", + "model":"gpt-4o-mini", + "streaming":{"enabled":false} + }`)) + reqChange.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(recChange, reqChange) + if recChange.Code != http.StatusOK { + t.Fatalf("change status = %d, want %d, body=%s", recChange.Code, http.StatusOK, recChange.Body.String()) + } + + afterChange, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() after change error = %v", err) + } + if afterChange.ModelList[0].Streaming.Enabled { + t.Fatal("streaming.enabled = true, want false after explicit update") + } +} + func TestHandleUpdateModel_PersistsProvider(t *testing.T) { configPath, cleanup := setupOAuthTestEnv(t) defer cleanup() @@ -1790,6 +1902,12 @@ func TestHandleListModels_ReturnsProviderOptionsWithoutPersistingLegacyMigration t.Fatal("openai provider option missing") } else if option.DefaultAPIBase != "https://api.openai.com/v1" { t.Fatalf("openai default_api_base = %q, want %q", option.DefaultAPIBase, "https://api.openai.com/v1") + } else if !option.SupportsFetch { + t.Fatal("openai provider option should report supports_fetch") + } else if option.DisplayName != "OpenAI" { + t.Fatalf("openai display_name = %q, want %q", option.DisplayName, "OpenAI") + } else if len(option.CommonModels) == 0 { + t.Fatal("openai common_models should not be empty") } if option, ok := optionsByID["anthropic"]; !ok { t.Fatal("anthropic provider option missing") @@ -1803,6 +1921,8 @@ func TestHandleListModels_ReturnsProviderOptionsWithoutPersistingLegacyMigration t.Fatal("github-copilot provider option missing") } else if option.DefaultAPIBase != "localhost:4321" { t.Fatalf("github-copilot default_api_base = %q, want %q", option.DefaultAPIBase, "localhost:4321") + } else if !option.Local { + t.Fatal("github-copilot should be marked local") } if option, ok := optionsByID["elevenlabs"]; !ok { t.Fatal("elevenlabs provider option missing") @@ -1819,6 +1939,19 @@ func TestHandleListModels_ReturnsProviderOptionsWithoutPersistingLegacyMigration } else if !option.EmptyAPIKeyAllowed { t.Fatal("lmstudio should allow empty api keys") } + if option, ok := optionsByID["gpt4free"]; !ok { + t.Fatal("gpt4free provider option missing") + } else { + if option.DefaultAPIBase != "http://localhost:1337/v1" { + t.Fatalf("gpt4free default_api_base = %q, want %q", option.DefaultAPIBase, "http://localhost:1337/v1") + } + if !option.EmptyAPIKeyAllowed { + t.Fatal("gpt4free should allow empty api keys") + } + if !option.SupportsFetch { + t.Fatal("gpt4free provider option should report supports_fetch") + } + } if option, ok := optionsByID["siliconflow"]; !ok { t.Fatal("siliconflow provider option missing") } else if option.DefaultAPIBase != "https://api.siliconflow.cn/v1" { @@ -1843,6 +1976,11 @@ func TestHandleListModels_ReturnsProviderOptionsWithoutPersistingLegacyMigration t.Fatal("antigravity auth method should be locked") } } + if option, ok := optionsByID["qwen-portal"]; !ok { + t.Fatal("qwen-portal provider option missing") + } else if len(option.Aliases) == 0 || option.Aliases[0] != "qwen" { + t.Fatalf("qwen-portal aliases = %#v, want to include qwen", option.Aliases) + } updated, err := config.LoadConfig(configPath) if err != nil { @@ -2453,3 +2591,106 @@ func TestHandleFetchModels_SiliconFlowUsesOpenAICompatibleEndpoint(t *testing.T) t.Fatalf("model id = %q, want %q", resp.Models[0].ID, "deepseek-ai/DeepSeek-V3") } } + +func TestHandleFetchModels_ModelIndexUsesStoredKey(t *testing.T) { + var gotAuth string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"data":[{"id":"gpt-4o","owned_by":"openai"}]}`) + })) + defer srv.Close() + + tmp := t.TempDir() + oldHome := os.Getenv("PICOCLAW_HOME") + t.Setenv("PICOCLAW_HOME", filepath.Join(tmp, ".picoclaw")) + defer func() { + if oldHome != "" { + os.Setenv("PICOCLAW_HOME", oldHome) + } else { + os.Unsetenv("PICOCLAW_HOME") + } + }() + + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "my-openai", + Provider: "openai", + Model: "gpt-4o", + APIKeys: config.SimpleSecureStrings("sk-stored-secret"), + APIBase: srv.URL, + }, + } + configPath := filepath.Join(tmp, "config.json") + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig error: %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + idx := 0 + body := fmt.Sprintf(`{"provider":"openai","api_base":"%s","model_index":%d}`, srv.URL, idx) + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/models/fetch", bytes.NewBufferString(body)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + if gotAuth != "Bearer sk-stored-secret" { + t.Fatalf("Authorization = %q, want stored key to be used", gotAuth) + } + + var resp struct { + Models []upstreamModel `json:"models"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil { + t.Fatalf("Unmarshal error: %v", err) + } + if len(resp.Models) != 1 || resp.Models[0].ID != "gpt-4o" { + t.Fatalf("unexpected response: %+v", resp) + } +} + +func TestHandleFetchModels_ModelIndexProviderMismatchRejectsKey(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "" { + t.Error("stored key should NOT be sent to mismatched provider") + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"data":[]}`) + })) + defer srv.Close() + + tmp := t.TempDir() + t.Setenv("PICOCLAW_HOME", filepath.Join(tmp, ".picoclaw")) + + cfg := config.DefaultConfig() + cfg.ModelList = []*config.ModelConfig{ + { + ModelName: "my-openai", + Provider: "openai", + Model: "gpt-4o", + APIKeys: config.SimpleSecureStrings("sk-openai-secret"), + APIBase: "https://api.openai.com/v1", + }, + } + configPath := filepath.Join(tmp, "config.json") + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig error: %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + body := fmt.Sprintf(`{"provider":"siliconflow","api_base":"%s","model_index":0}`, srv.URL) + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/models/fetch", bytes.NewBufferString(body)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) +} diff --git a/web/backend/api/oauth.go b/web/backend/api/oauth.go index 116e304b..7586a45b 100644 --- a/web/backend/api/oauth.go +++ b/web/backend/api/oauth.go @@ -767,7 +767,7 @@ func modelBelongsToProvider(provider string, modelCfg *config.ModelConfig) bool case oauthProviderAnthropic: return protocol == "anthropic" case oauthProviderGoogleAntigravity: - return protocol == "antigravity" || protocol == "google-antigravity" + return protocol == "antigravity" default: return false } diff --git a/web/backend/api/session.go b/web/backend/api/session.go index cc18ee6e..bab4ef79 100644 --- a/web/backend/api/session.go +++ b/web/backend/api/session.go @@ -50,6 +50,8 @@ type sessionChatMessage struct { Role string `json:"role"` Content string `json:"content"` Kind string `json:"kind,omitempty"` + ModelName string `json:"model_name,omitempty"` + CreatedAt *time.Time `json:"created_at,omitempty"` Media []string `json:"media,omitempty"` Attachments []sessionChatAttachment `json:"attachments,omitempty"` ToolCalls []utils.VisibleToolCall `json:"tool_calls,omitempty"` @@ -510,6 +512,8 @@ func sessionTranscriptMessages( chatMsg := sessionChatMessage{ Role: "user", Content: msg.Content, + ModelName: msg.ModelName, + CreatedAt: msg.CreatedAt, Media: append([]string(nil), msg.Media...), Attachments: attachments, } @@ -529,9 +533,11 @@ func sessionTranscriptMessages( toolCallsMsg, hasToolCallsMsg := assistantToolCallsMessage( msg.ToolCalls, + msg.ModelName, toolFeedbackMaxArgsLength, + msg.CreatedAt, ) - visibleToolMessages := visibleAssistantToolMessages(msg.ToolCalls) + visibleToolMessages := visibleAssistantToolMessages(msg.ToolCalls, msg.ModelName, msg.CreatedAt) // Pico web chat can persist both visible `message` tool output and a // later plain assistant reply in the same turn. Hide only the fixed @@ -556,6 +562,8 @@ func sessionTranscriptMessages( chatMsg := sessionChatMessage{ Role: "assistant", Content: content, + ModelName: msg.ModelName, + CreatedAt: msg.CreatedAt, Media: append([]string(nil), msg.Media...), Attachments: attachments, } @@ -682,15 +690,19 @@ func assistantThoughtMessage(msg providers.Message) (sessionChatMessage, bool) { return sessionChatMessage{}, false } return sessionChatMessage{ - Role: "assistant", - Content: reasoning, - Kind: "thought", + Role: "assistant", + Content: reasoning, + Kind: "thought", + ModelName: msg.ModelName, + CreatedAt: msg.CreatedAt, }, true } func assistantToolCallsMessage( toolCalls []providers.ToolCall, + modelName string, toolFeedbackMaxArgsLength int, + createdAt *time.Time, ) (sessionChatMessage, bool) { if len(toolCalls) == 0 { return sessionChatMessage{}, false @@ -707,6 +719,8 @@ func assistantToolCallsMessage( return sessionChatMessage{ Role: "assistant", Kind: "tool_calls", + ModelName: modelName, + CreatedAt: createdAt, ToolCalls: visibleToolCalls, }, true } @@ -718,7 +732,11 @@ func visibleAssistantToolArgsPreview( return utils.VisibleToolCallArgumentsPreview(tc, toolFeedbackMaxArgsLength) } -func visibleAssistantToolMessages(toolCalls []providers.ToolCall) []sessionChatMessage { +func visibleAssistantToolMessages( + toolCalls []providers.ToolCall, + modelName string, + createdAt *time.Time, +) []sessionChatMessage { if len(toolCalls) == 0 { return nil } @@ -734,8 +752,10 @@ func visibleAssistantToolMessages(toolCalls []providers.ToolCall) []sessionChatM continue } messages = append(messages, sessionChatMessage{ - Role: "assistant", - Content: content, + Role: "assistant", + Content: content, + ModelName: modelName, + CreatedAt: createdAt, }) } @@ -918,6 +938,11 @@ func (h *Handler) handleGetSession(w http.ResponseWriter, r *http.Request) { } } + for i := range sess.Messages { + if sess.Messages[i].CreatedAt == nil { + sess.Messages[i].CreatedAt = &sess.Updated + } + } messages := detailSessionMessages(sess.Messages, toolFeedbackMaxArgsLength) w.Header().Set("Content-Type", "application/json") diff --git a/web/backend/api/session_test.go b/web/backend/api/session_test.go index 760935db..8604f1df 100644 --- a/web/backend/api/session_test.go +++ b/web/backend/api/session_test.go @@ -564,7 +564,7 @@ func TestHandleGetSession_ReconstructsThoughtFromAssistantReasoningContent(t *te sessionKey := picoSessionPrefix + "detail-reasoning-content" for _, msg := range []providers.Message{ {Role: "user", Content: "hello"}, - {Role: "assistant", Content: "final visible answer", ReasoningContent: "internal chain of thought"}, + {Role: "assistant", Content: "final visible answer", ModelName: "gpt-5.4", ReasoningContent: "internal chain of thought"}, } { if err := store.AddFullMessage(nil, sessionKey, msg); err != nil { t.Fatalf("AddFullMessage() error = %v", err) @@ -597,9 +597,15 @@ func TestHandleGetSession_ReconstructsThoughtFromAssistantReasoningContent(t *te resp.Messages[1].Kind != "thought" { t.Fatalf("thought message = %#v, want assistant thought/internal chain of thought", resp.Messages[1]) } + if resp.Messages[1].ModelName != "gpt-5.4" { + t.Fatalf("thought model_name = %q, want %q", resp.Messages[1].ModelName, "gpt-5.4") + } if resp.Messages[2].Role != "assistant" || resp.Messages[2].Content != "final visible answer" { t.Fatalf("final message = %#v, want assistant/final visible answer", resp.Messages[2]) } + if resp.Messages[2].ModelName != "gpt-5.4" { + t.Fatalf("final model_name = %q, want %q", resp.Messages[2].ModelName, "gpt-5.4") + } } func TestHandleGetSession_ReconstructsRefreshMatrixForThoughtAndToolSummary(t *testing.T) { @@ -725,8 +731,9 @@ func TestHandleGetSession_ReconstructsVisibleMessageToolOutputWithoutDuplicateSu for _, msg := range []providers.Message{ {Role: "user", Content: "test"}, { - Role: "assistant", - Content: "", + Role: "assistant", + Content: "", + ModelName: "gpt-5.4-mini", ToolCalls: []providers.ToolCall{ { ID: "call_1", @@ -771,9 +778,15 @@ func TestHandleGetSession_ReconstructsVisibleMessageToolOutputWithoutDuplicateSu t.Fatalf("first message = %#v, want user/test", resp.Messages[0]) } assertVisibleToolCallMessage(t, resp.Messages[1], "message") + if resp.Messages[1].ModelName != "gpt-5.4-mini" { + t.Fatalf("tool_calls model_name = %q, want %q", resp.Messages[1].ModelName, "gpt-5.4-mini") + } if resp.Messages[2].Role != "assistant" || resp.Messages[2].Content != "visible tool output" { t.Fatalf("assistant message = %#v, want visible tool output", resp.Messages[2]) } + if resp.Messages[2].ModelName != "gpt-5.4-mini" { + t.Fatalf("visible tool output model_name = %q, want %q", resp.Messages[2].ModelName, "gpt-5.4-mini") + } } func TestHandleGetSession_PreservesFinalAssistantReplyAfterMessageToolOutput(t *testing.T) { diff --git a/web/frontend/src/api/channels.ts b/web/frontend/src/api/channels.ts index 42a3a060..36389ca8 100644 --- a/web/frontend/src/api/channels.ts +++ b/web/frontend/src/api/channels.ts @@ -77,6 +77,12 @@ export async function patchAppConfig( }) } +export async function resetAppConfig(): Promise { + return request("/api/config/reset", { + method: "POST", + }) +} + // WeChat QR login flow API export interface WeixinFlowResponse { diff --git a/web/frontend/src/api/models.ts b/web/frontend/src/api/models.ts index 9fd29e0f..a319eda5 100644 --- a/web/frontend/src/api/models.ts +++ b/web/frontend/src/api/models.ts @@ -20,6 +20,9 @@ export interface ModelInfo { request_timeout?: number thinking_level?: string tool_schema_transform?: string + streaming?: { + enabled?: boolean + } extra_body?: Record custom_headers?: Record // Meta @@ -33,12 +36,20 @@ export interface ModelInfo { export interface ModelProviderOption { id: string + display_name?: string + icon_slug?: string + domain?: string default_api_base: string empty_api_key_allowed: boolean create_allowed: boolean default_model_allowed: boolean + supports_fetch?: boolean default_auth_method?: string auth_method_locked?: boolean + local?: boolean + priority?: number + common_models?: string[] + aliases?: string[] } interface ModelsListResponse { @@ -149,12 +160,14 @@ export async function testModelInline( export interface UpstreamModel { id: string owned_by?: string + extra?: Record } export interface FetchModelsRequest { provider: string api_key?: string api_base?: string + model_index?: number } export interface FetchModelsResponse { diff --git a/web/frontend/src/api/sessions.ts b/web/frontend/src/api/sessions.ts index edd7d7c2..8c0c5e90 100644 --- a/web/frontend/src/api/sessions.ts +++ b/web/frontend/src/api/sessions.ts @@ -14,7 +14,9 @@ export interface SessionDetail { messages: { role: "user" | "assistant" content: string + created_at?: string kind?: "normal" | "thought" | "tool_calls" + model_name?: string media?: string[] attachments?: { type?: "image" | "audio" | "video" | "file" diff --git a/web/frontend/src/components/channels/channel-config-page.tsx b/web/frontend/src/components/channels/channel-config-page.tsx index 8a8300d0..10e2a225 100644 --- a/web/frontend/src/components/channels/channel-config-page.tsx +++ b/web/frontend/src/components/channels/channel-config-page.tsx @@ -661,6 +661,7 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) { configuredSecrets={configuredSecrets} hiddenKeys={[...hiddenKeys, "bot_id"]} requiredKeys={requiredKeys} + supportsStreaming fieldErrors={fieldErrors} registerArrayFieldFlusher={registerArrayFieldFlusher} arrayFieldResetVersion={arrayFieldResetVersion} @@ -675,6 +676,7 @@ export function ChannelConfigPage({ channelName }: ChannelConfigPageProps) { configuredSecrets={configuredSecrets} hiddenKeys={hiddenKeys} requiredKeys={requiredKeys} + supportsStreaming={channel?.name === "pico"} fieldErrors={fieldErrors} registerArrayFieldFlusher={registerArrayFieldFlusher} arrayFieldResetVersion={arrayFieldResetVersion} diff --git a/web/frontend/src/components/channels/channel-forms/generic-form.tsx b/web/frontend/src/components/channels/channel-forms/generic-form.tsx index c8ee3f69..ee429321 100644 --- a/web/frontend/src/components/channels/channel-forms/generic-form.tsx +++ b/web/frontend/src/components/channels/channel-forms/generic-form.tsx @@ -17,12 +17,15 @@ import { Field, KeyInput, SwitchCardField } from "@/components/shared-form" import { Card, CardContent } from "@/components/ui/card" import { Input } from "@/components/ui/input" +import { StreamingConfigField } from "./streaming-config-field" + interface GenericFormProps { config: ChannelConfig onChange: (key: string, value: unknown) => void configuredSecrets?: string[] hiddenKeys?: string[] requiredKeys?: string[] + supportsStreaming?: boolean fieldErrors?: Record registerArrayFieldFlusher?: ( fieldPath: string, @@ -43,6 +46,7 @@ const OBJECT_FIELDS = new Set([ "allow_from", "allow_origins", "groups", + "streaming", ]) function formatLabel(key: string): string { @@ -78,6 +82,7 @@ export function GenericForm({ configuredSecrets = [], hiddenKeys = [], requiredKeys = [], + supportsStreaming = false, fieldErrors = {}, registerArrayFieldFlusher, arrayFieldResetVersion, @@ -89,6 +94,9 @@ export function GenericForm({ const typingConfig = asRecord(config.typing) const placeholderConfig = asRecord(config.placeholder) const placeholderEnabled = asBool(placeholderConfig.enabled) + const showStreamingConfig = + (config.streaming !== undefined || supportsStreaming) && + !hiddenFieldSet.has("streaming") const rawFields = Object.keys(config).filter( (k) => @@ -264,7 +272,8 @@ export function GenericForm({ (config.group_trigger !== undefined && !hiddenFieldSet.has("group_trigger")) || (config.typing !== undefined && !hiddenFieldSet.has("typing")) || - (config.placeholder !== undefined && !hiddenFieldSet.has("placeholder")) + (config.placeholder !== undefined && !hiddenFieldSet.has("placeholder")) || + (config.streaming !== undefined && !hiddenFieldSet.has("streaming")) return (
@@ -375,6 +384,15 @@ export function GenericForm({
)} + {showStreamingConfig && ( +
+ onChange("streaming", value)} + /> +
+ )} + {config.placeholder !== undefined && !hiddenFieldSet.has("placeholder") && (
diff --git a/web/frontend/src/components/channels/channel-forms/streaming-config-field.tsx b/web/frontend/src/components/channels/channel-forms/streaming-config-field.tsx new file mode 100644 index 00000000..f14b55ac --- /dev/null +++ b/web/frontend/src/components/channels/channel-forms/streaming-config-field.tsx @@ -0,0 +1,98 @@ +import { useTranslation } from "react-i18next" + +import { Field, SwitchCardField } from "@/components/shared-form" +import { Input } from "@/components/ui/input" + +interface StreamingConfigFieldProps { + value: unknown + onChange: (value: Record) => void +} + +function asRecord(value: unknown): Record { + if (value && typeof value === "object" && !Array.isArray(value)) { + return value as Record + } + return {} +} + +function asBool(value: unknown): boolean { + return value === true +} + +function numberInputValue(value: unknown): string { + return typeof value === "number" && value > 0 ? String(value) : "" +} + +export function StreamingConfigField({ + value, + onChange, +}: StreamingConfigFieldProps) { + const { t } = useTranslation() + const streamingConfig = asRecord(value) + const streamingEnabled = asBool(streamingConfig.enabled) + + const update = (patch: Record) => { + onChange({ ...streamingConfig, ...patch }) + } + + const handleEnabledChange = (checked: boolean) => { + if (!checked) { + onChange({ + enabled: false, + throttle_seconds: null, + min_growth_chars: null, + }) + return + } + update({ enabled: true }) + } + + return ( + + {streamingEnabled && ( +
+ + + update({ + throttle_seconds: + e.target.value === "" ? 0 : Number(e.target.value), + }) + } + placeholder="0" + /> + + + + update({ + min_growth_chars: + e.target.value === "" ? 0 : Number(e.target.value), + }) + } + placeholder="0" + /> + +
+ )} + + ) +} diff --git a/web/frontend/src/components/channels/channel-forms/telegram-form.tsx b/web/frontend/src/components/channels/channel-forms/telegram-form.tsx index f9c7c778..31405be6 100644 --- a/web/frontend/src/components/channels/channel-forms/telegram-form.tsx +++ b/web/frontend/src/components/channels/channel-forms/telegram-form.tsx @@ -14,6 +14,8 @@ import { Field, KeyInput, SwitchCardField } from "@/components/shared-form" import { Card, CardContent } from "@/components/ui/card" import { Input } from "@/components/ui/input" +import { StreamingConfigField } from "./streaming-config-field" + interface TelegramFormProps { config: ChannelConfig onChange: (key: string, value: unknown) => void @@ -125,6 +127,13 @@ export function TelegramForm({ />
+
+ onChange("streaming", value)} + /> +
+
{!isCollapsedBlock && ( -
+
PicoClaw + {trimmedModelName && ( + <> + • + {trimmedModelName} + + )} {formattedTimestamp && ( <> • @@ -104,13 +113,21 @@ export function AssistantMessage({ )} {collapsedLabel} -
- {trimmedModelName} )} - /> +
+
+ {formattedTimestamp && ( + {formattedTimestamp} + )} + +
)} {(!isCollapsedBlock || isExpanded) && isToolCalls && hasToolCalls && ( diff --git a/web/frontend/src/components/chat/chat-composer.tsx b/web/frontend/src/components/chat/chat-composer.tsx index 569ed21e..43bc8a46 100644 --- a/web/frontend/src/components/chat/chat-composer.tsx +++ b/web/frontend/src/components/chat/chat-composer.tsx @@ -1,5 +1,5 @@ import { IconArrowUp, IconPhotoPlus, IconX } from "@tabler/icons-react" -import type { KeyboardEvent } from "react" +import { useRef, type KeyboardEvent as ReactKeyboardEvent } from "react" import { useTranslation } from "react-i18next" import TextareaAutosize from "react-textarea-autosize" @@ -52,14 +52,25 @@ export function ChatComposer({ }: ChatComposerProps) { const { t } = useTranslation() const canInput = inputDisabledReason === null + const composingRef = useRef(false) const disabledMessage = inputDisabledReason === null ? null : t(`chat.disabledPlaceholder.${inputDisabledReason}`) const placeholder = disabledMessage ?? t("chat.placeholder") - const handleKeyDown = (e: KeyboardEvent) => { - if (e.nativeEvent.isComposing) return + const handleKeyDown = (e: ReactKeyboardEvent) => { + const nativeEvent = e.nativeEvent as Event & { + isComposing?: boolean + keyCode?: number + } + if ( + composingRef.current || + nativeEvent.isComposing || + nativeEvent.keyCode === 229 + ) { + return + } if (e.key === "Enter" && !e.shiftKey) { e.preventDefault() onSend() @@ -98,6 +109,12 @@ export function ChatComposer({ onInputChange(e.target.value)} + onCompositionStart={() => { + composingRef.current = true + }} + onCompositionEnd={() => { + composingRef.current = false + }} onKeyDown={handleKeyDown} placeholder={placeholder} disabled={!canInput} diff --git a/web/frontend/src/components/chat/chat-page.tsx b/web/frontend/src/components/chat/chat-page.tsx index be30ff03..3b158843 100644 --- a/web/frontend/src/components/chat/chat-page.tsx +++ b/web/frontend/src/components/chat/chat-page.tsx @@ -351,7 +351,7 @@ export function ChatPage() {
{messages.length === 0 && !isTyping && ( @@ -376,6 +376,7 @@ export function ChatPage() { content={msg.content} attachments={msg.attachments} kind={msg.kind} + modelName={msg.modelName} toolCalls={msg.toolCalls} timestamp={msg.timestamp} /> @@ -383,6 +384,7 @@ export function ChatPage() { )}
diff --git a/web/frontend/src/components/chat/user-message.tsx b/web/frontend/src/components/chat/user-message.tsx index 44b873aa..80ebc7ca 100644 --- a/web/frontend/src/components/chat/user-message.tsx +++ b/web/frontend/src/components/chat/user-message.tsx @@ -1,17 +1,23 @@ import { IconCheck, IconCopy } from "@tabler/icons-react" +import { useTranslation } from "react-i18next" import { Button } from "@/components/ui/button" import { useCopyToClipboard } from "@/hooks/use-copy-to-clipboard" +import { formatMessageTime } from "@/hooks/use-pico-chat" import { cn } from "@/lib/utils" import type { ChatAttachment } from "@/store/chat" -import { useTranslation } from "react-i18next" interface UserMessageProps { content: string attachments?: ChatAttachment[] + timestamp?: string | number } -export function UserMessage({ content, attachments = [] }: UserMessageProps) { +export function UserMessage({ + content, + attachments = [], + timestamp = "", +}: UserMessageProps) { const { t } = useTranslation() const { copy, isCopied } = useCopyToClipboard() const hasText = content.trim().length > 0 @@ -22,6 +28,8 @@ export function UserMessage({ content, attachments = [] }: UserMessageProps) { const copyMessageLabel = isCopied ? t("chat.copiedLabel") : t("chat.copyMessage") + const formattedTimestamp = + timestamp !== "" ? formatMessageTime(timestamp) : "" return (
@@ -81,6 +89,12 @@ export function UserMessage({ content, attachments = [] }: UserMessageProps) {
)} + + {formattedTimestamp && ( + + {formattedTimestamp} + + )}
) } diff --git a/web/frontend/src/components/config/config-page.tsx b/web/frontend/src/components/config/config-page.tsx index 8af5815f..b3787726 100644 --- a/web/frontend/src/components/config/config-page.tsx +++ b/web/frontend/src/components/config/config-page.tsx @@ -5,7 +5,7 @@ import { useEffect, useState } from "react" import { useTranslation } from "react-i18next" import { toast } from "sonner" -import { patchAppConfig } from "@/api/channels" +import { patchAppConfig, resetAppConfig } from "@/api/channels" import { launcherFetch } from "@/api/http" import { postLauncherDashboardSetup } from "@/api/launcher-auth" import { @@ -32,6 +32,7 @@ import { EMPTY_LAUNCHER_FORM, type LauncherForm, type MCPServerForm, + type TurnProfileForm, buildFormFromConfig, parseCIDRText, parseFloatField, @@ -40,6 +41,17 @@ import { parseMultilineList, } from "@/components/config/form-model" import { PageHeader } from "@/components/page-header" +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, + AlertDialogTrigger, +} from "@/components/ui/alert-dialog" import { Badge } from "@/components/ui/badge" import { Button } from "@/components/ui/button" import { showSaveSuccessOrRestartToast } from "@/lib/restart-required" @@ -60,6 +72,33 @@ function buildStringMapMergePatch( return patch } +function buildTurnProfilePatch( + profile: TurnProfileForm, +): Record { + const result: Record = { + enabled: profile.enabled, + history: { mode: profile.historyMode }, + system_prompt: { mode: profile.systemPromptMode }, + skills: { mode: profile.skillsMode }, + tools: { mode: profile.toolsMode }, + } + + if (profile.skillsMode === "custom") { + result.skills = { + mode: "custom", + allow: parseMultilineList(profile.skillsAllowText), + } + } + if (profile.toolsMode === "custom") { + result.tools = { + mode: "custom", + allow: parseMultilineList(profile.toolsAllowText), + } + } + + return result +} + export function ConfigPage() { const { t } = useTranslation() const queryClient = useQueryClient() @@ -72,15 +111,24 @@ export function ConfigPage() { const [autoStartEnabled, setAutoStartEnabled] = useState(false) const [autoStartBaseline, setAutoStartBaseline] = useState(false) const [saving, setSaving] = useState(false) + const [showFactoryResetDialog, setShowFactoryResetDialog] = useState(false) const { data, isLoading, error } = useQuery({ queryKey: ["config"], queryFn: async () => { - const res = await launcherFetch("/api/config") - if (!res.ok) { - throw new Error("Failed to load config") + const controller = new AbortController() + const timer = setTimeout(() => controller.abort(), 5000) + try { + const res = await launcherFetch("/api/config", { + signal: controller.signal, + }) + if (!res.ok) { + throw new Error("Failed to load config") + } + return res.json() + } finally { + clearTimeout(timer) } - return res.json() }, }) @@ -201,6 +249,13 @@ export function ConfigPage() { ) } + const handleTurnProfileFieldChange = ( + key: K, + value: TurnProfileForm[K], + ) => { + updateField("turnProfile", { ...form.turnProfile, [key]: value }) + } + const handleReset = () => { setForm(baseline) setLauncherForm(launcherBaseline) @@ -208,6 +263,27 @@ export function ConfigPage() { toast.info(t("pages.config.reset_success")) } + const handleFactoryReset = async () => { + try { + await resetAppConfig() + const fresh = await launcherFetch("/api/config").then((r) => r.json()) + const parsed = buildFormFromConfig(fresh) + setForm(parsed) + setBaseline(parsed) + await queryClient.invalidateQueries({ queryKey: ["config"] }) + await refreshGatewayState() + toast.success(t("pages.config.factory_reset_success")) + } catch (err) { + toast.error( + err instanceof Error + ? err.message + : t("pages.config.factory_reset_error"), + ) + } finally { + setShowFactoryResetDialog(false) + } + } + const handleSave = async () => { try { setSaving(true) @@ -273,6 +349,7 @@ export function ConfigPage() { "Summarize token percent", { min: 1, max: 100 }, ) + const turnProfile = buildTurnProfilePatch(form.turnProfile) const heartbeatInterval = parseIntField( form.heartbeatInterval, "Heartbeat interval", @@ -507,6 +584,7 @@ export function ConfigPage() { max_tool_iterations: maxToolIterations, summarize_message_threshold: summarizeMessageThreshold, summarize_token_percent: summarizeTokenPercent, + turn_profile: turnProfile, }, }, session: { @@ -625,8 +703,38 @@ export function ConfigPage() { } } + const factoryResetButton = ( + + + + + + + + {t("pages.config.factory_reset_confirm_title")} + + + {t("pages.config.factory_reset_confirm_desc")} + + + + {t("common.cancel")} + + {t("pages.config.factory_reset_confirm")} + + + + + ) + const actionButtons = (
+ {factoryResetButton}