diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index dea2e5f..7bae4ae 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -27,7 +27,22 @@ jobs: - uses: actions/setup-go@v5 with: go-version: "1.26.x" - - run: go test ./... + # Persist the embedded-postgres binary so most runs skip the download + # entirely. The version is pinned by the embedded-postgres library (V16). + - name: Cache embedded-postgres binary + uses: actions/cache@v4 + with: + path: ~/.embedded-postgres-go + key: embedded-postgres-${{ runner.os }}-16.9.0 + # The DB integration tests each start their own embedded postgres. Run + # packages serially (-p 1): concurrent StartEmbedded calls otherwise race + # the shared cold-cache Maven download (throttled -> "no version found + # matching 16.9.0") and collide on ports (FreePort closes the listener + # before postgres binds), which fails StartEmbedded with the server never + # becoming ready. + - run: go test -p 1 ./... + env: + CAPTAIN_DB_EMBEDDED_TEST: "1" build: runs-on: ubuntu-latest diff --git a/.github/workflows/update-llm-model-registry.yml b/.github/workflows/update-llm-model-registry.yml new file mode 100644 index 0000000..22fa99d --- /dev/null +++ b/.github/workflows/update-llm-model-registry.yml @@ -0,0 +1,72 @@ +name: Update LLM Model Registry + +on: + schedule: + - cron: "0 3 * * 1" + workflow_dispatch: + +permissions: + contents: write + pull-requests: write + +concurrency: + group: update-llm-model-registry + cancel-in-progress: false + +jobs: + update: + runs-on: ubuntu-latest + env: + PR_BRANCH: chore/update-llm-model-registry + + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + ref: ${{ github.event.repository.default_branch }} + + - name: Install Go + uses: actions/setup-go@v5 + with: + go-version: "1.26.x" + + - name: Generate model registry + run: make models-update + + - name: Check for registry changes + id: changes + run: | + if git diff --quiet -- pkg/api/registry/models.json; then + echo "changed=false" >> "$GITHUB_OUTPUT" + else + git diff --stat -- pkg/api/registry/models.json + echo "changed=true" >> "$GITHUB_OUTPUT" + fi + + - name: Commit and push registry update + if: steps.changes.outputs.changed == 'true' + run: | + git config user.name "github-actions[bot]" + git config user.email "github-actions[bot]@users.noreply.github.com" + git checkout -B "$PR_BRANCH" + git add pkg/api/registry/models.json + git commit -m "chore: update llm model registry" + git push --force-with-lease origin "$PR_BRANCH" + + - name: Create pull request + if: steps.changes.outputs.changed == 'true' + env: + GH_TOKEN: ${{ github.token }} + run: | + if gh pr view "$PR_BRANCH" --repo "${{ github.repository }}" >/dev/null 2>&1; then + echo "Pull request already exists for $PR_BRANCH" + exit 0 + fi + + gh pr create \ + --repo "${{ github.repository }}" \ + --base "${{ github.event.repository.default_branch }}" \ + --head "$PR_BRANCH" \ + --title "chore: update llm model registry" \ + --body "Automated update from models.dev using \`make generate-llm-model-registry\`." diff --git a/.gitignore b/.gitignore index c49149d..8389c11 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ pkg/cli/webapp/node_modules/ pkg/cli/webapp/*.tsbuildinfo # goreleaser output at repo root. /dist/ +dist/ # The built webapp under pkg/cli/webapp/dist is gitignored — it's a large # minified bundle whose vite build needs a local clicky-ui link: dependency # that isn't available in CI. A global `dist/` ignore excludes the directory, @@ -23,6 +24,8 @@ pkg/cli/webapp/*.tsbuildinfo !pkg/cli/webapp/dist/ pkg/cli/webapp/dist/* !pkg/cli/webapp/dist/.gitkeep -!pkg/cli/webapp/dist/index.html .grite/ hack/ +.ok/ +.okignore +.ginkgo/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..024bc37 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,14 @@ +# captain — agent notes + +Shared ways of working, the gavel todo workflow, and global skills come from the root ~/.agents/AGENTS.md. + +## Memory +- [Commons-db migrations (HCL + SQL)](.agents/memory/commons-db-migrations.md) — HCL for tables, post-HCL SQL for views/triggers/deferred constraints; verify with a seeded PostgreSQL smoke, not dry apply. +- [AI model fallback, recommendations & whoami catalogs](.agents/memory/ai-model-catalog-fallback.md) — IsFallbackEligible/ErrNoAPIKey classification in pkg/ai, per-family model retention in model_lists.go, keyless cmux backends in whoami. +- [AI schema shaping, generation config & runtime logging](.agents/memory/ai-schema-generation-logging.md) — SchemaJSONForBackend provider transforms vs local validation, EffortConfig in generation_config.go, agent:model[:effort] log identity. +- [Commit grouping for mixed Captain diffs](.agents/memory/commit-grouping.md) — group by semantic product slice not directory tree, tests travel with features, every file accounted for exactly once. +- [Session viewer, Codex/Claude parsing & model normalization](.agents/memory/sessions-viewer-metadata.md) — SessionBrowser/SessionInspector seams, shared Codex extraction, synthetic event tools, unified Title/InitialPrompt derivation. +- [History discovery, performance & bash categorization](.agents/memory/history-discovery-performance.md) — case-insensitive session-dir matching in both finders, prune Claude/Codex discovery early, pkg/bash priority-based classification. +- [Prompt CLI runs: JSON contract, context.dir, run tables & Codex app-server](.agents/memory/prompt-cli-runs.md) — AIPromptResult full-request JSON, context.dir normalization into cmd.Dir, -M comparison tables, app-server JSON-RPC seams. +- [Prompt workbench, serve UI, catalog & backend modes](.agents/memory/prompt-workbench-serve-ui.md) — captain serve hosts the workbench, catalog-backed family-first selectors, cmux-style backend modes land across four seams, rebuild linked clicky-ui dist. +- [Lint cleanup: gavel lint sweeps & betterleaks fixes](.agents/memory/lint-cleanup.md) — snapshot gavel lint baseline, split by ownership slices, betterleaks comment hits fixed by neutral wording. diff --git a/Makefile b/Makefile index 4f14f80..68cfea6 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ # Makefile stub - forwards to Taskfile # Install task: https://taskfile.dev/installation/ -.PHONY: build lint test install fmt tidy clean all +.PHONY: build lint test install fmt tidy clean all models-update build: @task build @@ -26,3 +26,6 @@ clean: all: @task all + +models-update: + @task models:update diff --git a/README.md b/README.md index e3a63c4..6ef271e 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,46 @@ It provides tools to: - run a web UI for launching and chatting with AI agents - configure default backend, model, and AI safety toggles +```xml h=185px w=358px preview +
+ +
$2,500
+ +

Drag to adjust — the figure updates live.

+ +
+``` + +```html preview +
+
+ +
+``` + ## What it does From the codebase, Captain is organized around a few core capabilities: @@ -93,7 +133,7 @@ Supported backends are inferred from code and dependencies, including: ### 6. Adapter status and configuration - `whoami` — lists every AI adapter (API providers and CLI agents), how each is authenticated, whether its binary is installed, and the models it exposes -- `configure` — interactive wizard to set default backend, model, reasoning effort, budget, timeout, and feature toggles (caching, MCP, hooks, skills, user/project settings, memory) saved to `~/.captain.yaml` +- `configure` — configure per-provider agent, model, and effort defaults in `~/.captain.yaml`, or validate and save direct-provider API tokens in `~/.config/captain/vault` ### 7. Web UI and MCP server @@ -459,9 +499,18 @@ Lists every AI adapter (API providers and CLI agents: `anthropic`, `openai`, `ge ```bash captain configure +captain configure openai +captain configure openai --test +captain configure openai --agent codex-agent --model gpt-5.6-sol --effort high --active ``` -Interactive wizard that writes `~/.captain.yaml` with defaults for backend, model, reasoning effort, budget, timeout, and feature toggles (caching, MCP, hooks, skills, user/project settings, memory). These defaults apply to `captain ai prompt`, `captain ai agent`, `captain ai test`, and other AI commands. +The providerless interactive wizard writes `~/.captain.yaml` with defaults for backend, model, reasoning effort, budget, timeout, and feature toggles (caching, MCP, hooks, skills, user/project settings, memory). + +With a provider and any of `--agent`, `--model`, `--effort`, or `--active`, `configure` saves provider-specific runtime defaults in `~/.captain.yaml`. `--effort default` clears an explicit effort so the selected model chooses its own default. The active provider supplies the agent and model for completely flagless runs. Explicit command flags still win, and each fallback model independently inherits defaults from its own provider. + +Passing `anthropic`, `openai`, `gemini`, or `deepseek` securely prompts for an API token, validates it with the provider's model-list endpoint, and saves it to `~/.config/captain/vault` only when validation succeeds. `--test` validates the currently effective token without writing. Automation may use `--token`, but interactive input is preferred because command-line arguments can be retained in shell history or process listings. Explicit runtime keys take precedence over the vault, and the vault takes precedence over environment variables. + +The `/whoami` page in `captain serve` exposes the same token operations plus agent, model, effort, and active-provider controls. Stored tokens are never displayed; only a masked identifier and credential source are returned. ### Serve @@ -474,7 +523,7 @@ task www:dev task www:build ``` -Starts an HTTP API and embedded web UI. The UI launches `captain ai agent` operations and opens follow-up chat windows that resume the returned session. `--dev` starts the Vite dev server from `pkg/cli/webapp` and proxies `/api` back to the Go process. Use `task www:dev` for the local Go-backed Vite proxy with the browser opened, and `task www:build` to rebuild the embedded web UI assets. +Starts an HTTP API and embedded web UI. The UI launches `captain ai agent` operations and opens follow-up chat windows that resume the returned session. `--dev` starts the Vite dev server from `pkg/cli/webapp`, binds the Go API to a random free port, and proxies `/api` to it. Pass `--port` to use a specific API port in development. Use `task www:dev` for the local Go-backed Vite proxy with the browser opened, and `task www:build` to rebuild the embedded web UI assets. ### MCP server @@ -610,4 +659,4 @@ Start the web UI: - It is both an analysis tool and an execution/control tool. - The container/sandbox functionality is a major part of the project, not a side feature. - Many commands assume the presence of Claude local state under the user’s Claude config/projects directories. -- Configuration is persisted to `~/.captain.yaml` via `captain configure`. +- Runtime defaults are persisted to `~/.captain.yaml`; validated provider tokens are persisted separately to the private `~/.config/captain/vault` file. diff --git a/Taskfile.yaml b/Taskfile.yaml index 4e2e7ca..13f1d21 100644 --- a/Taskfile.yaml +++ b/Taskfile.yaml @@ -82,6 +82,11 @@ tasks: cmds: - go mod tidy + models:update: + desc: Regenerate the embedded LLM model registry from models.dev + local patches + cmds: + - go run ./pkg/ai/internal/gen-model-registry --patches pkg/ai/internal/gen-model-registry/patches.json --output pkg/api/registry/models.json + clean: desc: Remove build artifacts cmds: diff --git a/cmd/captain/main.go b/cmd/captain/main.go index 2a15202..0ccec63 100644 --- a/cmd/captain/main.go +++ b/cmd/captain/main.go @@ -35,6 +35,7 @@ func main() { } clicky.BindAllFlags(rootCmd.PersistentFlags(), "format") + cli.BindDatabaseURLFlag(rootCmd.PersistentFlags()) // Bind commons' -P/--properties flag so per-subsystem log levels and HTTP // wire logging can be toggled, e.g. -Plog.level.http=trace3. properties.BindFlags(rootCmd.PersistentFlags()) @@ -58,9 +59,9 @@ func main() { } rootCmd.AddCommand(historyAlias) - infoCmd := clicky.AddNamedCommand("info", rootCmd, cli.InfoOptions{}, cli.RunInfo) - infoCmd.Short = "Show current Claude Code session and project info" - infoCmd.Long = "Display metadata about the current Claude Code session including project root, active session ID, and configuration." + infoCmd := clicky.AddNamedCommandWithContext("info", rootCmd, cli.InfoOptions{}, cli.RunInfo) + infoCmd.Short = "Show the current agent session or discovered project sessions" + infoCmd.Long = "Detect the current Codex, Claude, Gemini, or Captain session from its environment, including matching Captain database sessions. Explicit discovery flags show project session history instead." costCmd := clicky.AddNamedCommand("cost", rootCmd, cli.CostOptions{}, cli.RunCost) costCmd.Short = "Show token usage and estimated costs" @@ -80,6 +81,11 @@ func main() { clicky.AddNamedCommandWithContext("live", sessionsCmd, cli.SessionLiveOptions{}, cli.RunSessionLive).Short = "List sessions with live process health" clicky.AddNamedCommandWithContext("get", sessionsCmd, cli.SessionGetOptions{}, cli.RunSessionGet).Short = "Show a session transcript" + psCmd := clicky.AddNamedCommandWithContext("ps", rootCmd, cli.PSOptions{}, cli.RunPS) + psCmd.Short = "List live agent sessions (claude/codex) with session id, agents, and cmux surface" + psCmd.Long = "Detect running claude/codex agent processes (via ps + lsof), resolve each one's session id, sub-agent ids, cmux surface (from CMUX_* env vars), and last activity, then augment with cached token/cost/context data. Only currently-active sessions are listed." + clicky.AddNamedCommandWithContext("ps", sessionsCmd, cli.PSOptions{}, cli.RunPS).Short = psCmd.Short + sandboxCmd := &cobra.Command{ Use: "sandbox", Short: "Sandbox configuration tools", @@ -101,7 +107,12 @@ func main() { "also include request/response bodies.", } rootCmd.AddCommand(aiCmd) - clicky.AddNamedCommand("prompt", aiCmd, cli.AIPromptOptions{}, cli.RunAIPrompt) + aiCmd.AddCommand(cli.NewCommandAlias(cli.CommandAliasOptions{ + Name: "prompt", + Short: "Alias for captain prompt run", + Root: rootCmd, + Target: []string{"prompt", "run"}, + })) clicky.AddNamedCommand("agent", aiCmd, cli.AIAgentOptions{}, cli.RunAIAgent).Short = "Run an iterative agent with verifiers, worktree, and commit" clicky.AddNamedCommand("models", aiCmd, cli.AIModelsOptions{}, cli.RunAIModels) clicky.AddNamedCommand("test", aiCmd, cli.AITestOptions{}, cli.RunAITest) @@ -109,14 +120,19 @@ func main() { whoamiCmd := clicky.AddNamedCommand("whoami", rootCmd, cli.WhoamiOptions{}, cli.RunWhoami) whoamiCmd.Short = "List agent adapters, auth methods, and available models" - whoamiCmd.Long = "Show every AI agent adapter (API providers and CLI agents), how each is authenticated (API-key env var or CLI login), whether its CLI binary is installed, and the models each provider exposes via a live API call. Pass --models=false to skip the network probes, or --backend to inspect a single adapter." + whoamiCmd.Long = "Show every AI agent adapter (API providers and CLI agents), how each is authenticated (Captain vault, API-key env var, or CLI login), whether its CLI binary is installed, and the models each provider exposes via a live API call. Pass --models=false to skip the network probes, or --backend to inspect a single adapter." - configureCmd := clicky.AddNamedCommand("configure", rootCmd, cli.ConfigureOptions{}, cli.RunConfigure) - configureCmd.Short = "Interactive wizard to set default model, backend, budget, and safety toggles" - configureCmd.Long = "Run an interactive form to configure ~/.captain.yaml. These defaults are applied to `captain ai prompt`, `captain ai test`, and other AI commands when corresponding flags are not passed." + configureCmd := clicky.AddNamedCommandWithContext("configure", rootCmd, cli.ConfigureOptions{}, cli.RunConfigure) + configureCmd.Use = "configure [provider]" + configureCmd.Short = "Configure provider defaults or validate and save an API token" + configureCmd.Long = "Run without a provider for the interactive ~/.captain.yaml wizard. Run `captain configure ` to securely prompt for, validate, and save an API token in ~/.config/captain/vault, or pass --agent, --model, --effort, and --active to save provider-specific runtime defaults. Token flags and runtime-default flags cannot be combined. Automation may pass the hidden --token flag, but command-line arguments can be retained in shell history or process listings. Use --test to validate the effective token without saving." rootCmd.AddCommand(cli.NewServeCommand(version)) + attachmentsCmd := &cobra.Command{Use: "attachments", Short: "Manage durable prompt attachments"} + rootCmd.AddCommand(attachmentsCmd) + clicky.AddNamedCommand("gc", attachmentsCmd, cli.AttachmentsGCOptions{}, cli.RunAttachmentsGC).Short = "Remove old unreferenced attachments" + dodCmd := &cobra.Command{ Use: "dod", Short: "Definition of Done checks", @@ -149,6 +165,28 @@ func main() { hookCmd.AddCommand(dodHookCmd) clicky.AddNamedCommand("install", dodHookCmd, cli.HookInstallOptions{}, cli.RunDodInstall) + monitorHookCmd := &cobra.Command{Use: "monitor", Short: "Session monitoring hooks (hooks-first session tracking)"} + hookCmd.AddCommand(monitorHookCmd) + monitorNotifyCmd := &cobra.Command{ + Use: "notify", + Short: "Forward one provider hook event to captain serve", + Long: "Hook receiver for session monitoring: Claude Code lifecycle hooks pipe their JSON payload " + + "on stdin, codex notify appends its payload as the final argument. The event is POSTed to the " + + "running captain serve instance ($CAPTAIN_SERVER_URL or http://localhost:9020) with a 1s timeout " + + "and always exits 0 — when serve is unreachable the event is dropped and the daily recon reconciles it.", + Args: cobra.ArbitraryArgs, + RunE: func(cmd *cobra.Command, args []string) error { + provider, _ := cmd.Flags().GetString("provider") + url, _ := cmd.Flags().GetString("url") + return cli.RunHookMonitorNotify(cli.HookMonitorNotifyOptions{Provider: provider, URL: url}, args) + }, + } + monitorNotifyCmd.Flags().String("provider", "claude", "Hook payload provider: claude or codex") + monitorNotifyCmd.Flags().String("url", "", "Captain serve base URL (default $CAPTAIN_SERVER_URL or http://localhost:9020)") + monitorHookCmd.AddCommand(monitorNotifyCmd) + monitorInstallCmd := clicky.AddNamedCommand("install", monitorHookCmd, cli.HookMonitorInstallOptions{}, cli.RunHookMonitorInstall) + monitorInstallCmd.Short = "Install session-monitoring hooks into ~/.claude/settings.json and ~/.codex/config.toml" + projectsCmd := &cobra.Command{Use: "projects", Short: "Manage Claude Code project sessions"} rootCmd.AddCommand(projectsCmd) clicky.AddNamedCommand("list", projectsCmd, cli.ProjectsListOptions{}, cli.RunProjectsList) diff --git a/docs/package.json b/docs/package.json index 96850da..88ae299 100644 --- a/docs/package.json +++ b/docs/package.json @@ -13,10 +13,7 @@ "@astrojs/mdx": "^7.0.2", "@astrojs/react": "^6.0.1", "@flanksource/clicky-ui": "link:../../clicky-ui/packages/ui", - "@flanksource/icons": "^1.0.60", "@mdxeditor/editor": "^4.0.4", - "@radix-ui/react-compose-refs": "^1.1.2", - "@radix-ui/react-slot": "^1.1.1", "@shikijs/langs": "^1.24.0", "@shikijs/themes": "^1.24.0", "@shikijs/transformers": "^1.24.0", @@ -24,10 +21,6 @@ "@tanstack/react-query": "^5.90.12", "ai": "^6.0.199", "astro": "^7.1.3", - "class-variance-authority": "^0.7.1", - "clsx": "^2.1.1", - "dompurify": "^3.4.7", - "jotai": "^2.16.0", "marked": "^17.0.1", "react": "^18.3.1", "react-dom": "^18.3.1", @@ -35,7 +28,6 @@ "recharts": "^3.8.1", "shiki": "^1.24.0", "streamdown": "^2.5.0", - "tailwind-merge": "^2.6.0", "tailwindcss": "^4.3.2" }, "devDependencies": { diff --git a/docs/pnpm-lock.yaml b/docs/pnpm-lock.yaml index 9bfbb93..f659981 100644 --- a/docs/pnpm-lock.yaml +++ b/docs/pnpm-lock.yaml @@ -4,6 +4,9 @@ settings: autoInstallPeers: true excludeLinksFromLockfile: false +overrides: + js-yaml: '>=4.3.0' + importers: .: @@ -17,9 +20,6 @@ importers: '@flanksource/clicky-ui': specifier: link:../../clicky-ui/packages/ui version: link:../../clicky-ui/packages/ui - '@flanksource/icons': - specifier: ^1.0.60 - version: 1.0.63(react@18.3.1) '@mdxeditor/editor': specifier: ^4.0.4 version: 4.0.4(@codemirror/language@6.12.4)(@lezer/highlight@1.2.3)(@types/react-dom@18.3.7(@types/react@18.3.31))(@types/react@18.3.31)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(yjs@13.6.31) @@ -50,18 +50,6 @@ importers: astro: specifier: ^7.1.3 version: 7.1.3(@astrojs/markdown-remark@7.2.1)(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.2)(jiti@2.7.0) - class-variance-authority: - specifier: ^0.7.1 - version: 0.7.1 - clsx: - specifier: ^2.1.1 - version: 2.1.1 - dompurify: - specifier: ^3.4.7 - version: 3.4.11 - jotai: - specifier: ^2.16.0 - version: 2.20.1(@babel/core@7.29.7)(@babel/template@7.29.7)(@types/react@18.3.31)(react@18.3.1) marked: specifier: ^17.0.1 version: 17.0.6 @@ -637,11 +625,6 @@ packages: cpu: [x64] os: [win32] - '@flanksource/icons@1.0.63': - resolution: {integrity: sha512-QJ3N49jltwT4xJVvajr/6V++LqTEYKGamjtLA+2NNECDG8I/guhgH3l984rIoE4+dkYTrUPwvoUh4SiKLic8FQ==} - peerDependencies: - react: '*' - '@floating-ui/core@1.7.5': resolution: {integrity: sha512-1Ih4WTWyw0+lKyFMcBHGbb5U5FtuHJuujoyyr5zTaWS5EYMeT6Jb2AuDeftsCsEuchO+mM2ij5+q9crhydzLhQ==} @@ -1929,9 +1912,6 @@ packages: resolution: {integrity: sha512-77PSwercCZU2Fc4sX94eF8k8Pxte6JAwL4/ICZLFjJLqegs7kCuAsqqj/70NQF6TvDpgFjkubQB2FW2ZZddvQg==} engines: {node: '>=8'} - class-variance-authority@0.7.1: - resolution: {integrity: sha512-Ka+9Trutv7G8M6WT6SeiRWz792K5qEqIGEGzXKhAE6xOWAY6pPH8U+9IY3oCMv6kqTmLsv7Xh/2w2RigkePMsg==} - classnames@2.5.1: resolution: {integrity: sha512-saHYOzhIQs6wy2sVxTM6bUDsQO4F50V9RQ22qBpEdCW+I+/Wmke2HOl6lS6dTpdxVhb88/I6+Hs+438c3lfUow==} @@ -2496,31 +2476,9 @@ packages: resolution: {integrity: sha512-AC/7JofJvZGrrneWNaEnJeOLUx+JlGt7tNa0wZiRPT4MY1wmfKjt2+6O2p2uz2+skll8OZZmJMNqeke7kKbNgQ==} hasBin: true - jotai@2.20.1: - resolution: {integrity: sha512-dnuKfU/GLi8B28RRMjQ3AfoN7kfzP8o41+AX2FmITZqEMY8PHnjABq+VkEooomLwYaGjda+pgy0yFSjaHX/ZPg==} - engines: {node: '>=12.20.0'} - peerDependencies: - '@babel/core': '>=7.0.0' - '@babel/template': '>=7.0.0' - '@types/react': '>=17.0.0' - react: '>=17.0.0' - peerDependenciesMeta: - '@babel/core': - optional: true - '@babel/template': - optional: true - '@types/react': - optional: true - react: - optional: true - js-tokens@4.0.0: resolution: {integrity: sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==} - js-yaml@4.2.0: - resolution: {integrity: sha512-ePWsvanv0DWuDRsW8dnt+R4jQ31SCRCQ7hhNcPXZPsoBZiemuZNYGf7adZdqX2D86j6rvKp3RpCxVTSb8WQlOw==} - hasBin: true - js-yaml@4.3.0: resolution: {integrity: sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==} hasBin: true @@ -4289,10 +4247,6 @@ snapshots: '@esbuild/win32-x64@0.28.1': optional: true - '@flanksource/icons@1.0.63(react@18.3.1)': - dependencies: - react: 18.3.1 - '@floating-ui/core@1.7.5': dependencies: '@floating-ui/utils': 0.2.11 @@ -4754,7 +4708,7 @@ snapshots: cm6-theme-basic-light: 0.2.0(@codemirror/language@6.12.4)(@codemirror/state@6.7.0)(@codemirror/view@6.43.4)(@lezer/highlight@1.2.3) codemirror: 6.0.2 downshift: 7.6.2(react@18.3.1) - js-yaml: 4.2.0 + js-yaml: 4.3.0 lexical: 0.35.0 mdast-util-directive: 3.1.0 mdast-util-from-markdown: 2.0.3 @@ -5780,10 +5734,6 @@ snapshots: ci-info@4.4.0: {} - class-variance-authority@0.7.1: - dependencies: - clsx: 2.1.1 - classnames@2.5.1: {} clsx@2.1.1: {} @@ -6459,19 +6409,8 @@ snapshots: jiti@2.7.0: {} - jotai@2.20.1(@babel/core@7.29.7)(@babel/template@7.29.7)(@types/react@18.3.31)(react@18.3.1): - optionalDependencies: - '@babel/core': 7.29.7 - '@babel/template': 7.29.7 - '@types/react': 18.3.31 - react: 18.3.1 - js-tokens@4.0.0: {} - js-yaml@4.2.0: - dependencies: - argparse: 2.0.1 - js-yaml@4.3.0: dependencies: argparse: 2.0.1 diff --git a/docs/pnpm-workspace.yaml b/docs/pnpm-workspace.yaml new file mode 100644 index 0000000..2245205 --- /dev/null +++ b/docs/pnpm-workspace.yaml @@ -0,0 +1,5 @@ +packages: + - . + +minimumReleaseAge: 10080 +trustPolicy: no-downgrade diff --git a/docs/src/components/PromptApiShapes.tsx b/docs/src/components/PromptApiShapes.tsx index ed0da08..4d3e749 100644 --- a/docs/src/components/PromptApiShapes.tsx +++ b/docs/src/components/PromptApiShapes.tsx @@ -1,6 +1,37 @@ import { Badge, CodeBlock, JsonView, KeyValueList, MethodBadge } from "@flanksource/clicky-ui/data"; import { ClickyProviders } from "./ClickyProviders"; +const commitInputSchema = { + type: "object", + additionalProperties: false, + required: ["patch", "maxBodyLines"], + properties: { + patch: { type: "string", description: "Git patch to summarize" }, + maxBodyLines: { + type: "integer", + description: "Maximum commit-message body lines; zero omits the cap", + }, + }, +}; + +const commitOutputSchema = { + type: "object", + additionalProperties: false, + required: ["type", "subject"], + properties: { + type: { + type: "string", + description: "Conventional commit type: feat|fix|perf|refactor|test|docs|build|ci|chore|revert", + }, + scope: { type: "string", description: "Optional scope, e.g. db, api, fe, kubernetes" }, + subject: { + type: "string", + description: "Imperative subject line, max 100 chars, no trailing period", + }, + body: { type: "string", description: "Optional body explaining why and impact" }, + }, +}; + const promptSummary = { id: "embedded\\u0000embedded\\u0000testdata/commit.prompt", name: "commit", @@ -9,7 +40,15 @@ const promptSummary = { relPath: "testdata/commit.prompt", writable: false, model: "claude-sonnet-4-6", - variables: [{ name: "diff", type: "string" }], + variables: [ + { + name: "maxBodyLines", + type: "integer", + description: "Maximum commit-message body lines; zero omits the cap", + required: true, + }, + { name: "patch", type: "string", description: "Git patch to summarize", required: true }, + ], }; const renderResult = { @@ -17,17 +56,20 @@ const renderResult = { name: "commit", model: "claude-sonnet-4-6", backend: "anthropic", - user: "Diff:\\n...", - system: "You write Conventional Commit messages.", + user: "DIFF INPUT:\\n\\n...\\n\\nREQUIREMENTS:\\n- type: one of feat|fix|perf|refactor|test|docs|build|ci|chore|revert\\n...", + system: "You are a commit message generator. Analyze the diff below and produce a Conventional Commit message.", validationError: "", input: { prompt: { source: "commit.prompt", - user: "Diff:\\n...", - system: "You write Conventional Commit messages.", + user: "DIFF INPUT:\\n\\n...\\n\\nREQUIREMENTS:\\n- type: one of feat|fix|perf|refactor|test|docs|build|ci|chore|revert\\n...", + system: "You are a commit message generator. Analyze the diff below and produce a Conventional Commit message.", + schemaJSON: commitOutputSchema, }, budget: { maxTokens: 1024, timeout: "2h" }, }, + inputSchema: commitInputSchema, + outputSchema: commitOutputSchema, }; export default function PromptApiShapes() { @@ -65,4 +107,3 @@ export default function PromptApiShapes() { ); } - diff --git a/docs/src/components/RuntimeSpecDemo.tsx b/docs/src/components/RuntimeSpecDemo.tsx index 2f42831..7b3803e 100644 --- a/docs/src/components/RuntimeSpecDemo.tsx +++ b/docs/src/components/RuntimeSpecDemo.tsx @@ -8,28 +8,28 @@ const tools: ToolMeta[] = [ name: "Read", label: "Read", group: "Files", - defaultMode: "ask", + defaultPermission: "ask", description: "Read files from the workspace.", }, { name: "Edit", label: "Edit", group: "Files", - defaultMode: "ask", + defaultPermission: "ask", description: "Apply targeted edits.", }, { name: "Bash", label: "Bash", group: "Shell", - defaultMode: "ask", + defaultPermission: "ask", description: "Run shell commands.", }, { name: "WebSearch", label: "Web search", group: "Web", - defaultMode: "disabled", + defaultPermission: "off", description: "Search external documentation.", }, ]; @@ -48,7 +48,6 @@ const initialValue: AISpecRuntimeValue = { prompt: { system: "You are a careful refactoring assistant.", user: "Refactor: {{target}}", - source: "options.prompt", }, permissions: { mode: "acceptEdits", @@ -96,4 +95,3 @@ export default function RuntimeSpecDemo() { ); } - diff --git a/docs/src/pages/prompts/format.mdx b/docs/src/pages/prompts/format.mdx index a91eb9b..9393384 100644 --- a/docs/src/pages/prompts/format.mdx +++ b/docs/src/pages/prompts/format.mdx @@ -37,11 +37,21 @@ A `.prompt` file has YAML frontmatter followed by a Handlebars body. dotprompt o dotprompt metadata Used by the workbench to discover variables and defaults. + + output.schema + dotprompt and Captain providers + Exposed as prompt metadata and attached to the rendered request as Prompt.SchemaJSON. + permissions, memory, budget, setup Captain spec Decoded into `api.Spec` with strict field validation. + + runtimes + Captain execution metadata + Declares the default parallel runtime matrix without becoming part of the request sent to each model. + @@ -98,5 +108,4 @@ Refactor: ## Structured output -Captain does not turn a bare dotprompt `output` schema into a provider schema target. Structured output is attached by Go callers that pass a concrete output type to `Render`; providers derive the JSON schema from that Go type. - +When no Go output target is passed to `Render`, Captain marshals a dotprompt `output.schema` into `Prompt.SchemaJSON` and sends it to the provider unchanged. A concrete Go output target takes precedence and is attached as `Prompt.Schema` for provider-side schema generation. diff --git a/docs/src/pages/prompts/index.mdx b/docs/src/pages/prompts/index.mdx index 7210b4b..95d9990 100644 --- a/docs/src/pages/prompts/index.mdx +++ b/docs/src/pages/prompts/index.mdx @@ -35,7 +35,7 @@ Captain prompts are `.prompt` templates rendered through Google dotprompt and fo -## Minimal prompt +## Structured prompt ## Where to go next @@ -59,4 +98,3 @@ Diff: - [Prompt Format](/prompts/format/) covers frontmatter, roles, variables, and config precedence. - [Runtime Overlay](/prompts/runtime/) covers `api.Spec`, budgets, permissions, memory, model selection, and setup. - [Sources and API](/prompts/sources-api/) covers embedded/local sources and render/run API payloads. - diff --git a/docs/src/pages/prompts/runtime.mdx b/docs/src/pages/prompts/runtime.mdx index 634c510..c620d96 100644 --- a/docs/src/pages/prompts/runtime.mdx +++ b/docs/src/pages/prompts/runtime.mdx @@ -74,3 +74,21 @@ The prompt file provides the base request, but it is not the last word. Captain The prompt workbench keeps runtime selection catalog-backed. The UI groups by provider family first, then mode/backend, then the model list filtered to that backend. Prompt files can still name `model` and `backend` directly when the runtime should be pinned. +## Parallel runtimes + +A prompt can declare its default comparison matrix at the root. Running the file without `-M` dispatches all declared runtimes concurrently, and the workbench opens them as editable runtime rows. Each entry accepts the compact model selector or object form, and a matrix must contain at least two distinct runtimes. + + + +The matrix supplies defaults rather than locking the prompt. Explicit CLI `-M` selectors or HTTP/workbench `runtimes` replace the declared list for that run. A singular `model` is optional when the prompt declares a runtime matrix. diff --git a/docs/src/pages/prompts/sources-api.mdx b/docs/src/pages/prompts/sources-api.mdx index c5b49d0..1febf69 100644 --- a/docs/src/pages/prompts/sources-api.mdx +++ b/docs/src/pages/prompts/sources-api.mdx @@ -25,6 +25,8 @@ Captain discovers prompt sources in this order: Embedded prompts are read-only. Saving an embedded prompt from the workbench writes a local fork, stripping the embedded walk root so `testdata/commit.prompt` becomes `commit.prompt`. +CLI prompt commands accept an opaque prompt ID, an explicit file path, or a bare filename with or without the `.prompt` extension. Bare filenames are resolved across the discovered sources. An existing file in the current directory takes precedence; if more than one discovered source contains the same filename, Captain reports the ambiguity and requires an ID or explicit path. + ## Entity records can never match a + # NULL row) and takes the index, and its insert-time upkeep, to zero. + index "captain_messages_model_call_id_idx" { + columns = [column.model_call_id] + where = "model_call_id IS NOT NULL" + } + + check "captain_messages_sequence_nonnegative" { + expr = "sequence >= 0" + } + check "captain_messages_role_nonempty" { + expr = "length(btrim(role)) > 0" + } + check "captain_messages_schema_version_positive" { + expr = "schema_version > 0" + } + check "captain_messages_source_line_positive" { + expr = "source_line IS NULL OR source_line > 0" + } +} + +table "captain_events" { + schema = schema.public + + column "id" { + null = false + type = uuid + default = sql("gen_random_uuid()") + } + column "session_id" { + null = false + type = uuid + } + column "turn_id" { + null = true + type = uuid + } + column "prompt_run_id" { + null = true + type = uuid + } + column "iteration_id" { + null = true + type = uuid + } + column "model_call_id" { + null = true + type = uuid + } + column "parent_event_id" { + null = true + type = uuid + } + column "event_key" { + null = true + type = text + } + column "stream" { + null = false + type = text + default = "runtime" + } + column "sequence" { + null = true + type = bigint + } + column "kind" { + null = false + type = text + } + column "scope" { + null = false + type = text + default = "session" + } + column "payload" { + null = false + type = jsonb + default = sql("'{}'::jsonb") + } + column "schema_version" { + null = false + type = integer + default = 1 + } + column "occurred_at" { + null = true + type = timestamptz + } + column "recorded_at" { + null = false + type = timestamptz + default = sql("now()") + } + + primary_key { + columns = [column.id] + } + + foreign_key "captain_events_session_id_fkey" { + columns = [column.session_id] + ref_columns = [table.captain_sessions.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + foreign_key "captain_events_turn_id_fkey" { + columns = [column.turn_id] + ref_columns = [table.captain_turns.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + foreign_key "captain_events_prompt_run_id_fkey" { + columns = [column.prompt_run_id] + ref_columns = [table.captain_prompt_runs.column.id] + on_update = NO_ACTION + on_delete = NO_ACTION + } + foreign_key "captain_events_iteration_id_fkey" { + columns = [ + column.prompt_run_id, + column.iteration_id, + ] + ref_columns = [ + table.captain_prompt_run_iterations.column.prompt_run_id, + table.captain_prompt_run_iterations.column.id, + ] + on_update = NO_ACTION + on_delete = NO_ACTION + } + foreign_key "captain_events_model_call_id_fkey" { + columns = [column.model_call_id] + ref_columns = [table.captain_model_calls.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + foreign_key "captain_events_parent_event_id_fkey" { + columns = [column.parent_event_id] + ref_columns = [table.captain_events.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + + index "captain_events_event_key" { + unique = true + columns = [column.session_id, column.event_key] + where = "event_key IS NOT NULL" + } + index "captain_events_stream_sequence_key" { + unique = true + columns = [column.session_id, column.stream, column.sequence] + where = "sequence IS NOT NULL" + } + index "captain_events_turn_id_idx" { + columns = [column.turn_id] + } + index "captain_events_prompt_run_id_idx" { + columns = [column.prompt_run_id] + } + index "captain_events_iteration_id_idx" { + columns = [column.iteration_id] + } + index "captain_events_model_call_id_idx" { + columns = [column.model_call_id] + } + index "captain_events_kind_recorded_at_idx" { + columns = [column.kind, column.recorded_at] + } + + check "captain_events_sequence_nonnegative" { + expr = "sequence IS NULL OR sequence >= 0" + } + check "captain_events_schema_version_positive" { + expr = "schema_version > 0" + } + check "captain_events_parent_not_self" { + expr = "parent_event_id IS NULL OR parent_event_id <> id" + } + check "captain_events_iteration_has_run" { + expr = "iteration_id IS NULL OR prompt_run_id IS NOT NULL" + } +} + +table "captain_turn_requests" { + schema = schema.public + + column "id" { + null = false + type = uuid + default = sql("gen_random_uuid()") + } + column "session_id" { + null = false + type = uuid + } + column "turn_id" { + null = true + type = uuid + } + column "prompt_run_id" { + null = true + type = uuid + } + column "plan_id" { + null = true + type = uuid + } + column "model_call_id" { + null = true + type = uuid + } + column "tool_call_id" { + null = true + type = text + } + column "kind" { + null = false + type = enum.captain_turn_request_kind + } + column "state" { + null = false + type = enum.captain_turn_request_state + default = "pending" + } + column "request" { + null = false + type = jsonb + default = sql("'{}'::jsonb") + } + column "response" { + null = true + type = jsonb + } + column "idempotency_key" { + null = true + type = text + } + column "requested_by" { + null = true + type = text + } + column "resolved_by" { + null = true + type = text + } + column "reason" { + null = true + type = text + } + column "version" { + null = false + type = bigint + default = 0 + } + column "expires_at" { + null = true + type = timestamptz + } + column "created_at" { + null = false + type = timestamptz + default = sql("now()") + } + column "resolved_at" { + null = true + type = timestamptz + } + + primary_key { + columns = [column.id] + } + + foreign_key "captain_turn_requests_session_id_fkey" { + columns = [column.session_id] + ref_columns = [table.captain_sessions.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + foreign_key "captain_turn_requests_turn_id_fkey" { + columns = [column.turn_id] + ref_columns = [table.captain_turns.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + foreign_key "captain_turn_requests_prompt_run_id_fkey" { + columns = [column.prompt_run_id] + ref_columns = [table.captain_prompt_runs.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + foreign_key "captain_turn_requests_plan_id_fkey" { + columns = [column.plan_id] + ref_columns = [table.captain_plans.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + foreign_key "captain_turn_requests_model_call_id_fkey" { + columns = [column.model_call_id] + ref_columns = [table.captain_model_calls.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + + index "captain_turn_requests_idempotency_key" { + unique = true + columns = [column.session_id, column.idempotency_key] + where = "idempotency_key IS NOT NULL" + } + index "captain_turn_requests_pending_session_idx" { + columns = [column.session_id, column.kind, column.created_at] + where = "state = 'pending'" + } + index "captain_turn_requests_turn_id_idx" { + columns = [column.turn_id] + } + index "captain_turn_requests_prompt_run_id_idx" { + columns = [column.prompt_run_id] + } + + check "captain_turn_requests_version_nonnegative" { + expr = "version >= 0" + } + check "captain_turn_requests_resolution" { + expr = "(state = 'pending' AND resolved_at IS NULL) OR (state <> 'pending' AND resolved_at IS NOT NULL)" + } + check "captain_turn_requests_time_order" { + expr = "resolved_at IS NULL OR resolved_at >= created_at" + } +} diff --git a/migrations/40_artifacts.pg.hcl b/migrations/40_artifacts.pg.hcl new file mode 100644 index 0000000..ec5fa89 --- /dev/null +++ b/migrations/40_artifacts.pg.hcl @@ -0,0 +1,186 @@ +table "captain_artifacts" { + schema = schema.public + + column "id" { + null = false + type = uuid + default = sql("gen_random_uuid()") + } + column "session_id" { + null = false + type = uuid + } + column "turn_id" { + null = true + type = uuid + } + column "model_call_id" { + null = true + type = uuid + } + column "prompt_run_id" { + null = true + type = uuid + } + column "kind" { + null = false + type = text + } + column "path" { + null = true + type = text + } + column "digest" { + null = true + type = text + } + column "content_type" { + null = true + type = text + } + column "metadata" { + null = false + type = jsonb + default = sql("'{}'::jsonb") + } + column "occurred_at" { + null = true + type = timestamptz + } + column "created_at" { + null = false + type = timestamptz + default = sql("now()") + } + + primary_key { + columns = [column.id] + } + + foreign_key "captain_artifacts_session_id_fkey" { + columns = [column.session_id] + ref_columns = [table.captain_sessions.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + foreign_key "captain_artifacts_turn_id_fkey" { + columns = [column.turn_id] + ref_columns = [table.captain_turns.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + foreign_key "captain_artifacts_model_call_id_fkey" { + columns = [column.model_call_id] + ref_columns = [table.captain_model_calls.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + foreign_key "captain_artifacts_prompt_run_id_fkey" { + columns = [column.prompt_run_id] + ref_columns = [table.captain_prompt_runs.column.id] + on_update = NO_ACTION + on_delete = SET_NULL + } + + index "captain_artifacts_session_kind_idx" { + columns = [column.session_id, column.kind, column.occurred_at] + } + index "captain_artifacts_turn_id_idx" { + columns = [column.turn_id] + } + index "captain_artifacts_model_call_id_idx" { + columns = [column.model_call_id] + } + index "captain_artifacts_prompt_run_id_idx" { + columns = [column.prompt_run_id] + } + index "captain_artifacts_digest_idx" { + columns = [column.digest] + where = "digest IS NOT NULL" + } +} + +table "captain_session_sources" { + schema = schema.public + + column "id" { + null = false + type = uuid + default = sql("gen_random_uuid()") + } + column "session_id" { + null = false + type = uuid + } + column "source_kind" { + null = false + type = text + } + column "path" { + null = false + type = text + } + column "source_identity" { + null = true + type = text + } + column "parser_version" { + null = false + type = integer + default = 1 + } + column "byte_offset" { + null = false + type = bigint + default = 0 + } + column "observed_size" { + null = false + type = bigint + default = 0 + } + column "observed_mod_time" { + null = true + type = timestamptz + } + column "last_event_key" { + null = true + type = text + } + column "created_at" { + null = false + type = timestamptz + default = sql("now()") + } + column "updated_at" { + null = false + type = timestamptz + default = sql("now()") + } + + primary_key { + columns = [column.id] + } + + foreign_key "captain_session_sources_session_id_fkey" { + columns = [column.session_id] + ref_columns = [table.captain_sessions.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + + index "captain_session_sources_source_path_key" { + unique = true + columns = [column.source_kind, column.path] + } + index "captain_session_sources_session_id_idx" { + columns = [column.session_id] + } + + check "captain_session_sources_parser_version_positive" { + expr = "parser_version > 0" + } + check "captain_session_sources_offsets_nonnegative" { + expr = "byte_offset >= 0 AND observed_size >= 0" + } +} diff --git a/migrations/50_constraints.sql b/migrations/50_constraints.sql new file mode 100644 index 0000000..cfdb766 --- /dev/null +++ b/migrations/50_constraints.sql @@ -0,0 +1,39 @@ +-- phase: post + +-- Paired provenance foreign keys form intentional cycles across runs, plans, +-- iterations, calls and events. Deferring their NO ACTION checks allows a +-- complete session aggregate to cascade in one statement while a standalone +-- deletion of referenced provenance still fails at transaction commit. Atlas +-- OSS does not model this PostgreSQL constraint property, so the post-HCL SQL +-- phase owns it. +-- +-- This script is hash-gated: it re-runs only when its content changes. If an +-- HCL change recreates one of these constraints, bump this script so the +-- deferrable property is restored. +ALTER TABLE public.captain_prompt_runs + ALTER CONSTRAINT captain_prompt_runs_input_plan_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_prompt_runs + ALTER CONSTRAINT captain_prompt_runs_input_plan_revision_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_plans + ALTER CONSTRAINT captain_plans_source_prompt_run_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_plans + ALTER CONSTRAINT captain_plans_source_iteration_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_plans + ALTER CONSTRAINT captain_plans_approved_revision_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_model_calls + ALTER CONSTRAINT captain_model_calls_prompt_run_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_model_calls + ALTER CONSTRAINT captain_model_calls_iteration_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_events + ALTER CONSTRAINT captain_events_prompt_run_id_fkey + DEFERRABLE INITIALLY DEFERRED; +ALTER TABLE public.captain_events + ALTER CONSTRAINT captain_events_iteration_id_fkey + DEFERRABLE INITIALLY DEFERRED; diff --git a/migrations/51_state_triggers.sql b/migrations/51_state_triggers.sql new file mode 100644 index 0000000..3e59d16 --- /dev/null +++ b/migrations/51_state_triggers.sql @@ -0,0 +1,320 @@ +-- phase: post + +CREATE OR REPLACE FUNCTION public.captain_set_updated_at() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +BEGIN + NEW.updated_at := clock_timestamp(); + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_set_session_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +DECLARE + observed_at timestamptz := clock_timestamp(); +BEGIN + IF TG_OP = 'UPDATE' THEN + IF ROW( + NEW.lifecycle_status, + NEW.activity_state, + NEW.health_state, + NEW.state_reason + ) IS DISTINCT FROM ROW( + OLD.lifecycle_status, + OLD.activity_state, + OLD.health_state, + OLD.state_reason + ) THEN + NEW.state_version := OLD.state_version + 1; + NEW.state_observed_at := observed_at; + ELSE + NEW.state_version := OLD.state_version; + NEW.state_observed_at := OLD.state_observed_at; + END IF; + END IF; + + IF NEW.lifecycle_status = 'running' THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.created_at, observed_at); + NEW.ended_at := NULL; + ELSIF NEW.lifecycle_status IN ('succeeded', 'partial', 'failed', 'cancelled', 'interrupted') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.ended_at, NEW.created_at, observed_at); + NEW.ended_at := COALESCE(NEW.ended_at, observed_at); + END IF; + + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_set_prompt_run_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +DECLARE + observed_at timestamptz := clock_timestamp(); +BEGIN + IF TG_OP = 'UPDATE' THEN + IF ROW( + NEW.phase, + NEW.state, + NEW.current_iteration, + NEW.execution_session_id, + NEW.rendered_spec, + NEW.result_text, + NEW.result_json, + NEW.error + ) IS DISTINCT FROM ROW( + OLD.phase, + OLD.state, + OLD.current_iteration, + OLD.execution_session_id, + OLD.rendered_spec, + OLD.result_text, + OLD.result_json, + OLD.error + ) THEN + NEW.version := OLD.version + 1; + ELSE + NEW.version := OLD.version; + END IF; + END IF; + + IF NEW.state IN ('running', 'waiting') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.created_at, observed_at); + NEW.finished_at := NULL; + ELSIF NEW.state IN ('succeeded', 'failed', 'cancelled') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.finished_at, NEW.created_at, observed_at); + NEW.finished_at := COALESCE(NEW.finished_at, observed_at); + END IF; + + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_set_prompt_iteration_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +DECLARE + observed_at timestamptz := clock_timestamp(); +BEGIN + IF NEW.state = 'running' THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.created_at, observed_at); + NEW.finished_at := NULL; + ELSIF NEW.state IN ('succeeded', 'failed', 'cancelled') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.finished_at, NEW.created_at, observed_at); + NEW.finished_at := COALESCE(NEW.finished_at, observed_at); + END IF; + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_set_turn_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +DECLARE + observed_at timestamptz := clock_timestamp(); +BEGIN + IF NEW.status = 'open' THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.created_at, observed_at); + NEW.ended_at := NULL; + ELSIF NEW.status IN ('ended', 'error', 'interrupted') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.ended_at, NEW.created_at, observed_at); + NEW.ended_at := COALESCE(NEW.ended_at, observed_at); + END IF; + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_set_model_call_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +DECLARE + observed_at timestamptz := clock_timestamp(); +BEGIN + IF NEW.status = 'running' THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.created_at, observed_at); + NEW.ended_at := NULL; + ELSIF NEW.status IN ('succeeded', 'failed', 'cancelled') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.ended_at, NEW.created_at, observed_at); + NEW.ended_at := COALESCE(NEW.ended_at, observed_at); + END IF; + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_set_turn_request_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +BEGIN + IF TG_OP = 'UPDATE' THEN + IF ROW( + NEW.state, + NEW.response, + NEW.resolved_by, + NEW.reason + ) IS DISTINCT FROM ROW( + OLD.state, + OLD.response, + OLD.resolved_by, + OLD.reason + ) THEN + NEW.version := OLD.version + 1; + ELSE + NEW.version := OLD.version; + END IF; + END IF; + + IF NEW.state = 'pending' THEN + NEW.resolved_at := NULL; + ELSE + NEW.resolved_at := COALESCE(NEW.resolved_at, clock_timestamp()); + IF NEW.resolved_at < NEW.created_at THEN + NEW.created_at := NEW.resolved_at; + END IF; + END IF; + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_set_plan_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +BEGIN + IF TG_OP = 'INSERT' THEN + IF NEW.approval_state IN ('approved', 'rejected') THEN + NEW.approval_created_at := COALESCE(NEW.approval_created_at, clock_timestamp()); + ELSIF NEW.approval_state = 'revision_requested' THEN + NEW.feedback_at := COALESCE(NEW.feedback_at, clock_timestamp()); + END IF; + ELSIF NEW.approval_state IS DISTINCT FROM OLD.approval_state THEN + IF NEW.approval_state IN ('approved', 'rejected') THEN + NEW.approval_created_at := COALESCE(NEW.approval_created_at, clock_timestamp()); + ELSIF NEW.approval_state = 'revision_requested' THEN + NEW.feedback_at := COALESCE(NEW.feedback_at, clock_timestamp()); + END IF; + END IF; + RETURN NEW; +END; +$$; + +CREATE OR REPLACE FUNCTION public.captain_sync_prompt_run_iteration() +RETURNS trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, public +AS $$ +BEGIN + UPDATE public.captain_prompt_runs + SET current_iteration = GREATEST(current_iteration, NEW.iteration) + WHERE id = NEW.prompt_run_id + AND current_iteration < NEW.iteration; + RETURN NEW; +END; +$$; + +DROP TRIGGER IF EXISTS captain_sessions_state_before ON public.captain_sessions; +CREATE TRIGGER captain_sessions_state_before +BEFORE INSERT OR UPDATE ON public.captain_sessions +FOR EACH ROW EXECUTE FUNCTION public.captain_set_session_state(); + +DROP TRIGGER IF EXISTS captain_prompt_runs_state_before ON public.captain_prompt_runs; +CREATE TRIGGER captain_prompt_runs_state_before +BEFORE INSERT OR UPDATE ON public.captain_prompt_runs +FOR EACH ROW EXECUTE FUNCTION public.captain_set_prompt_run_state(); + +DROP TRIGGER IF EXISTS captain_prompt_run_iterations_state_before ON public.captain_prompt_run_iterations; +CREATE TRIGGER captain_prompt_run_iterations_state_before +BEFORE INSERT OR UPDATE ON public.captain_prompt_run_iterations +FOR EACH ROW EXECUTE FUNCTION public.captain_set_prompt_iteration_state(); + +DROP TRIGGER IF EXISTS captain_turns_state_before ON public.captain_turns; +CREATE TRIGGER captain_turns_state_before +BEFORE INSERT OR UPDATE ON public.captain_turns +FOR EACH ROW EXECUTE FUNCTION public.captain_set_turn_state(); + +DROP TRIGGER IF EXISTS captain_model_calls_state_before ON public.captain_model_calls; +CREATE TRIGGER captain_model_calls_state_before +BEFORE INSERT OR UPDATE ON public.captain_model_calls +FOR EACH ROW EXECUTE FUNCTION public.captain_set_model_call_state(); + +DROP TRIGGER IF EXISTS captain_turn_requests_state_before ON public.captain_turn_requests; +CREATE TRIGGER captain_turn_requests_state_before +BEFORE INSERT OR UPDATE ON public.captain_turn_requests +FOR EACH ROW EXECUTE FUNCTION public.captain_set_turn_request_state(); + +DROP TRIGGER IF EXISTS captain_plans_state_before ON public.captain_plans; +CREATE TRIGGER captain_plans_state_before +BEFORE INSERT OR UPDATE ON public.captain_plans +FOR EACH ROW EXECUTE FUNCTION public.captain_set_plan_state(); + +DROP TRIGGER IF EXISTS captain_sessions_updated_at_before ON public.captain_sessions; +CREATE TRIGGER captain_sessions_updated_at_before +BEFORE UPDATE ON public.captain_sessions +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_session_processes_updated_at_before ON public.captain_session_processes; +CREATE TRIGGER captain_session_processes_updated_at_before +BEFORE UPDATE ON public.captain_session_processes +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_prompt_runs_updated_at_before ON public.captain_prompt_runs; +CREATE TRIGGER captain_prompt_runs_updated_at_before +BEFORE UPDATE ON public.captain_prompt_runs +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_prompt_run_iterations_updated_at_before ON public.captain_prompt_run_iterations; +CREATE TRIGGER captain_prompt_run_iterations_updated_at_before +BEFORE UPDATE ON public.captain_prompt_run_iterations +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_plans_updated_at_before ON public.captain_plans; +CREATE TRIGGER captain_plans_updated_at_before +BEFORE UPDATE ON public.captain_plans +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_turns_updated_at_before ON public.captain_turns; +CREATE TRIGGER captain_turns_updated_at_before +BEFORE UPDATE ON public.captain_turns +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_model_calls_updated_at_before ON public.captain_model_calls; +CREATE TRIGGER captain_model_calls_updated_at_before +BEFORE UPDATE ON public.captain_model_calls +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_session_sources_updated_at_before ON public.captain_session_sources; +CREATE TRIGGER captain_session_sources_updated_at_before +BEFORE UPDATE ON public.captain_session_sources +FOR EACH ROW EXECUTE FUNCTION public.captain_set_updated_at(); + +DROP TRIGGER IF EXISTS captain_prompt_run_iterations_sync_after ON public.captain_prompt_run_iterations; +CREATE TRIGGER captain_prompt_run_iterations_sync_after +AFTER INSERT OR UPDATE OF iteration, prompt_run_id ON public.captain_prompt_run_iterations +FOR EACH ROW EXECUTE FUNCTION public.captain_sync_prompt_run_iteration(); + +-- Internal trigger functions are not PostgREST RPC endpoints. +REVOKE ALL ON FUNCTION public.captain_set_updated_at() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_set_session_state() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_set_prompt_run_state() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_set_prompt_iteration_state() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_set_turn_state() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_set_model_call_state() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_set_turn_request_state() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_set_plan_state() FROM PUBLIC; +REVOKE ALL ON FUNCTION public.captain_sync_prompt_run_iteration() FROM PUBLIC; diff --git a/migrations/52_session_activity_triggers.sql b/migrations/52_session_activity_triggers.sql new file mode 100644 index 0000000..92fce49 --- /dev/null +++ b/migrations/52_session_activity_triggers.sql @@ -0,0 +1,134 @@ +-- phase: post + +-- Remove the retired delivery queue and every producer attached by earlier +-- Captain migrations. These statements also make upgrades reclaim existing +-- outbox storage instead of merely omitting the table from fresh installs. +DROP TABLE IF EXISTS public.captain_outbox; + +DROP TRIGGER IF EXISTS captain_sessions_emit_after ON public.captain_sessions; +DROP TRIGGER IF EXISTS captain_session_processes_emit_after ON public.captain_session_processes; +DROP TRIGGER IF EXISTS captain_prompt_runs_emit_after ON public.captain_prompt_runs; +DROP TRIGGER IF EXISTS captain_prompt_run_iterations_emit_after ON public.captain_prompt_run_iterations; +DROP TRIGGER IF EXISTS captain_plans_emit_after ON public.captain_plans; +DROP TRIGGER IF EXISTS captain_plan_revisions_emit_after ON public.captain_plan_revisions; +DROP TRIGGER IF EXISTS captain_turns_emit_after ON public.captain_turns; +DROP TRIGGER IF EXISTS captain_model_calls_emit_after ON public.captain_model_calls; +DROP TRIGGER IF EXISTS captain_messages_emit_after ON public.captain_messages; +DROP TRIGGER IF EXISTS captain_events_emit_after ON public.captain_events; +DROP TRIGGER IF EXISTS captain_turn_requests_emit_after ON public.captain_turn_requests; +DROP TRIGGER IF EXISTS captain_artifacts_emit_after ON public.captain_artifacts; +DROP TRIGGER IF EXISTS captain_session_sources_emit_after ON public.captain_session_sources; + +DROP FUNCTION IF EXISTS public.captain_emit_session_change(); +DROP FUNCTION IF EXISTS public.captain_notify_outbox(); + +CREATE OR REPLACE FUNCTION public.captain_touch_session_activity() +RETURNS trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, public +AS $$ +DECLARE + row_data jsonb := to_jsonb(NEW); + session_id_value uuid; + agent_activity_at timestamptz; +BEGIN + -- Historical backfills validate and checksum the final rows themselves. + -- Their archive-derived activity timestamps must remain deterministic. + IF current_setting('captain.suppress_session_change', true) = 'on' THEN + RETURN NEW; + END IF; + + IF pg_trigger_depth() > 1 THEN + RETURN NEW; + END IF; + + CASE TG_TABLE_NAME + WHEN 'captain_prompt_run_iterations' THEN + SELECT r.session_id + INTO session_id_value + FROM public.captain_prompt_runs r + WHERE r.id = NULLIF(row_data ->> 'prompt_run_id', '')::uuid; + WHEN 'captain_model_calls' THEN + SELECT t.session_id + INTO session_id_value + FROM public.captain_turns t + WHERE t.id = NULLIF(row_data ->> 'turn_id', '')::uuid; + ELSE + session_id_value := NULLIF(row_data ->> 'session_id', '')::uuid; + END CASE; + + -- Only timestamps that represent agent work can advance last_activity_at. + -- Write-time fallbacks would make replayed historical sessions look active + -- and the monotonic projection could never correct them. + agent_activity_at := COALESCE( + NULLIF(row_data ->> 'occurred_at', '')::timestamptz, + NULLIF(row_data ->> 'state_observed_at', '')::timestamptz, + NULLIF(row_data ->> 'started_at', '')::timestamptz, + NULLIF(row_data ->> 'resolved_at', '')::timestamptz + ); + + -- Keep this as an allowlist: host telemetry, ingest bookkeeping, and any new + -- table without an explicit activity contract must not affect idle checks. + IF agent_activity_at IS NOT NULL AND TG_OP <> 'DELETE' AND TG_TABLE_NAME IN ( + 'captain_messages', + 'captain_turns', + 'captain_turn_requests', + 'captain_model_calls', + 'captain_events', + 'captain_prompt_runs', + 'captain_prompt_run_iterations', + 'captain_artifacts' + ) THEN + UPDATE public.captain_sessions + SET last_activity_at = GREATEST(last_activity_at, agent_activity_at) + WHERE id = session_id_value + AND (last_activity_at IS NULL OR last_activity_at < agent_activity_at); + END IF; + + RETURN NEW; +END; +$$; + +DROP TRIGGER IF EXISTS captain_prompt_runs_activity_after ON public.captain_prompt_runs; +CREATE TRIGGER captain_prompt_runs_activity_after +AFTER INSERT OR UPDATE ON public.captain_prompt_runs +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +DROP TRIGGER IF EXISTS captain_prompt_run_iterations_activity_after ON public.captain_prompt_run_iterations; +CREATE TRIGGER captain_prompt_run_iterations_activity_after +AFTER INSERT OR UPDATE ON public.captain_prompt_run_iterations +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +DROP TRIGGER IF EXISTS captain_turns_activity_after ON public.captain_turns; +CREATE TRIGGER captain_turns_activity_after +AFTER INSERT OR UPDATE ON public.captain_turns +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +DROP TRIGGER IF EXISTS captain_model_calls_activity_after ON public.captain_model_calls; +CREATE TRIGGER captain_model_calls_activity_after +AFTER INSERT OR UPDATE ON public.captain_model_calls +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +DROP TRIGGER IF EXISTS captain_messages_activity_after ON public.captain_messages; +CREATE TRIGGER captain_messages_activity_after +AFTER INSERT OR UPDATE ON public.captain_messages +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +DROP TRIGGER IF EXISTS captain_events_activity_after ON public.captain_events; +CREATE TRIGGER captain_events_activity_after +AFTER INSERT OR UPDATE ON public.captain_events +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +DROP TRIGGER IF EXISTS captain_turn_requests_activity_after ON public.captain_turn_requests; +CREATE TRIGGER captain_turn_requests_activity_after +AFTER INSERT OR UPDATE ON public.captain_turn_requests +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +DROP TRIGGER IF EXISTS captain_artifacts_activity_after ON public.captain_artifacts; +CREATE TRIGGER captain_artifacts_activity_after +AFTER INSERT OR UPDATE ON public.captain_artifacts +FOR EACH ROW EXECUTE FUNCTION public.captain_touch_session_activity(); + +-- Internal trigger functions are not PostgREST RPC endpoints. +REVOKE ALL ON FUNCTION public.captain_touch_session_activity() FROM PUBLIC; diff --git a/migrations/60_view_session_overview.sql b/migrations/60_view_session_overview.sql new file mode 100644 index 0000000..e2db594 --- /dev/null +++ b/migrations/60_view_session_overview.sql @@ -0,0 +1,211 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_overview +WITH (security_barrier = true) +AS +SELECT + s.id, + s.provider_session_id, + s.source, + s.provider, + s.host_id, + s.parent_session_id, + s.root_session_id, + COALESCE(s.root_session_id, s.id) AS thread_id, + s.path, + COALESCE(source.path, s.path) AS history_file, + source.source_kind, + source.parser_version, + s.project, + s.cwd, + s.title, + s.initial_prompt, + s.slug, + s.agent_type, + s.description, + s.cli_version, + s.lifecycle_status, + s.activity_state, + s.health_state, + s.state_reason, + s.state_version, + s.state_observed_at, + s.git, + s.metadata, + s.started_at, + s.ended_at, + s.last_activity_at, + s.created_at, + s.updated_at, + CASE + WHEN s.started_at IS NULL THEN NULL + ELSE EXTRACT(EPOCH FROM COALESCE(s.ended_at, clock_timestamp()) - s.started_at) + END AS duration_seconds, + process.id AS process_id, + process.pid, + process.status AS process_status, + process.command, + process.cwd AS process_cwd, + process.surface, + process.process_started_at, + process.last_heartbeat_at, + process.lease_owner, + process.lease_expires_at, + process.ended_at AS process_ended_at, + process.id IS NOT NULL AND process.ended_at IS NULL AS process_active, + COALESCE(message_stats.message_count, 0) AS message_count, + COALESCE(message_stats.tool_call_count, 0) AS tool_call_count, + COALESCE(event_stats.event_count, 0) AS event_count, + COALESCE(call_stats.turn_count, 0) AS turn_count, + COALESCE(call_stats.model_call_count, 0) AS model_call_count, + COALESCE(agent_stats.agent_count, 1) AS agent_count, + COALESCE(prompt_stats.prompt_run_count, 0) AS prompt_run_count, + COALESCE(plan_stats.plan_count, 0) AS plan_count, + COALESCE(request_stats.pending_request_count, 0) AS pending_request_count, + COALESCE(request_stats.approved_request_count, 0) AS approved_request_count, + COALESCE(request_stats.denied_request_count, 0) AS denied_request_count, + COALESCE(file_stats.file_read_count, 0) AS file_read_count, + COALESCE(file_stats.file_written_count, 0) AS file_written_count, + latest_call.model, + latest_call.backend, + latest_call.effort, + latest_call.context_tokens, + latest_call.context_window_tokens, + CASE + WHEN latest_call.context_window_tokens > 0 THEN + GREATEST( + 0, + LEAST( + 100, + round( + (1 - latest_call.context_tokens::numeric / latest_call.context_window_tokens::numeric) * 100 + )::integer + ) + ) + ELSE NULL + END AS context_free_percent, + COALESCE(call_stats.input_tokens, 0) AS input_tokens, + COALESCE(call_stats.output_tokens, 0) AS output_tokens, + COALESCE(call_stats.reasoning_tokens, 0) AS reasoning_tokens, + COALESCE(call_stats.cache_read_tokens, 0) AS cache_read_tokens, + COALESCE(call_stats.cache_write_tokens, 0) AS cache_write_tokens, + COALESCE(call_stats.input_tokens, 0) + + COALESCE(call_stats.output_tokens, 0) + + COALESCE(call_stats.cache_read_tokens, 0) + + COALESCE(call_stats.cache_write_tokens, 0) AS total_tokens, + COALESCE(call_stats.cost_usd, 0::numeric) AS cost_usd, + -- Appended after initial release: CREATE OR REPLACE VIEW only allows adding + -- columns at the end, so later additions must stay below this line. + process.source AS process_source, + process.cpu_percent, + process.memory_percent, + process.memory_rss_bytes, + process.sampled_at AS process_sampled_at +FROM public.captain_sessions s +LEFT JOIN LATERAL ( + SELECT p.* + FROM public.captain_session_processes p + WHERE p.session_id = s.id + AND p.ended_at IS NULL + ORDER BY COALESCE(p.last_heartbeat_at, p.process_started_at) DESC, p.id DESC + LIMIT 1 +) process ON true +LEFT JOIN LATERAL ( + SELECT src.path, src.source_kind, src.parser_version + FROM public.captain_session_sources src + WHERE src.session_id = s.id + ORDER BY src.updated_at DESC, src.id DESC + LIMIT 1 +) source ON true +LEFT JOIN LATERAL ( + SELECT + count(*)::bigint AS message_count, + COALESCE(sum(( + SELECT count(*) + FROM jsonb_array_elements( + CASE + WHEN jsonb_typeof(m.parts) = 'array' THEN m.parts + ELSE '[]'::jsonb + END + ) part + WHERE ( + part ->> 'type' = 'dynamic-tool' + OR part ->> 'type' LIKE 'tool-%' + ) + AND COALESCE(part ->> 'toolName', '') <> '' + )), 0)::bigint AS tool_call_count + FROM public.captain_messages m + WHERE m.session_id = s.id +) message_stats ON true +LEFT JOIN LATERAL ( + SELECT count(*)::bigint AS event_count + FROM public.captain_events e + WHERE e.session_id = s.id +) event_stats ON true +LEFT JOIN LATERAL ( + SELECT + count(DISTINCT t.id)::bigint AS turn_count, + count(c.id)::bigint AS model_call_count, + COALESCE(sum(c.input_tokens), 0)::bigint AS input_tokens, + COALESCE(sum(c.output_tokens), 0)::bigint AS output_tokens, + COALESCE(sum(c.reasoning_tokens), 0)::bigint AS reasoning_tokens, + COALESCE(sum(c.cache_read_tokens), 0)::bigint AS cache_read_tokens, + COALESCE(sum(c.cache_write_tokens), 0)::bigint AS cache_write_tokens, + COALESCE(sum( + c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + ) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cost_usd + FROM public.captain_turns t + LEFT JOIN public.captain_model_calls c ON c.turn_id = t.id + WHERE t.session_id = s.id +) call_stats ON true +LEFT JOIN LATERAL ( + SELECT + c.model, + c.backend, + c.effort, + c.context_tokens, + c.context_window_tokens + FROM public.captain_turns t + JOIN public.captain_model_calls c ON c.turn_id = t.id + WHERE t.session_id = s.id + ORDER BY COALESCE(c.ended_at, c.started_at, c.created_at) DESC, c.call_index DESC + LIMIT 1 +) latest_call ON true +LEFT JOIN LATERAL ( + SELECT count(*)::bigint + 1 AS agent_count + FROM public.captain_sessions child + WHERE child.root_session_id = s.id +) agent_stats ON true +LEFT JOIN LATERAL ( + SELECT count(*)::bigint AS prompt_run_count + FROM public.captain_prompt_runs r + WHERE r.session_id = s.id +) prompt_stats ON true +LEFT JOIN LATERAL ( + SELECT count(*)::bigint AS plan_count + FROM public.captain_plans p + WHERE p.source_session_id = s.id +) plan_stats ON true +LEFT JOIN LATERAL ( + SELECT + count(*) FILTER (WHERE r.state = 'pending')::bigint AS pending_request_count, + count(*) FILTER (WHERE r.state IN ('approved', 'answered'))::bigint AS approved_request_count, + count(*) FILTER (WHERE r.state = 'denied')::bigint AS denied_request_count + FROM public.captain_turn_requests r + WHERE r.session_id = s.id +) request_stats ON true +LEFT JOIN LATERAL ( + SELECT + count(*) FILTER (WHERE a.kind = 'file.read')::bigint AS file_read_count, + count(*) FILTER (WHERE a.kind IN ('file.write', 'file.edit', 'file.delete'))::bigint AS file_written_count + FROM public.captain_artifacts a + WHERE a.session_id = s.id + AND a.kind LIKE 'file.%' +) file_stats ON true; + +COMMENT ON VIEW public.captain_session_overview IS + 'One row per session for PostgREST list, metadata, health, live-process, usage and cost surfaces.'; diff --git a/migrations/61_view_session_transcript.sql b/migrations/61_view_session_transcript.sql new file mode 100644 index 0000000..445b75c --- /dev/null +++ b/migrations/61_view_session_transcript.sql @@ -0,0 +1,30 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_transcript +WITH (security_barrier = true) +AS +SELECT + m.id, + m.session_id, + m.turn_id, + m.model_call_id, + m.provider_message_id, + m.sequence, + m.role, + m.parts, + m.raw, + m.schema_version, + m.occurred_at, + m.recorded_at, + c.model, + c.backend, + c.effort, + c.status AS model_call_status, + -- Appended after initial release: CREATE OR REPLACE VIEW only allows adding + -- columns at the end, so later additions must stay below this line. + m.source_line +FROM public.captain_messages m +LEFT JOIN public.captain_model_calls c ON c.id = m.model_call_id; + +COMMENT ON VIEW public.captain_session_transcript IS + 'Message rows for the SessionInspector transcript tab; order by sequence through PostgREST.'; diff --git a/migrations/62_view_session_turns.sql b/migrations/62_view_session_turns.sql new file mode 100644 index 0000000..c67a47a --- /dev/null +++ b/migrations/62_view_session_turns.sql @@ -0,0 +1,98 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_turns +WITH (security_barrier = true) +AS +SELECT + t.id, + t.session_id, + t.provider_turn_id, + t.turn_index, + t.description, + t.status, + t.stop_reason, + t.error, + t.started_at, + t.ended_at, + t.created_at, + t.updated_at, + CASE + WHEN t.started_at IS NULL THEN NULL + ELSE EXTRACT(EPOCH FROM COALESCE(t.ended_at, clock_timestamp()) - t.started_at) + END AS duration_seconds, + latest_call.model, + latest_call.backend, + latest_call.effort, + latest_call.context_tokens, + latest_call.context_window_tokens, + CASE + WHEN latest_call.context_window_tokens > 0 THEN + GREATEST( + 0, + LEAST( + 100, + round( + (1 - latest_call.context_tokens::numeric / latest_call.context_window_tokens::numeric) * 100 + )::integer + ) + ) + ELSE NULL + END AS context_free_percent, + COALESCE(call_stats.model_call_count, 0) AS model_call_count, + COALESCE(call_stats.input_tokens, 0) AS input_tokens, + COALESCE(call_stats.output_tokens, 0) AS output_tokens, + COALESCE(call_stats.reasoning_tokens, 0) AS reasoning_tokens, + COALESCE(call_stats.cache_read_tokens, 0) AS cache_read_tokens, + COALESCE(call_stats.cache_write_tokens, 0) AS cache_write_tokens, + COALESCE(call_stats.input_tokens, 0) + + COALESCE(call_stats.output_tokens, 0) + + COALESCE(call_stats.cache_read_tokens, 0) + + COALESCE(call_stats.cache_write_tokens, 0) AS total_tokens, + COALESCE(call_stats.cost_usd, 0::numeric) AS cost_usd, + COALESCE(message_stats.message_count, 0) AS message_count, + COALESCE(message_stats.message_ids, ARRAY[]::uuid[]) AS message_ids, + COALESCE(event_stats.event_count, 0) AS event_count, + COALESCE(event_stats.event_ids, ARRAY[]::uuid[]) AS event_ids +FROM public.captain_turns t +LEFT JOIN LATERAL ( + SELECT + count(*)::bigint AS model_call_count, + COALESCE(sum(c.input_tokens), 0)::bigint AS input_tokens, + COALESCE(sum(c.output_tokens), 0)::bigint AS output_tokens, + COALESCE(sum(c.reasoning_tokens), 0)::bigint AS reasoning_tokens, + COALESCE(sum(c.cache_read_tokens), 0)::bigint AS cache_read_tokens, + COALESCE(sum(c.cache_write_tokens), 0)::bigint AS cache_write_tokens, + COALESCE(sum( + c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + ) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cost_usd + FROM public.captain_model_calls c + WHERE c.turn_id = t.id +) call_stats ON true +LEFT JOIN LATERAL ( + SELECT c.model, c.backend, c.effort, c.context_tokens, c.context_window_tokens + FROM public.captain_model_calls c + WHERE c.turn_id = t.id + ORDER BY COALESCE(c.ended_at, c.started_at, c.created_at) DESC, c.call_index DESC + LIMIT 1 +) latest_call ON true +LEFT JOIN LATERAL ( + SELECT + count(*)::bigint AS message_count, + array_agg(m.id ORDER BY m.sequence) AS message_ids + FROM public.captain_messages m + WHERE m.turn_id = t.id +) message_stats ON true +LEFT JOIN LATERAL ( + SELECT + count(*)::bigint AS event_count, + array_agg(e.id ORDER BY COALESCE(e.occurred_at, e.recorded_at), e.id) AS event_ids + FROM public.captain_events e + WHERE e.turn_id = t.id +) event_stats ON true; + +COMMENT ON VIEW public.captain_session_turns IS + 'Per-turn usage, cost, context and message/event attribution for the SessionInspector turns tab.'; diff --git a/migrations/63_view_session_agents.sql b/migrations/63_view_session_agents.sql new file mode 100644 index 0000000..a898132 --- /dev/null +++ b/migrations/63_view_session_agents.sql @@ -0,0 +1,60 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_agents +WITH (security_barrier = true) +AS +SELECT + s.id, + s.id AS session_id, + s.parent_session_id, + s.root_session_id, + COALESCE(s.root_session_id, s.id) AS thread_id, + s.parent_session_id IS NULL AS is_root, + s.agent_type, + s.description, + s.path AS history_file, + s.source, + s.provider, + s.lifecycle_status, + s.activity_state, + s.health_state, + s.started_at, + s.ended_at, + COALESCE(child_stats.child_count, 0) AS child_count, + COALESCE(call_stats.input_tokens, 0) AS input_tokens, + COALESCE(call_stats.output_tokens, 0) AS output_tokens, + COALESCE(call_stats.reasoning_tokens, 0) AS reasoning_tokens, + COALESCE(call_stats.cache_read_tokens, 0) AS cache_read_tokens, + COALESCE(call_stats.cache_write_tokens, 0) AS cache_write_tokens, + COALESCE(call_stats.input_tokens, 0) + + COALESCE(call_stats.output_tokens, 0) + + COALESCE(call_stats.cache_read_tokens, 0) + + COALESCE(call_stats.cache_write_tokens, 0) AS total_tokens, + COALESCE(call_stats.cost_usd, 0::numeric) AS cost_usd +FROM public.captain_sessions s +LEFT JOIN LATERAL ( + SELECT count(*)::bigint AS child_count + FROM public.captain_sessions child + WHERE child.parent_session_id = s.id +) child_stats ON true +LEFT JOIN LATERAL ( + SELECT + COALESCE(sum(c.input_tokens), 0)::bigint AS input_tokens, + COALESCE(sum(c.output_tokens), 0)::bigint AS output_tokens, + COALESCE(sum(c.reasoning_tokens), 0)::bigint AS reasoning_tokens, + COALESCE(sum(c.cache_read_tokens), 0)::bigint AS cache_read_tokens, + COALESCE(sum(c.cache_write_tokens), 0)::bigint AS cache_write_tokens, + COALESCE(sum( + c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + ) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cost_usd + FROM public.captain_turns t + JOIN public.captain_model_calls c ON c.turn_id = t.id + WHERE t.session_id = s.id +) call_stats ON true; + +COMMENT ON VIEW public.captain_session_agents IS + 'Session hierarchy rows with usage and cost for the SessionInspector agents tab.'; diff --git a/migrations/64_view_session_files.sql b/migrations/64_view_session_files.sql new file mode 100644 index 0000000..55d88da --- /dev/null +++ b/migrations/64_view_session_files.sql @@ -0,0 +1,25 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_files +WITH (security_barrier = true) +AS +SELECT + a.id, + a.session_id, + a.turn_id, + a.model_call_id, + a.prompt_run_id, + a.kind, + split_part(a.kind, '.', 2) AS operation, + a.path, + a.digest, + a.content_type, + a.metadata, + a.occurred_at, + a.created_at +FROM public.captain_artifacts a +WHERE a.kind LIKE 'file.%' + AND a.path IS NOT NULL; + +COMMENT ON VIEW public.captain_session_files IS + 'Normalized file.* artifact rows for the SessionInspector files tab.'; diff --git a/migrations/65_view_session_plans.sql b/migrations/65_view_session_plans.sql new file mode 100644 index 0000000..7d80aa2 --- /dev/null +++ b/migrations/65_view_session_plans.sql @@ -0,0 +1,51 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_plans +WITH (security_barrier = true) +AS +SELECT + p.id, + p.source_session_id AS session_id, + p.source_prompt_run_id, + p.source_iteration_id, + p.source_turn_id, + p.title, + p.slug, + p.path, + p.variant, + p.spec_profile, + p.approval_state, + p.approved_revision_id, + p.approval_comment, + p.approved_by, + p.approval_created_at, + p.feedback_at, + p.created_at, + p.updated_at, + latest_revision.id AS latest_revision_id, + latest_revision.revision AS latest_revision, + latest_revision.plan_markdown AS latest_plan_markdown, + latest_revision.content_hash AS latest_content_hash, + latest_revision.feedback AS latest_feedback, + latest_revision.created_by AS latest_created_by, + latest_revision.created_at AS latest_created_at, + approved_revision.revision AS approved_revision, + approved_revision.plan_markdown AS approved_plan_markdown, + approved_revision.content_hash AS approved_content_hash, + COALESCE(approved_revision.id, latest_revision.id) AS current_revision_id, + COALESCE(approved_revision.revision, latest_revision.revision) AS current_revision, + COALESCE(approved_revision.plan_markdown, latest_revision.plan_markdown) AS plan_markdown, + COALESCE(approved_revision.content_hash, latest_revision.content_hash) AS content_hash +FROM public.captain_plans p +LEFT JOIN LATERAL ( + SELECT r.* + FROM public.captain_plan_revisions r + WHERE r.plan_id = p.id + ORDER BY r.revision DESC, r.created_at DESC, r.id DESC + LIMIT 1 +) latest_revision ON true +LEFT JOIN public.captain_plan_revisions approved_revision + ON approved_revision.id = p.approved_revision_id; + +COMMENT ON VIEW public.captain_session_plans IS + 'Plan rows with latest and approved immutable revisions for the SessionInspector plan tab.'; diff --git a/migrations/66_view_session_approvals.sql b/migrations/66_view_session_approvals.sql new file mode 100644 index 0000000..310d5eb --- /dev/null +++ b/migrations/66_view_session_approvals.sql @@ -0,0 +1,30 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_approvals +WITH (security_barrier = true) +AS +SELECT + r.id, + r.session_id, + r.turn_id, + r.prompt_run_id, + r.plan_id, + r.model_call_id, + r.tool_call_id, + r.kind, + r.state, + r.request, + r.response, + r.idempotency_key, + r.requested_by, + r.resolved_by, + r.reason, + r.version, + r.expires_at, + r.created_at, + r.resolved_at +FROM public.captain_turn_requests r +WHERE r.kind IN ('tool_approval', 'plan_exit_approval'); + +COMMENT ON VIEW public.captain_session_approvals IS + 'Tool and plan-exit approval rows for the SessionInspector approvals tab.'; diff --git a/migrations/67_view_session_costs.sql b/migrations/67_view_session_costs.sql new file mode 100644 index 0000000..40c8e19 --- /dev/null +++ b/migrations/67_view_session_costs.sql @@ -0,0 +1,56 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_costs +WITH (security_barrier = true) +AS +SELECT + concat_ws( + ':', + t.session_id::text, + c.model, + c.backend, + COALESCE(c.effort, 'default'), + upper(c.currency) + ) AS id, + t.session_id, + c.model, + c.backend, + c.effort, + upper(c.currency) AS currency, + count(*)::bigint AS model_call_count, + sum(c.input_tokens)::bigint AS input_tokens, + sum(c.output_tokens)::bigint AS output_tokens, + sum(c.reasoning_tokens)::bigint AS reasoning_tokens, + sum(c.cache_read_tokens)::bigint AS cache_read_tokens, + sum(c.cache_write_tokens)::bigint AS cache_write_tokens, + ( + sum(c.input_tokens) + + sum(c.output_tokens) + + sum(c.cache_read_tokens) + + sum(c.cache_write_tokens) + )::bigint AS total_tokens, + sum(c.input_cost) AS input_cost, + sum(c.output_cost) AS output_cost, + sum(c.reasoning_cost) AS reasoning_cost, + sum(c.cache_read_cost) AS cache_read_cost, + sum(c.cache_write_cost) AS cache_write_cost, + sum( + c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + ) AS total_cost, + min(c.started_at) AS first_call_at, + max(c.ended_at) AS last_call_at +FROM public.captain_turns t +JOIN public.captain_model_calls c ON c.turn_id = t.id +GROUP BY + t.session_id, + c.model, + c.backend, + c.effort, + upper(c.currency); + +COMMENT ON VIEW public.captain_session_costs IS + 'Per-session model/backend/effort token and cost totals for the SessionInspector costs tab.'; diff --git a/migrations/68_view_session_events.sql b/migrations/68_view_session_events.sql new file mode 100644 index 0000000..177e7d7 --- /dev/null +++ b/migrations/68_view_session_events.sql @@ -0,0 +1,26 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_session_events +WITH (security_barrier = true) +AS +SELECT + e.id, + e.session_id, + e.turn_id, + e.prompt_run_id, + e.iteration_id, + e.model_call_id, + e.parent_event_id, + e.event_key, + e.stream, + e.sequence, + e.kind, + e.scope, + e.payload, + e.schema_version, + e.occurred_at, + e.recorded_at +FROM public.captain_events e; + +COMMENT ON VIEW public.captain_session_events IS + 'Durable session and turn event rows for SessionInspector metadata and raw-event surfaces.'; diff --git a/migrations/69_view_prompt_run_overview.sql b/migrations/69_view_prompt_run_overview.sql new file mode 100644 index 0000000..55ae798 --- /dev/null +++ b/migrations/69_view_prompt_run_overview.sql @@ -0,0 +1,115 @@ +-- phase: post + +CREATE OR REPLACE VIEW public.captain_prompt_run_overview +WITH (security_barrier = true) +AS +SELECT + r.id, + r.session_id, + r.root_session_id, + r.batch_id, + r.parent_run_id, + r.input_plan_id, + r.input_plan_revision_id, + r.origin, + r.spec_profile, + r.admission_key, + r.rendered_spec, + r.prompt_markdown, + r.verification_markdown, + r.phase, + r.state, + r.current_iteration, + r.result_text, + r.result_json, + r.error, + r.version, + r.queued_at, + r.started_at, + r.finished_at, + r.created_at, + r.updated_at, + CASE + WHEN r.started_at IS NULL THEN NULL + ELSE EXTRACT(EPOCH FROM COALESCE(r.finished_at, clock_timestamp()) - r.started_at) + END AS duration_seconds, + COALESCE(iteration_stats.iteration_count, 0) AS iteration_count, + COALESCE(iteration_stats.succeeded_iteration_count, 0) AS succeeded_iteration_count, + COALESCE(iteration_stats.failed_iteration_count, 0) AS failed_iteration_count, + latest_iteration.id AS latest_iteration_id, + latest_iteration.iteration AS latest_iteration, + latest_iteration.state AS latest_iteration_state, + latest_iteration.feedback AS latest_iteration_feedback, + latest_iteration.verification_result AS latest_verification_result, + latest_iteration.error AS latest_iteration_error, + latest_iteration.started_at AS latest_iteration_started_at, + latest_iteration.finished_at AS latest_iteration_finished_at, + COALESCE(call_stats.model_call_count, 0) AS model_call_count, + COALESCE(call_stats.input_tokens, 0) AS input_tokens, + COALESCE(call_stats.output_tokens, 0) AS output_tokens, + COALESCE(call_stats.reasoning_tokens, 0) AS reasoning_tokens, + COALESCE(call_stats.cache_read_tokens, 0) AS cache_read_tokens, + COALESCE(call_stats.cache_write_tokens, 0) AS cache_write_tokens, + COALESCE(call_stats.input_tokens, 0) + + COALESCE(call_stats.output_tokens, 0) + + COALESCE(call_stats.cache_read_tokens, 0) + + COALESCE(call_stats.cache_write_tokens, 0) AS total_tokens, + COALESCE(call_stats.cost_usd, 0::numeric) AS cost_usd, + COALESCE(plan_stats.plan_count, 0) AS plan_count, + plan_stats.latest_plan_id, + plan_stats.latest_plan_approval_state, + plan_stats.latest_plan_revision, + r.execution_session_id +FROM public.captain_prompt_runs r +LEFT JOIN LATERAL ( + SELECT + count(*)::bigint AS iteration_count, + count(*) FILTER (WHERE i.state = 'succeeded')::bigint AS succeeded_iteration_count, + count(*) FILTER (WHERE i.state = 'failed')::bigint AS failed_iteration_count + FROM public.captain_prompt_run_iterations i + WHERE i.prompt_run_id = r.id +) iteration_stats ON true +LEFT JOIN LATERAL ( + SELECT i.* + FROM public.captain_prompt_run_iterations i + WHERE i.prompt_run_id = r.id + ORDER BY i.iteration DESC, i.created_at DESC, i.id DESC + LIMIT 1 +) latest_iteration ON true +LEFT JOIN LATERAL ( + SELECT + count(*)::bigint AS model_call_count, + COALESCE(sum(c.input_tokens), 0)::bigint AS input_tokens, + COALESCE(sum(c.output_tokens), 0)::bigint AS output_tokens, + COALESCE(sum(c.reasoning_tokens), 0)::bigint AS reasoning_tokens, + COALESCE(sum(c.cache_read_tokens), 0)::bigint AS cache_read_tokens, + COALESCE(sum(c.cache_write_tokens), 0)::bigint AS cache_write_tokens, + COALESCE(sum( + c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + ) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cost_usd + FROM public.captain_model_calls c + WHERE c.prompt_run_id = r.id +) call_stats ON true +LEFT JOIN LATERAL ( + SELECT + count(*)::bigint AS plan_count, + (array_agg(p.id ORDER BY p.created_at DESC, p.id DESC))[1] AS latest_plan_id, + (array_agg(p.approval_state ORDER BY p.created_at DESC, p.id DESC))[1] AS latest_plan_approval_state, + (array_agg(pr.revision ORDER BY p.created_at DESC, p.id DESC))[1] AS latest_plan_revision + FROM public.captain_plans p + LEFT JOIN LATERAL ( + SELECT revision.revision + FROM public.captain_plan_revisions revision + WHERE revision.plan_id = p.id + ORDER BY revision.revision DESC, revision.created_at DESC, revision.id DESC + LIMIT 1 + ) pr ON true + WHERE p.source_prompt_run_id = r.id +) plan_stats ON true; + +COMMENT ON VIEW public.captain_prompt_run_overview IS + 'Prompt-run control-plane state with iteration, plan, usage, cost and verification summaries.'; diff --git a/migrations/70_prompt_run_runtime.sql b/migrations/70_prompt_run_runtime.sql new file mode 100644 index 0000000..1001398 --- /dev/null +++ b/migrations/70_prompt_run_runtime.sql @@ -0,0 +1,75 @@ +-- phase: post + +UPDATE public.captain_prompt_runs +SET + runtime = CASE + WHEN jsonb_typeof(rendered_spec -> 'runtime') = 'object' THEN rendered_spec -> 'runtime' + ELSE jsonb_strip_nulls(jsonb_build_object( + 'mode', 'run', + 'resolved', jsonb_strip_nulls(jsonb_build_object( + 'backend', NULLIF(rendered_spec ->> 'backend', ''), + 'model', NULLIF(rendered_spec ->> 'model', ''), + 'effort', NULLIF(COALESCE( + rendered_spec #>> '{config,effort}', + rendered_spec #>> '{input,effort}' + ), '') + )) + )) + END, + rendered_spec = rendered_spec - 'runtime' +WHERE runtime = '{}'::jsonb + AND ( + jsonb_typeof(rendered_spec -> 'runtime') = 'object' + OR NULLIF(rendered_spec ->> 'backend', '') IS NOT NULL + OR NULLIF(rendered_spec ->> 'model', '') IS NOT NULL + ); + +CREATE OR REPLACE FUNCTION public.captain_set_prompt_run_state() +RETURNS trigger +LANGUAGE plpgsql +SET search_path = pg_catalog, public +AS $$ +DECLARE + observed_at timestamptz := clock_timestamp(); +BEGIN + IF TG_OP = 'UPDATE' THEN + IF ROW( + NEW.phase, + NEW.state, + NEW.current_iteration, + NEW.execution_session_id, + NEW.rendered_spec, + NEW.runtime, + NEW.result_text, + NEW.result_json, + NEW.error + ) IS DISTINCT FROM ROW( + OLD.phase, + OLD.state, + OLD.current_iteration, + OLD.execution_session_id, + OLD.rendered_spec, + OLD.runtime, + OLD.result_text, + OLD.result_json, + OLD.error + ) THEN + NEW.version := OLD.version + 1; + ELSE + NEW.version := OLD.version; + END IF; + END IF; + + IF NEW.state IN ('running', 'waiting') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.created_at, observed_at); + NEW.finished_at := NULL; + ELSIF NEW.state IN ('succeeded', 'failed', 'cancelled') THEN + NEW.started_at := COALESCE(NEW.started_at, NEW.finished_at, NEW.created_at, observed_at); + NEW.finished_at := COALESCE(NEW.finished_at, observed_at); + END IF; + + RETURN NEW; +END; +$$; + +REVOKE ALL ON FUNCTION public.captain_set_prompt_run_state() FROM PUBLIC; diff --git a/migrations/71_session_storage_params.sql b/migrations/71_session_storage_params.sql new file mode 100644 index 0000000..3321f8a --- /dev/null +++ b/migrations/71_session_storage_params.sql @@ -0,0 +1,33 @@ +-- phase: post + +-- captain_sessions is a heartbeat table: every live session repeatedly rewrites +-- last_activity_at, activity_state, health_state and state_version on a row +-- averaging ~950 bytes across 29 columns, and three row-level triggers fire on +-- each of those updates. At the default fillfactor of 100 a heap page is packed +-- with no room for a new row version, so the heartbeat cannot stay HOT and each +-- update migrates the tuple to a fresh page. Plain VACUUM then frees space +-- inside pages it can never hand back, so the heap only ever ratchets upward. +-- +-- Measured on a developer database before this change: 7,791 live rows spread +-- over 40,458 heap pages (316 MB) with just 8.8% of pages holding any live +-- tuple -- roughly 40x the ~10 MB the live set warrants. It was also the single +-- largest source of physical reads in the database, at 1.47M of 1.76M total +-- block reads and a 25% heap cache hit ratio while every other table sat above +-- 92%. +-- +-- fillfactor reserves in-page room so the heartbeat stays HOT, and the lowered +-- autovacuum scale factors reclaim that space long before a fifth of the table +-- is dead. Atlas OSS does not model PostgreSQL storage parameters, so the +-- post-HCL SQL phase owns them, the same way 50_constraints.sql owns +-- constraint deferrability. +-- +-- This script is hash-gated: it re-runs only when its content changes. If an +-- HCL change recreates captain_sessions, bump this script so the storage +-- parameters are restored. Setting the parameters does not rewrite the table -- +-- existing bloat is reclaimed by a separate VACUUM FULL or pg_repack. +ALTER TABLE public.captain_sessions SET ( + fillfactor = 70, + autovacuum_vacuum_scale_factor = 0.02, + autovacuum_analyze_scale_factor = 0.02, + autovacuum_vacuum_cost_delay = 0 +); diff --git a/migrations/72_ingest_storage_params.sql b/migrations/72_ingest_storage_params.sql new file mode 100644 index 0000000..893faf6 --- /dev/null +++ b/migrations/72_ingest_storage_params.sql @@ -0,0 +1,58 @@ +-- phase: post + +-- The transcript ingest tables need visibility-map upkeep that the shipped +-- autovacuum defaults never deliver, and two of the three are update-churned +-- rather than append-only. +-- +-- captain_messages is written once per row and never updated: insertMessages +-- upserts with ON CONFLICT DO NOTHING. Its default fillfactor of 100 is +-- therefore correct -- there are no updates to keep HOT, and dense pages are +-- what an append-only table wants. What it does need is a far tighter +-- insert-driven vacuum. autovacuum_vacuum_insert_scale_factor defaults to 0.2, +-- so on a measured 714,045-row table roughly 143,000 new rows must accumulate +-- before an insert vacuum fires -- about two weeks at the observed ~10,000 +-- messages/day. Across that window the visibility map rots, and because +-- captain_messages_session_sequence_key is the most-scanned index in the +-- database (2,649,235 scans), every one of those index-only scans quietly +-- degrades into full heap fetches. +-- +-- Measured on a developer database before this change, a 200-session lookup +-- reported Heap Fetches: 28,614 and shared hit=11,517; after a plain VACUUM +-- restored the map the same plan reported Heap Fetches: 2 and shared hit=2,029 +-- -- 5.7x fewer buffers for an identical result. +-- +-- captain_turns and captain_model_calls are a different shape. upsertTurns and +-- the model-call upsert both use ON CONFLICT DO UPDATE, and a live session is +-- re-ingested each time its transcript grows, so every turn already persisted +-- is rewritten on every pass. That is update churn, not appends, and it showed: +-- captain_turns sat at 26.9% and captain_model_calls at 38.6% all-visible page +-- coverage. These two get fillfactor to keep the repeated rewrites HOT, the +-- same reasoning as 71_session_storage_params.sql, plus the tightened insert +-- threshold for map upkeep. +-- +-- Atlas OSS does not model PostgreSQL storage parameters, so the post-HCL SQL +-- phase owns them. This script is hash-gated: it re-runs only when its content +-- changes. If an HCL change recreates any of these tables, bump this script so +-- the storage parameters are restored. Setting the parameters does not rewrite +-- a table or backfill the visibility map -- the first autovacuum pass does. +ALTER TABLE public.captain_messages SET ( + autovacuum_vacuum_insert_scale_factor = 0.02, + autovacuum_analyze_scale_factor = 0.02, + autovacuum_vacuum_cost_delay = 0 +); + +ALTER TABLE public.captain_turns SET ( + fillfactor = 70, + autovacuum_vacuum_scale_factor = 0.02, + autovacuum_vacuum_insert_scale_factor = 0.05, + autovacuum_analyze_scale_factor = 0.02, + autovacuum_vacuum_cost_delay = 0 +); + +ALTER TABLE public.captain_model_calls SET ( + fillfactor = 70, + autovacuum_vacuum_scale_factor = 0.02, + autovacuum_vacuum_insert_scale_factor = 0.05, + autovacuum_analyze_scale_factor = 0.02, + autovacuum_vacuum_cost_delay = 0 +); diff --git a/migrations/73_normalize_session_cwd.sql b/migrations/73_normalize_session_cwd.sql new file mode 100644 index 0000000..0b16c58 --- /dev/null +++ b/migrations/73_normalize_session_cwd.sql @@ -0,0 +1,23 @@ +-- phase: post + +-- FindSessionIDByCWD used to compare rtrim(cwd, '/') = $2 because callers were +-- free to store either spelling of a directory. No index can serve that +-- expression, so the process-to-session heuristic seq-scanned captain_sessions +-- every time an agent process was discovered without a session id in its +-- command line. Measured on a developer database: 2,374 sequential scans had +-- read 18.8M tuples from an 8,000-row table, and a single lookup scanned 8,141 +-- rows to return 133. +-- +-- normalizeCWD now settles the spelling at both write sites (CreateOrGetSession +-- and projectSessionColumns), the query is a plain equality test, and +-- captain_sessions_source_cwd_idx backs it. This backfills the rows written +-- before that, so the equality test cannot miss them. +-- +-- Root is deliberately left alone: rtrim('/', '/') is the empty string, which +-- would erase the directory rather than normalize it. Rows already holding the +-- empty string are equally untouched -- there is nothing to normalize. +UPDATE public.captain_sessions +SET cwd = rtrim(btrim(cwd), '/') +WHERE cwd IS NOT NULL + AND cwd <> rtrim(btrim(cwd), '/') + AND rtrim(btrim(cwd), '/') <> ''; diff --git a/migrations/concurrency_integration_test.go b/migrations/concurrency_integration_test.go new file mode 100644 index 0000000..9546f6c --- /dev/null +++ b/migrations/concurrency_integration_test.go @@ -0,0 +1,81 @@ +package migrations + +import ( + "context" + "os" + "path/filepath" + "sync" + "testing" + "time" + + commonsdb "github.com/flanksource/commons-db/db" + "github.com/stretchr/testify/require" +) + +func TestConcurrentApplySerializesCaptainMigrations(t *testing.T) { + if os.Getenv("CAPTAIN_DB_EMBEDDED_TEST") == "" { + t.Skip("set CAPTAIN_DB_EMBEDDED_TEST=1 to run embedded-postgres migration tests") + } + + dsn, stop, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: filepath.Join(t.TempDir(), "postgres"), + Database: "captain_concurrent_migrations", + }) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, stop()) }) + + // Hold the same session lock before releasing a group of Apply calls. This + // proves every caller enters through the advisory-lock boundary rather than + // racing the Atlas inspect/diff/apply window. + blocker, err := acquireMigrationLock(t.Context(), dsn) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, blocker.Close()) }) + + const callers = 6 + results := make(chan error, callers) + start := make(chan struct{}) + var ready sync.WaitGroup + ready.Add(callers) + for range callers { + go func() { + ready.Done() + <-start + results <- Apply(t.Context(), dsn) + }() + } + ready.Wait() + close(start) + + select { + case err := <-results: + t.Fatalf("Apply completed while the Captain migration lock was held: %v", err) + case <-time.After(250 * time.Millisecond): + } + + require.NoError(t, blocker.Close()) + for range callers { + select { + case err := <-results: + require.NoError(t, err) + case <-time.After(90 * time.Second): + t.Fatal("timed out waiting for serialized Captain migrations") + } + } + + db, err := commonsdb.NewDB(dsn) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, db.Close()) }) + var sessionsTable *string + require.NoError(t, db.QueryRowContext(t.Context(), + `SELECT to_regclass('public.captain_sessions')::text`).Scan(&sessionsTable)) + require.NotNil(t, sessionsTable) + require.Equal(t, "captain_sessions", *sessionsTable) + + // All Apply calls returned, so their session locks must have been released. + // Bound the reacquisition to catch a leaked dedicated connection cleanly. + reacquireCtx, cancel := context.WithTimeout(t.Context(), 2*time.Second) + defer cancel() + reacquired, err := acquireMigrationLock(reacquireCtx, dsn) + require.NoError(t, err) + require.NoError(t, reacquired.Close()) +} diff --git a/migrations/enum_upgrade_integration_test.go b/migrations/enum_upgrade_integration_test.go new file mode 100644 index 0000000..917aef5 --- /dev/null +++ b/migrations/enum_upgrade_integration_test.go @@ -0,0 +1,57 @@ +package migrations + +import ( + "os" + "path/filepath" + + commonsdb "github.com/flanksource/commons-db/db" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Captain migration upgrades", func() { + It("appends partial to the existing session lifecycle enum", func(ctx SpecContext) { + if os.Getenv("CAPTAIN_DB_EMBEDDED_TEST") == "" { + Skip("set CAPTAIN_DB_EMBEDDED_TEST=1 to run embedded-postgres migration tests") + } + + dsn, stop, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: filepath.Join(GinkgoT().TempDir(), "postgres"), + Database: "captain_enum_upgrade", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(stop) + + db, err := commonsdb.NewDB(dsn) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + _, err = db.ExecContext(ctx, `CREATE TYPE public.captain_session_lifecycle_status AS ENUM ( + 'created', 'running', 'succeeded', 'failed', 'cancelled', 'interrupted' + )`) + Expect(err).NotTo(HaveOccurred()) + + Expect(Apply(ctx, dsn)).To(Succeed()) + Expect(Apply(ctx, dsn)).To(Succeed()) + + rows, err := db.QueryContext(ctx, ` + SELECT enumlabel + FROM pg_catalog.pg_enum + JOIN pg_catalog.pg_type ON pg_type.oid = pg_enum.enumtypid + WHERE pg_type.typname = 'captain_session_lifecycle_status' + ORDER BY enumsortorder + `) + Expect(err).NotTo(HaveOccurred()) + defer rows.Close() + + var values []string + for rows.Next() { + var value string + Expect(rows.Scan(&value)).To(Succeed()) + values = append(values, value) + } + Expect(rows.Err()).NotTo(HaveOccurred()) + Expect(values).To(Equal( + []string{"created", "running", "succeeded", "partial", "failed", "cancelled", "interrupted"}, + )) + }) +}) diff --git a/migrations/index_pruning_integration_test.go b/migrations/index_pruning_integration_test.go new file mode 100644 index 0000000..8c40ab8 --- /dev/null +++ b/migrations/index_pruning_integration_test.go @@ -0,0 +1,80 @@ +package migrations + +import ( + "os" + "path/filepath" + + commonsdb "github.com/flanksource/commons-db/db" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Removing an index from the HCL only prunes a real database if Atlas treats it +// as drift to reconcile rather than as an unmanaged object to leave alone. The +// indexes this schema deliberately no longer declares were measured dead, but +// they sit on the two most write-heavy tables, so leaving them behind on an +// existing database would keep paying their upkeep forever. +var _ = Describe("Captain index reconciliation", func() { + It("prunes indexes the schema no longer declares and narrows the ones it redeclares", func(ctx SpecContext) { + if os.Getenv("CAPTAIN_DB_EMBEDDED_TEST") == "" { + Skip("set CAPTAIN_DB_EMBEDDED_TEST=1 to run embedded-postgres migration tests") + } + + dsn, stop, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: filepath.Join(GinkgoT().TempDir(), "postgres"), + Database: "captain_index_pruning", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(stop) + + db, err := commonsdb.NewDB(dsn) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + Expect(Apply(ctx, dsn)).To(Succeed()) + + // Reproduce what a database migrated by an earlier Captain looks like: + // the dead indexes present, and model_call_id indexed in full rather + // than only where it is set. + for _, statement := range []string{ + `CREATE INDEX captain_sessions_project_idx ON public.captain_sessions (project)`, + `CREATE INDEX captain_model_calls_started_at_idx ON public.captain_model_calls (started_at)`, + `CREATE INDEX captain_model_calls_model_idx ON public.captain_model_calls (model, backend)`, + `CREATE INDEX captain_model_calls_iteration_id_idx ON public.captain_model_calls (iteration_id)`, + `DROP INDEX public.captain_messages_model_call_id_idx`, + `CREATE INDEX captain_messages_model_call_id_idx ON public.captain_messages (model_call_id)`, + } { + _, err := db.ExecContext(ctx, statement) + Expect(err).NotTo(HaveOccurred(), statement) + } + + Expect(Apply(ctx, dsn)).To(Succeed()) + + indexDefinition := func(name string) string { + var definition *string + Expect(db.QueryRowContext(ctx, + `SELECT indexdef FROM pg_indexes WHERE schemaname = 'public' AND indexname = $1`, name). + Scan(&definition)).To(Or(Succeed(), MatchError(ContainSubstring("no rows")))) + if definition == nil { + return "" + } + return *definition + } + + for _, dropped := range []string{ + "captain_sessions_project_idx", + "captain_model_calls_started_at_idx", + "captain_model_calls_model_idx", + "captain_model_calls_iteration_id_idx", + } { + Expect(indexDefinition(dropped)).To(BeEmpty(), dropped+" is no longer declared and must be dropped") + } + + // The FK this backs is ON DELETE SET NULL, so it must survive -- just + // without the NULL rows that made up every entry in it. + Expect(indexDefinition("captain_messages_model_call_id_idx")). + To(ContainSubstring("WHERE (model_call_id IS NOT NULL)")) + // Still the referencing side of two foreign keys. + Expect(indexDefinition("captain_model_calls_prompt_run_id_idx")).NotTo(BeEmpty()) + }) +}) diff --git a/migrations/migrations.go b/migrations/migrations.go new file mode 100644 index 0000000..98d5559 --- /dev/null +++ b/migrations/migrations.go @@ -0,0 +1,144 @@ +// Package migrations embeds and applies Captain's authoritative PostgreSQL +// schema. Consumers that share a database with Captain should call Apply before +// creating cross-schema references to Captain-owned tables. +package migrations + +import ( + "context" + "database/sql" + "embed" + "errors" + "fmt" + "strings" + "sync" + "time" + + commonsdb "github.com/flanksource/commons-db/db" + commonsmigrate "github.com/flanksource/commons-db/migrate" +) + +const Scope = "captain" + +const ( + // captainMigrationLockNamespace and captainMigrationLockKey are stable, + // Captain-specific PostgreSQL advisory-lock identifiers ("CAPT", "MIGR"). + // Gavel uses a different outer lifecycle key, so a host can safely hold its + // own lock while Captain serializes this independently owned migration scope. + captainMigrationLockNamespace int32 = 0x43415054 + captainMigrationLockKey int32 = 0x4d494752 + + migrationUnlockTimeout = 5 * time.Second +) + +// schemaFS contains the complete Captain-owned schema. SQL files are colocated +// with the HCL so commons-db/migrate applies their declared pre/post phases in +// the same migration scope. +// +//go:embed *.hcl *.sql +var schemaFS embed.FS + +type migrationLockHandle interface { + Close() error +} + +type applyDependencies struct { + acquireLock func(context.Context, string) (migrationLockHandle, error) + migrate func(context.Context, string) error +} + +var defaultApplyDependencies = applyDependencies{ + acquireLock: acquireMigrationLock, + migrate: func(ctx context.Context, connection string) error { + return commonsmigrate.Apply(ctx, connection, schemaFS, + commonsmigrate.WithName(Scope), + commonsmigrate.WithExclude("todo_*"), + ) + }, +} + +// Apply reconciles the checked-in HCL schema, then applies the colocated SQL +// migrations. A Captain-specific session advisory lock serializes the complete +// migration bundle across processes. It is safe to call repeatedly and uses a +// stable scope so Captain can share a database with other independently +// migrated applications. +func Apply(ctx context.Context, connection string) error { + return apply(ctx, connection, defaultApplyDependencies) +} + +func apply(ctx context.Context, connection string, deps applyDependencies) (resultErr error) { + if strings.TrimSpace(connection) == "" { + return errors.New("captain migration connection string is empty") + } + + lock, err := deps.acquireLock(ctx, connection) + if err != nil { + return fmt.Errorf("acquire Captain migration lock: %w", err) + } + defer func() { + if err := lock.Close(); err != nil { + resultErr = errors.Join(resultErr, fmt.Errorf("release Captain migration lock: %w", err)) + } + }() + + if err := deps.migrate(ctx, connection); err != nil { + return fmt.Errorf("migrate Captain database: %w", err) + } + return nil +} + +type migrationLock struct { + db *sql.DB + conn *sql.Conn + + once sync.Once + err error +} + +func acquireMigrationLock(ctx context.Context, connection string) (migrationLockHandle, error) { + db, err := commonsdb.NewDB(connection) + if err != nil { + return nil, fmt.Errorf("open advisory-lock database: %w", err) + } + conn, err := db.Conn(ctx) + if err != nil { + _ = db.Close() + return nil, fmt.Errorf("reserve advisory-lock connection: %w", err) + } + if _, err := conn.ExecContext(ctx, `SELECT pg_advisory_lock($1, $2)`, + captainMigrationLockNamespace, captainMigrationLockKey); err != nil { + _ = conn.Close() + _ = db.Close() + return nil, fmt.Errorf("lock Captain migration scope: %w", err) + } + return &migrationLock{db: db, conn: conn}, nil +} + +func (lock *migrationLock) Close() error { + if lock == nil { + return nil + } + lock.once.Do(func() { + var cleanupErrors []error + if lock.conn != nil { + ctx, cancel := context.WithTimeout(context.Background(), migrationUnlockTimeout) + var unlocked bool + if err := lock.conn.QueryRowContext(ctx, `SELECT pg_advisory_unlock($1, $2)`, + captainMigrationLockNamespace, captainMigrationLockKey).Scan(&unlocked); err != nil { + cleanupErrors = append(cleanupErrors, fmt.Errorf("unlock Captain migration scope: %w", err)) + } else if !unlocked { + cleanupErrors = append(cleanupErrors, errors.New("captain migration advisory lock was not held")) + } + cancel() + if err := lock.conn.Close(); err != nil { + cleanupErrors = append(cleanupErrors, fmt.Errorf("close Captain migration lock connection: %w", err)) + } + } + if lock.db != nil { + if err := lock.db.Close(); err != nil { + cleanupErrors = append(cleanupErrors, fmt.Errorf("close Captain migration lock database: %w", err)) + } + } + lock.err = errors.Join(cleanupErrors...) + }) + return lock.err +} diff --git a/migrations/migrations_test.go b/migrations/migrations_test.go new file mode 100644 index 0000000..cee76ef --- /dev/null +++ b/migrations/migrations_test.go @@ -0,0 +1,337 @@ +package migrations + +import ( + "context" + "errors" + "io/fs" + "reflect" + "strings" + "testing" +) + +func TestSchemaBundleContainsGavelIntegrationContract(t *testing.T) { + t.Parallel() + + expectedFiles := []string{ + "00_types.pg.hcl", + "01_session_lifecycle_partial.sql", + "10_sessions.pg.hcl", + "20_prompt_runs_and_plans.pg.hcl", + "21_plans.pg.hcl", + "30_execution.pg.hcl", + "40_artifacts.pg.hcl", + "50_constraints.sql", + "51_state_triggers.sql", + "52_session_activity_triggers.sql", + "60_view_session_overview.sql", + "61_view_session_transcript.sql", + "62_view_session_turns.sql", + "63_view_session_agents.sql", + "64_view_session_files.sql", + "65_view_session_plans.sql", + "66_view_session_approvals.sql", + "67_view_session_costs.sql", + "68_view_session_events.sql", + "69_view_prompt_run_overview.sql", + "70_prompt_run_runtime.sql", + "71_session_storage_params.sql", + "72_ingest_storage_params.sql", + } + for _, name := range expectedFiles { + if _, err := fs.Stat(schemaFS, name); err != nil { + t.Errorf("embedded migration %s: %v", name, err) + } + } + assertContainsAll(t, "00_types.pg.hcl", + `values = ["created", "running", "succeeded", "partial", "failed", "cancelled", "interrupted"]`, + ) + assertContainsAll(t, "01_session_lifecycle_partial.sql", + "-- phase: pre", + "ADD VALUE IF NOT EXISTS 'partial' BEFORE 'failed'", + ) + + assertContainsAll(t, "10_sessions.pg.hcl", + `table "captain_sessions"`, + `column "id"`, + `column "lifecycle_status"`, + `column "activity_state"`, + `column "health_state"`, + `column "state_version"`, + ) + assertContainsAll(t, "20_prompt_runs_and_plans.pg.hcl", + `table "captain_prompt_runs"`, + `column "session_id"`, + `column "execution_session_id"`, + `column "root_session_id"`, + `column "phase"`, + `column "state"`, + `column "runtime"`, + `table "captain_prompt_run_iterations"`, + `column "prompt_run_id"`, + ) + assertContainsAll(t, "21_plans.pg.hcl", + `table "captain_plans"`, + `column "approved_revision_id"`, + `table "captain_plan_revisions"`, + `column "plan_id"`, + ) + assertContainsAll(t, "30_execution.pg.hcl", + `table "captain_turns"`, + `column "session_id"`, + `table "captain_model_calls"`, + `column "turn_id"`, + `column "prompt_run_id"`, + `column "iteration_id"`, + `table "captain_events"`, + `table "captain_turn_requests"`, + `column "state"`, + `column "version"`, + ) + assertContainsAll(t, "50_constraints.sql", + "-- phase: post", + "ALTER CONSTRAINT captain_prompt_runs_input_plan_id_fkey", + "DEFERRABLE INITIALLY DEFERRED", + ) + assertContainsAll(t, "51_state_triggers.sql", + "-- phase: post", + "CREATE OR REPLACE FUNCTION public.captain_set_session_state()", + "CREATE TRIGGER captain_sessions_state_before", + "CREATE TRIGGER captain_sessions_updated_at_before", + "REVOKE ALL ON FUNCTION public.captain_sync_prompt_run_iteration() FROM PUBLIC;", + ) + assertContainsAll(t, "52_session_activity_triggers.sql", + "-- phase: post", + "DROP TABLE IF EXISTS public.captain_outbox;", + "DROP FUNCTION IF EXISTS public.captain_emit_session_change();", + "DROP FUNCTION IF EXISTS public.captain_notify_outbox();", + "CREATE OR REPLACE FUNCTION public.captain_touch_session_activity()", + "CREATE TRIGGER captain_messages_activity_after", + "REVOKE ALL ON FUNCTION public.captain_touch_session_activity() FROM PUBLIC;", + // last_activity_at is agent-work only, and the write is monotonic. An + // allowlist keeps a table that lacks an activity column (whose activity_at + // falls through to updated_at = now) from permanently poisoning it. + "SET last_activity_at = GREATEST(last_activity_at, agent_activity_at)", + "AND (last_activity_at IS NULL OR last_activity_at < agent_activity_at)", + "IF agent_activity_at IS NOT NULL AND TG_OP <> 'DELETE' AND TG_TABLE_NAME IN (", + ) + // captain_sessions is heartbeat-updated, so it needs in-page room for HOT + // updates and a far tighter autovacuum trigger than the defaults. Atlas OSS + // cannot express storage parameters, so post-phase SQL owns them and the + // HCL table definition must stay free of them. + assertContainsAll(t, "71_session_storage_params.sql", + "-- phase: post", + "ALTER TABLE public.captain_sessions SET (", + "fillfactor = 70", + "autovacuum_vacuum_scale_factor = 0.02", + ) + assertContainsNone(t, "10_sessions.pg.hcl", "fillfactor", "autovacuum_") + // captain_messages is append-only (ON CONFLICT DO NOTHING), so it keeps the + // dense default fillfactor and only tightens the insert-driven vacuum that + // keeps its visibility map -- and therefore its index-only scans -- intact. + // captain_turns and captain_model_calls are re-upserted on every re-ingest + // pass, so they need in-page room for HOT updates as well. + assertContainsAll(t, "72_ingest_storage_params.sql", + "-- phase: post", + "ALTER TABLE public.captain_messages SET (", + "ALTER TABLE public.captain_turns SET (", + "ALTER TABLE public.captain_model_calls SET (", + "autovacuum_vacuum_insert_scale_factor = 0.02", + ) + assertContainsNone(t, "72_ingest_storage_params.sql", + // An append-only table gains nothing from reserved in-page space. + "ALTER TABLE public.captain_messages SET (\n fillfactor", + ) + assertContainsNone(t, "30_execution.pg.hcl", "fillfactor", "autovacuum_") + for _, name := range expectedFiles { + assertContainsNone(t, name, + `table "captain_outbox"`, + "INSERT INTO public.captain_outbox", + "CREATE OR REPLACE FUNCTION public.captain_notify_outbox()", + "CREATE TRIGGER captain_outbox_notify_after", + "pg_notify('captain_outbox'", + ) + } + for file, view := range map[string]string{ + "60_view_session_overview.sql": "captain_session_overview", + "61_view_session_transcript.sql": "captain_session_transcript", + "62_view_session_turns.sql": "captain_session_turns", + "63_view_session_agents.sql": "captain_session_agents", + "64_view_session_files.sql": "captain_session_files", + "65_view_session_plans.sql": "captain_session_plans", + "66_view_session_approvals.sql": "captain_session_approvals", + "67_view_session_costs.sql": "captain_session_costs", + "68_view_session_events.sql": "captain_session_events", + "69_view_prompt_run_overview.sql": "captain_prompt_run_overview", + } { + assertContainsAll(t, file, + "-- phase: post", + "CREATE OR REPLACE VIEW public."+view, + "COMMENT ON VIEW public."+view, + ) + } + assertContainsAll(t, "70_prompt_run_runtime.sql", + "-- phase: post", + "UPDATE public.captain_prompt_runs", + "NEW.runtime", + "REVOKE ALL ON FUNCTION public.captain_set_prompt_run_state() FROM PUBLIC;", + ) +} + +// Hash-gated run-once scripts keep steady-state applies free of DDL; views are +// restored via commons-db view-dependency invalidation, so no script may opt +// back into re-running on every apply. +func TestSchemaBundleHasNoAlwaysRunScripts(t *testing.T) { + t.Parallel() + + entries, err := fs.Glob(schemaFS, "*.sql") + if err != nil { + t.Fatalf("glob embedded migrations: %v", err) + } + if len(entries) == 0 { + t.Fatal("no embedded SQL migrations found") + } + for _, name := range entries { + data, err := fs.ReadFile(schemaFS, name) + if err != nil { + t.Fatalf("read embedded migration %s: %v", name, err) + } + if strings.Contains(string(data), "-- runs: always") { + t.Errorf("%s declares '-- runs: always'; scripts must be hash-gated run-once", name) + } + } +} + +func TestApplyRejectsEmptyConnection(t *testing.T) { + t.Parallel() + if err := Apply(t.Context(), " "); err == nil { + t.Fatal("Apply unexpectedly accepted an empty connection string") + } +} + +func TestApplyHoldsMigrationLockAcrossMigration(t *testing.T) { + t.Parallel() + + var events []string + err := apply(t.Context(), "postgres://captain", applyDependencies{ + acquireLock: func(context.Context, string) (migrationLockHandle, error) { + events = append(events, "lock") + return &recordingMigrationLock{events: &events}, nil + }, + migrate: func(context.Context, string) error { + events = append(events, "migrate") + return nil + }, + }) + if err != nil { + t.Fatalf("apply: %v", err) + } + assertEventsEqual(t, events, []string{"lock", "migrate", "unlock"}) +} + +func TestApplyReleasesMigrationLockOnMigrationFailure(t *testing.T) { + t.Parallel() + + var events []string + migrationErr := errors.New("atlas failed") + err := apply(t.Context(), "postgres://captain", applyDependencies{ + acquireLock: func(context.Context, string) (migrationLockHandle, error) { + events = append(events, "lock") + return &recordingMigrationLock{events: &events}, nil + }, + migrate: func(context.Context, string) error { + events = append(events, "migrate") + return migrationErr + }, + }) + if !errors.Is(err, migrationErr) { + t.Fatalf("apply error = %v, want errors.Is(_, %v)", err, migrationErr) + } + assertEventsEqual(t, events, []string{"lock", "migrate", "unlock"}) +} + +func TestApplyReportsLockAcquisitionAndReleaseErrors(t *testing.T) { + t.Parallel() + + t.Run("acquire", func(t *testing.T) { + t.Parallel() + wantErr := errors.New("lock unavailable") + err := apply(t.Context(), "postgres://captain", applyDependencies{ + acquireLock: func(context.Context, string) (migrationLockHandle, error) { + return nil, wantErr + }, + }) + if !errors.Is(err, wantErr) || !strings.Contains(err.Error(), "acquire Captain migration lock") { + t.Fatalf("apply error = %v, want acquisition error", err) + } + }) + + t.Run("release joined with migration error", func(t *testing.T) { + t.Parallel() + migrationErr := errors.New("migration failed") + releaseErr := errors.New("unlock failed") + err := apply(t.Context(), "postgres://captain", applyDependencies{ + acquireLock: func(context.Context, string) (migrationLockHandle, error) { + return &recordingMigrationLock{err: releaseErr}, nil + }, + migrate: func(context.Context, string) error { return migrationErr }, + }) + if !errors.Is(err, migrationErr) || !errors.Is(err, releaseErr) { + t.Fatalf("apply error = %v, want joined migration and release errors", err) + } + }) +} + +func TestCaptainMigrationLockKeyIsStable(t *testing.T) { + t.Parallel() + if captainMigrationLockNamespace != 0x43415054 || captainMigrationLockKey != 0x4d494752 { + t.Fatalf("Captain migration lock key changed: namespace=%#x key=%#x", + captainMigrationLockNamespace, captainMigrationLockKey) + } +} + +type recordingMigrationLock struct { + events *[]string + err error +} + +func (lock *recordingMigrationLock) Close() error { + if lock.events != nil { + *lock.events = append(*lock.events, "unlock") + } + return lock.err +} + +func assertEventsEqual(t *testing.T, got, want []string) { + t.Helper() + if !reflect.DeepEqual(got, want) { + t.Fatalf("events = %v, want %v", got, want) + } +} + +func assertContainsAll(t *testing.T, name string, expected ...string) { + t.Helper() + data, err := fs.ReadFile(schemaFS, name) + if err != nil { + t.Fatalf("read embedded migration %s: %v", name, err) + } + content := string(data) + for _, value := range expected { + if !strings.Contains(content, value) { + t.Errorf("%s does not contain %q", name, value) + } + } +} + +func assertContainsNone(t *testing.T, name string, unexpected ...string) { + t.Helper() + data, err := fs.ReadFile(schemaFS, name) + if err != nil { + t.Fatalf("read embedded migration %s: %v", name, err) + } + content := string(data) + for _, value := range unexpected { + if strings.Contains(content, value) { + t.Errorf("%s unexpectedly contains %q", name, value) + } + } +} diff --git a/migrations/suite_test.go b/migrations/suite_test.go new file mode 100644 index 0000000..0636fd7 --- /dev/null +++ b/migrations/suite_test.go @@ -0,0 +1,13 @@ +package migrations + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestMigrations(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Migrations Suite") +} diff --git a/pkg/ai/adapter_models.go b/pkg/ai/adapter_models.go new file mode 100644 index 0000000..06c0e3b --- /dev/null +++ b/pkg/ai/adapter_models.go @@ -0,0 +1,276 @@ +package ai + +import ( + "context" + "fmt" + "strings" + "sync" + + "github.com/flanksource/captain/pkg/api" +) + +type modelFetch struct { + models []ModelDef + err error +} + +// resolveModelRows is the live model resolver, a package var so tests can +// substitute deterministic rows without hitting a provider API. +var resolveModelRows = ResolveModels + +// fetchAPIModels resolves each direct provider backend's live /v1/models +// endpoint once, concurrently. Local CLI/agent/cmux adapters deliberately do +// not participate: their model catalogs must describe the runtime they execute, +// independent of whether the parent provider's API key happens to be present. +// The resolver is Captain's cached model path, so repeated probes reuse a fresh +// cache instead of hitting providers every time. +func fetchAPIModels(backends []Backend, probe AuthProbe) map[Backend]modelFetch { + apis := map[Backend]bool{} + for _, b := range backends { + if b.Kind() != "api" { + continue + } + source := modelSourceBackend(b) + if source == "" { + continue + } + if effectiveAPIKey(source, probe) != "" { + apis[source] = true + } + } + + out := make(map[Backend]modelFetch, len(apis)) + var mu sync.Mutex + var wg sync.WaitGroup + for b := range apis { + wg.Add(1) + go func(backend Backend) { + defer wg.Done() + rows, err := resolveModelRows(context.Background(), ResolveOptions{Backend: backend, UseTokens: true}) + m := liveModelDefs(rows, backend) + mu.Lock() + out[backend] = modelFetch{models: m, err: err} + mu.Unlock() + }(b) + } + wg.Wait() + return out +} + +func fetchCodexModels(backends []Backend, probe AuthProbe) modelFetch { + if probe.CodexModels == nil { + return modelFetch{} + } + wanted := false + for _, backend := range backends { + if isCodexBackend(backend) { + wanted = true + break + } + } + if !wanted { + return modelFetch{} + } + binary, err := probe.LookPath("codex") + if err != nil || strings.TrimSpace(binary) == "" { + return modelFetch{err: fmt.Errorf("codex not in PATH")} + } + models, err := probe.CodexModels(context.Background(), binary) + return modelFetch{models: models, err: err} +} + +func liveModelDefs(rows []ResolvedModel, backend Backend) []ModelDef { + out := make([]ModelDef, 0, len(rows)) + for _, row := range rows { + if !row.Live { + continue + } + id := row.RuntimeID() + if id == "" { + continue + } + name := row.Label + if name == "" { + name = id + } + out = append(out, ModelDef{ + ID: id, + Name: name, + Backend: backend, + ReleaseDate: row.ReleaseDate, + CapabilitiesKnown: true, + Reasoning: row.Reasoning, + Temperature: row.Temperature, + SupportedEfforts: append([]api.Effort(nil), row.SupportedEfforts...), + DefaultEffort: row.DefaultEffort, + Priority: row.Priority, + }) + } + return out +} + +// applyModels fills in the model listing (or the reason it is unavailable) for +// a single adapter. Direct API backends use live provider rows. Local adapters +// use the catalog of the runtime they execute: Codex's installed catalog when +// available, otherwise Captain's backend-specific registry projection. +func applyModels(st *AdapterStatus, b Backend, cache map[Backend]modelFetch, codex modelFetch, probe AuthProbe) { + if isCodexBackend(b) { + if codex.err == nil && len(codex.models) > 0 { + models := make([]ModelDef, len(codex.models)) + for i, model := range codex.models { + model.Backend = b + models[i] = model + } + setModels(st, models, true) + return + } + setRegistryModels(st, b, codex.err) + return + } + if b.Kind() == "cli" { + setRegistryModels(st, b, nil) + return + } + + source := modelSourceBackend(b) + if source == "" { + st.ModelError = fmt.Sprintf("backend %s has no model listing", b) + return + } + + envVars := AuthEnvVars(source) + if effectiveAPIKey(source, probe) == "" { + st.ModelError = "configure a Captain vault token or set " + strings.Join(envVars, " or ") + " to list models" + return + } + + fetch, ok := cache[source] + if !ok { + return + } + if fetch.err != nil { + st.ModelError = fetch.err.Error() + return + } + setModels(st, modelsForAdapterBackend(b, fetch.models), false) +} + +func effectiveAPIKey(backend Backend, probe AuthProbe) string { + if probe.APICredentials != nil { + return probe.APICredentials[backend].Token + } + return firstEnv(AuthEnvVars(backend), probe.Getenv) +} + +func setRegistryModels(st *AdapterStatus, backend Backend, discoveryErr error) { + setModels(st, RegistryModelDefs(backend), true) + if len(st.Models) > 0 { + return + } + if discoveryErr != nil { + st.ModelError = fmt.Sprintf("runtime model discovery failed: %v; registry has no models for %s", discoveryErr, backend) + return + } + st.ModelError = fmt.Sprintf("registry has no models for %s", backend) +} + +// modelSourceBackend maps any backend onto the API backend whose model list it +// draws from: a CLI/agent/cmux adapter serves its provider family's models. +func modelSourceBackend(backend Backend) Backend { + return backend.Provider() +} + +func modelsForAdapterBackend(backend Backend, models []ModelDef) []ModelDef { + out := make([]ModelDef, 0, len(models)) + positions := map[string]int{} + for _, model := range models { + if model.Backend == BackendOpenAI { + if known, available := RegistryModelAvailability(backend, bareProviderModelID(model.ID)); known && !available { + continue + } + if IsIgnoredOpenAIModelID(model.ID) { + if _, ok := RegistryModelDef(backend, bareProviderModelID(model.ID)); !ok { + continue + } + } + } + id := modelIDForAdapterBackend(backend, model.ID) + if id == "" { + continue + } + name := model.Name + if name == "" { + name = id + } + next := ModelDef{ + ID: id, + Name: name, + Backend: backend, + ReleaseDate: model.ReleaseDate, + CapabilitiesKnown: model.CapabilitiesKnown, + Reasoning: model.Reasoning, + Temperature: model.Temperature, + SupportedEfforts: append([]api.Effort(nil), model.SupportedEfforts...), + DefaultEffort: model.DefaultEffort, + Priority: model.Priority, + } + if idx, ok := positions[id]; ok { + if modelDefNewer(next, out[idx]) { + out[idx] = next + } + continue + } + positions[id] = len(out) + out = append(out, next) + } + return out +} + +func modelIDForAdapterBackend(backend Backend, id string) string { + return NormalizeModelForBackend(backend, bareProviderModelID(id)) +} + +func modelDefNewer(left, right ModelDef) bool { + if left.Priority != right.Priority && (left.Priority > 0 || right.Priority > 0) { + if left.Priority == 0 { + return false + } + if right.Priority == 0 { + return true + } + return left.Priority < right.Priority + } + if left.ReleaseDate == "" { + return false + } + if right.ReleaseDate == "" { + return true + } + return left.ReleaseDate > right.ReleaseDate +} + +func bareProviderModelID(id string) string { + id = strings.TrimSpace(id) + if i := strings.LastIndex(id, "/"); i >= 0 { + return id[i+1:] + } + return id +} + +// setModels filters legacy entries and copies the sorted model list onto the +// adapter status as a count plus id list. The richer details are retained only +// for pretty output; JSON stays as the historical []string model list. +func setModels(st *AdapterStatus, models []ModelDef, curated bool) { + if curated { + models = CurrentCuratedModelsByReleaseDate(models) + } else { + models = CurrentModelsByReleaseDate(models) + } + st.ModelCount = len(models) + ids := make([]string, 0, len(models)) + for _, m := range models { + ids = append(ids, m.ID) + } + st.Models = ids + st.ModelDetails = models +} diff --git a/pkg/ai/adapters.go b/pkg/ai/adapters.go new file mode 100644 index 0000000..46104a5 --- /dev/null +++ b/pkg/ai/adapters.go @@ -0,0 +1,244 @@ +package ai + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/credentials" +) + +// WhoamiOptions selects which adapters to probe and whether to list their +// models. The flag tags drive `captain whoami`; the struct lives here so the +// probe and its caching can be reused by non-CLI consumers (e.g. the aichat +// server's model menu) without importing pkg/cli. +type WhoamiOptions struct { + Backend string `flag:"backend" help:"Show only this backend: anthropic|openai|gemini|deepseek|claude-cli|claude-agent|claude-cmux|codex-cli|codex-agent|codex-cmux|gemini-cli" short:"b"` + Models bool `flag:"models" help:"List models from provider APIs or installed CLI catalogs" default:"true" short:"m"` + Limit int `flag:"limit" help:"Max sample model IDs to show per adapter in pretty output after per-prefix filtering (0 = all)" default:"0" short:"l"` +} + +// AdapterStatus is the resolved auth/availability of a single agent adapter +// (backend). Type is "api" for HTTP providers called with a key, "cli" for +// backends delegated to an installed coding-agent binary. +type AdapterStatus struct { + Backend string `json:"backend"` + Type string `json:"type"` + Authenticated bool `json:"authenticated"` + AuthMethod string `json:"authMethod,omitempty"` + AuthDetail string `json:"authDetail,omitempty"` + Binary string `json:"binary,omitempty"` + BinaryMissing string `json:"binaryMissing,omitempty"` + ModelCount int `json:"modelCount"` + Models []string `json:"models,omitempty"` + ModelError string `json:"modelError,omitempty"` + + ModelDetails []ModelDef `json:"modelDetails,omitempty"` +} + +// Ready reports whether the adapter can actually run: authenticated, and (for +// CLI backends) with its binary present in PATH. +func (a AdapterStatus) Ready() bool { + if !a.Authenticated { + return false + } + if a.Type == "cli" { + return a.Binary != "" + } + return true +} + +// AuthProbe abstracts the host environment (env vars, PATH, credential files) +// so resolveAdapter stays pure and testable. Fields are exported so callers in +// other packages (and their tests) can construct a hermetic probe. +type AuthProbe struct { + Getenv func(string) string + LookPath func(string) (string, error) + FileExists func(string) bool + CodexModels func(context.Context, string) ([]ModelDef, error) + APICredentials map[Backend]api.ResolvedAPIKey + ProbeError error + Home string +} + +// OSAuthProbe wires AuthProbe to the real host environment. +func OSAuthProbe() AuthProbe { + home, _ := os.UserHomeDir() + probe := AuthProbe{ + Getenv: os.Getenv, + LookPath: exec.LookPath, + CodexModels: FetchCodexDebugModels, + FileExists: func(p string) bool { + _, err := os.Stat(p) + return err == nil + }, + Home: home, + } + probe.APICredentials = make(map[Backend]api.ResolvedAPIKey, len(apiBackends)) + for _, backend := range apiBackends { + resolved, err := ResolveAPIKey(backend) + if err != nil { + probe.ProbeError = err + break + } + probe.APICredentials[backend] = resolved + } + return probe +} + +// loginFile is a credential file whose presence indicates a CLI has been logged +// in out-of-band (subscription/OAuth) rather than via an API-key env var. +type loginFile struct { + rel string // path relative to the user's home directory + label string // human label, e.g. "codex login" +} + +// cliAdapter holds the CLI-only metadata for a backend: the binary that must be +// on PATH and the credential files that signal a completed login. +type cliAdapter struct { + binary string + logins []loginFile +} + +func cliAdapters() map[Backend]cliAdapter { + claude := cliAdapter{ + binary: "claude", + logins: []loginFile{ + {rel: filepath.Join(".claude", ".credentials.json"), label: "claude login"}, + {rel: ".claude.json", label: "claude login"}, + }, + } + return map[Backend]cliAdapter{ + BackendClaudeAgent: claude, + BackendClaudeCLI: claude, + BackendClaudeCmux: claude, + BackendCodexCLI: { + binary: "codex", + logins: []loginFile{{rel: filepath.Join(".codex", "auth.json"), label: "codex login"}}, + }, + BackendCodexAgent: { + binary: "codex", + logins: []loginFile{{rel: filepath.Join(".codex", "auth.json"), label: "codex login"}}, + }, + BackendCodexCmux: { + binary: "codex", + logins: []loginFile{{rel: filepath.Join(".codex", "auth.json"), label: "codex login"}}, + }, + BackendGeminiCLI: { + binary: "gemini", + logins: []loginFile{ + {rel: filepath.Join(".gemini", "oauth_creds.json"), label: "gemini login"}, + {rel: filepath.Join(".gemini", "google_accounts.json"), label: "gemini login"}, + }, + }, + } +} + +// resolveAdapter determines a backend's auth method and (for CLI backends) +// binary availability from the probed environment. An API-key env var always +// wins over a CLI login file because that is the path NewProvider/ListModels +// actually take. +func resolveAdapter(backend Backend, p AuthProbe) AdapterStatus { + st := AdapterStatus{Backend: string(backend), Type: backend.Kind()} + + if backend.Kind() == "api" && p.APICredentials != nil { + if resolved := p.APICredentials[backend]; resolved.Token != "" { + st.Authenticated = true + st.AuthDetail = MaskKey(resolved.Token) + if resolved.Source == credentials.SourceVault { + st.AuthMethod = "Captain vault" + } else { + st.AuthMethod = resolved.Detail + " (env)" + } + } + } else { + for _, v := range AuthEnvVars(backend) { + if val := p.Getenv(v); strings.TrimSpace(val) != "" { + st.Authenticated = true + st.AuthMethod = v + " (env)" + st.AuthDetail = MaskKey(val) + break + } + } + } + + if cli, ok := cliAdapters()[backend]; ok { + if path, err := p.LookPath(cli.binary); err == nil { + st.Binary = path + } else { + st.BinaryMissing = cli.binary + } + if !st.Authenticated { + for _, lf := range cli.logins { + full := filepath.Join(p.Home, lf.rel) + if p.FileExists(full) { + st.Authenticated = true + st.AuthMethod = lf.label + st.AuthDetail = full + break + } + } + } + } + + return st +} + +// MaskKey renders a secret as the first and last four characters so the output +// is identifiable without ever exposing the full token. +func MaskKey(s string) string { + s = strings.TrimSpace(s) + if len(s) <= 8 { + return "****" + } + return s[:4] + "…" + s[len(s)-4:] +} + +func firstEnv(vars []string, getenv func(string) string) string { + for _, v := range vars { + if val := getenv(v); strings.TrimSpace(val) != "" { + return val + } + } + return "" +} + +// ProbeAdapters resolves each backend's auth/availability and (when opts.Models) +// its model listing against the supplied environment probe. It is the shared, +// injectable core behind `captain whoami`, the prompt --schema builder, and the +// aichat server model menu, so passing a stub AuthProbe keeps callers hermetic +// (no live API calls when the probe reports no API keys). +func ProbeAdapters(opts WhoamiOptions, probe AuthProbe) ([]AdapterStatus, error) { + if probe.ProbeError != nil { + return nil, probe.ProbeError + } + backends := AllBackends() + if opts.Backend != "" { + b := Backend(opts.Backend) + if !b.Valid() { + return nil, fmt.Errorf("--backend must be one of: %s (got %q)", BackendList(), opts.Backend) + } + backends = []Backend{b} + } + + var models map[Backend]modelFetch + var codexModels modelFetch + if opts.Models { + models = fetchAPIModels(backends, probe) + codexModels = fetchCodexModels(backends, probe) + } + + adapters := make([]AdapterStatus, 0, len(backends)) + for _, b := range backends { + st := resolveAdapter(b, probe) + if opts.Models { + applyModels(&st, b, models, codexModels, probe) + } + adapters = append(adapters, st) + } + return adapters, nil +} diff --git a/pkg/ai/adapters_cache.go b/pkg/ai/adapters_cache.go new file mode 100644 index 0000000..76687f2 --- /dev/null +++ b/pkg/ai/adapters_cache.go @@ -0,0 +1,42 @@ +package ai + +import ( + "sync" + "time" +) + +// adapterCacheTTL bounds how long a probed adapter set is reused. A long-running +// server (whoami-backed model menu, prompt schema) refreshes on this cadence so +// key/login/model changes surface without a probe per request. +const adapterCacheTTL = 60 * time.Second + +// adapterProbe is the live probe sourcing the cache. It is a package var so +// tests can substitute a deterministic, network-free stub. +var adapterProbe = func() ([]AdapterStatus, error) { + return ProbeAdapters(WhoamiOptions{Models: true}, OSAuthProbe()) +} + +var ( + adapterCacheMu sync.Mutex + adapterCache []AdapterStatus + adapterCacheAt time.Time +) + +// CachedAdapters returns the probed adapters, reusing a cached probe within the +// TTL. A probe error is never cached: the next call retries so a transient +// failure does not permanently empty the catalog. `now` is a parameter so tests +// can advance time deterministically. +func CachedAdapters(now time.Time) ([]AdapterStatus, error) { + adapterCacheMu.Lock() + defer adapterCacheMu.Unlock() + if adapterCache != nil && now.Sub(adapterCacheAt) < adapterCacheTTL { + return adapterCache, nil + } + adapters, err := adapterProbe() + if err != nil { + return nil, err + } + adapterCache = adapters + adapterCacheAt = now + return adapters, nil +} diff --git a/pkg/ai/adapters_cache_test.go b/pkg/ai/adapters_cache_test.go new file mode 100644 index 0000000..3acb9b5 --- /dev/null +++ b/pkg/ai/adapters_cache_test.go @@ -0,0 +1,77 @@ +package ai + +import ( + "errors" + "testing" + "time" +) + +func TestCachedAdaptersReusesWithinTTL(t *testing.T) { + prevProbe := adapterProbe + prevCache, prevAt := adapterCache, adapterCacheAt + t.Cleanup(func() { + adapterProbe = prevProbe + adapterCache, adapterCacheAt = prevCache, prevAt + }) + + stub := []AdapterStatus{{Backend: string(BackendAnthropic), Type: "api"}} + calls := 0 + adapterProbe = func() ([]AdapterStatus, error) { + calls++ + return stub, nil + } + adapterCache, adapterCacheAt = nil, time.Time{} + + base := time.Unix(1_000_000, 0) + if _, err := CachedAdapters(base); err != nil { + t.Fatalf("CachedAdapters: %v", err) + } + if _, err := CachedAdapters(base.Add(10 * time.Second)); err != nil { + t.Fatalf("CachedAdapters: %v", err) + } + if calls != 1 { + t.Errorf("probe called %d times within TTL, want 1", calls) + } + if _, err := CachedAdapters(base.Add(2 * adapterCacheTTL)); err != nil { + t.Fatalf("CachedAdapters: %v", err) + } + if calls != 2 { + t.Errorf("probe called %d times after TTL expiry, want 2", calls) + } +} + +func TestCachedAdaptersDoesNotCacheErrors(t *testing.T) { + prevProbe := adapterProbe + prevCache, prevAt := adapterCache, adapterCacheAt + t.Cleanup(func() { + adapterProbe = prevProbe + adapterCache, adapterCacheAt = prevCache, prevAt + }) + adapterCache, adapterCacheAt = nil, time.Time{} + + calls := 0 + adapterProbe = func() ([]AdapterStatus, error) { + calls++ + if calls == 1 { + return nil, errors.New("transient probe failure") + } + return []AdapterStatus{{Backend: string(BackendOpenAI), Type: "api"}}, nil + } + + base := time.Unix(2_000_000, 0) + if _, err := CachedAdapters(base); err == nil { + t.Fatal("expected the transient probe error to surface") + } + // The next call within the TTL window must retry rather than serve a cached + // (empty) failure. + got, err := CachedAdapters(base.Add(time.Second)) + if err != nil { + t.Fatalf("CachedAdapters retry: %v", err) + } + if len(got) != 1 || got[0].Backend != string(BackendOpenAI) { + t.Fatalf("retry adapters = %+v, want the successful probe result", got) + } + if calls != 2 { + t.Errorf("probe called %d times, want 2 (error not cached)", calls) + } +} diff --git a/pkg/ai/adapters_test.go b/pkg/ai/adapters_test.go new file mode 100644 index 0000000..5e8c512 --- /dev/null +++ b/pkg/ai/adapters_test.go @@ -0,0 +1,405 @@ +package ai + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/credentials" +) + +func TestMaskKey(t *testing.T) { + cases := map[string]string{ + "sk-ant-api03-ABCDEFGH": "sk-a…EFGH", + "AIzaSyD-xyz123": "AIza…z123", + "short": "****", + "": "****", + "12345678": "****", + } + for in, want := range cases { + if got := MaskKey(in); got != want { + t.Errorf("MaskKey(%q) = %q, want %q", in, got, want) + } + } +} + +// fakeProbe builds an AuthProbe whose env, PATH, and filesystem are fully +// controlled so resolveAdapter can be exercised without touching the host. +func fakeProbe(env map[string]string, binaries map[string]string, files map[string]bool, home string) AuthProbe { + return AuthProbe{ + Getenv: func(k string) string { return env[k] }, + LookPath: func(b string) (string, error) { + if p, ok := binaries[b]; ok { + return p, nil + } + return "", os.ErrNotExist + }, + FileExists: func(p string) bool { return files[p] }, + Home: home, + } +} + +func TestResolveAdapter_APIKeyFromEnv(t *testing.T) { + st := resolveAdapter(BackendAnthropic, fakeProbe( + map[string]string{"ANTHROPIC_API_KEY": "sk-ant-api03-SECRETKEY"}, + nil, nil, "/home/u")) + + if st.Type != "api" { + t.Errorf("Type = %q, want api", st.Type) + } + if !st.Authenticated || !st.Ready() { + t.Errorf("expected authenticated+ready, got %+v", st) + } + if st.AuthMethod != "ANTHROPIC_API_KEY (env)" { + t.Errorf("AuthMethod = %q", st.AuthMethod) + } + if st.AuthDetail != "sk-a…TKEY" { + t.Errorf("AuthDetail = %q (key should be masked, never printed in full)", st.AuthDetail) + } +} + +func TestResolveAdapter_APINotConfigured(t *testing.T) { + st := resolveAdapter(BackendOpenAI, fakeProbe(nil, nil, nil, "/home/u")) + if st.Authenticated || st.Ready() { + t.Errorf("expected unauthenticated, got %+v", st) + } +} + +func TestOSAuthProbeUsesVaultCredential(t *testing.T) { + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) + for _, name := range api.AuthEnvVars(BackendOpenAI) { + t.Setenv(name, "") + } + vault, err := credentials.DefaultVault() + if err != nil { + t.Fatalf("DefaultVault: %v", err) + } + if err := vault.Set("openai", "vault-openai-secret"); err != nil { + t.Fatalf("Set: %v", err) + } + + adapters, err := ProbeAdapters(WhoamiOptions{Backend: "openai"}, OSAuthProbe()) + if err != nil { + t.Fatalf("ProbeAdapters: %v", err) + } + if len(adapters) != 1 { + t.Fatalf("adapters = %+v", adapters) + } + if got := adapters[0]; !got.Authenticated || got.AuthMethod != "Captain vault" || got.AuthDetail != "vaul…cret" { + t.Fatalf("adapter = %+v, want masked Captain vault credential", got) + } +} + +func TestResolveAdapter_CLILoginFile(t *testing.T) { + home := "/home/u" + authFile := filepath.Join(home, ".codex", "auth.json") + st := resolveAdapter(BackendCodexCLI, fakeProbe( + nil, + map[string]string{"codex": "/usr/local/bin/codex"}, + map[string]bool{authFile: true}, + home)) + + if st.Type != "cli" { + t.Errorf("Type = %q, want cli", st.Type) + } + if st.Binary != "/usr/local/bin/codex" { + t.Errorf("Binary = %q", st.Binary) + } + if !st.Authenticated || !st.Ready() { + t.Errorf("expected authenticated+ready via login file, got %+v", st) + } + if st.AuthMethod != "codex login" || st.AuthDetail != authFile { + t.Errorf("AuthMethod=%q AuthDetail=%q", st.AuthMethod, st.AuthDetail) + } +} + +func TestResolveAdapter_CmuxUsesLocalLoginWithoutAPIKey(t *testing.T) { + home := "/home/u" + claudeAuth := filepath.Join(home, ".claude.json") + codexAuth := filepath.Join(home, ".codex", "auth.json") + + claude := resolveAdapter(BackendClaudeCmux, fakeProbe( + nil, + map[string]string{"claude": "/usr/local/bin/claude"}, + map[string]bool{claudeAuth: true}, + home)) + if claude.Type != "cli" || !claude.Ready() { + t.Fatalf("claude-cmux should be ready through local claude login, got %+v", claude) + } + if claude.AuthMethod != "claude login" || claude.AuthDetail != claudeAuth { + t.Fatalf("claude-cmux auth = %q %q", claude.AuthMethod, claude.AuthDetail) + } + + codex := resolveAdapter(BackendCodexCmux, fakeProbe( + nil, + map[string]string{"codex": "/usr/local/bin/codex"}, + map[string]bool{codexAuth: true}, + home)) + if codex.Type != "cli" || !codex.Ready() { + t.Fatalf("codex-cmux should be ready through local codex login, got %+v", codex) + } + if codex.AuthMethod != "codex login" || codex.AuthDetail != codexAuth { + t.Fatalf("codex-cmux auth = %q %q", codex.AuthMethod, codex.AuthDetail) + } +} + +func TestResolveAdapter_CLIBinaryMissingNotReady(t *testing.T) { + home := "/home/u" + st := resolveAdapter(BackendGeminiCLI, fakeProbe( + map[string]string{"GEMINI_API_KEY": "AIzaSyD-aaaaaaaa"}, // authenticated... + nil, // ...but no gemini binary + nil, home)) + + if !st.Authenticated { + t.Fatalf("expected authenticated via env key, got %+v", st) + } + if st.Binary != "" || st.BinaryMissing != "gemini" { + t.Errorf("expected BinaryMissing=gemini, got Binary=%q BinaryMissing=%q", st.Binary, st.BinaryMissing) + } + if st.Ready() { + t.Error("a CLI adapter with no binary in PATH must not be Ready") + } +} + +func TestResolveAdapter_EnvKeyPreferredOverLogin(t *testing.T) { + home := "/home/u" + st := resolveAdapter(BackendClaudeAgent, fakeProbe( + map[string]string{"ANTHROPIC_API_KEY": "sk-ant-PREFERREDKEY"}, + map[string]string{"claude": "/usr/local/bin/claude"}, + map[string]bool{filepath.Join(home, ".claude.json"): true}, + home)) + + if st.AuthMethod != "ANTHROPIC_API_KEY (env)" { + t.Errorf("env key should win over login file, got AuthMethod=%q", st.AuthMethod) + } +} + +func TestProbeAdaptersUsesRegistryModelsForClaudeCmuxRegardlessOfAPIKey(t *testing.T) { + prev := resolveModelRows + resolveModelRows = func(_ context.Context, opts ResolveOptions) ([]ResolvedModel, error) { + t.Fatalf("local Claude adapter should not resolve provider API models: %+v", opts) + return nil, nil + } + t.Cleanup(func() { resolveModelRows = prev }) + + var withoutKey []string + for _, env := range []map[string]string{nil, {"ANTHROPIC_API_KEY": "sk-ant-test"}} { + adapters, err := ProbeAdapters(WhoamiOptions{Backend: string(BackendClaudeCmux), Models: true}, fakeProbe( + env, + map[string]string{"claude": "/usr/local/bin/claude"}, + nil, + "/home/u", + )) + if err != nil { + t.Fatalf("ProbeAdapters: %v", err) + } + if len(adapters) != 1 || len(adapters[0].Models) == 0 { + t.Fatalf("adapters = %+v, want one adapter with registry models", adapters) + } + if !stringSliceContains(adapters[0].Models, "claude-fable-5") { + t.Fatalf("models = %v, want preferred Fable model", adapters[0].Models) + } + var fable *ModelDef + for i := range adapters[0].ModelDetails { + if adapters[0].ModelDetails[i].ID == "claude-fable-5" { + fable = &adapters[0].ModelDetails[i] + break + } + } + if fable == nil || !fable.CapabilitiesKnown || !fable.Reasoning || fable.Temperature || len(fable.SupportedEfforts) != 5 { + t.Fatalf("fable model details = %+v", fable) + } + if withoutKey == nil { + withoutKey = append([]string(nil), adapters[0].Models...) + continue + } + if strings.Join(adapters[0].Models, "\x00") != strings.Join(withoutKey, "\x00") { + t.Fatalf("models with key = %v, without key = %v", adapters[0].Models, withoutKey) + } + } +} + +func TestProbeAdaptersFiltersNoisyOpenAIModelsForDirectAPI(t *testing.T) { + prev := resolveModelRows + resolveModelRows = func(_ context.Context, opts ResolveOptions) ([]ResolvedModel, error) { + if opts.Backend != BackendOpenAI || !opts.UseTokens { + t.Fatalf("resolve opts = %+v, want openai live token resolve", opts) + } + return []ResolvedModel{ + {Model: Model{ID: "openai/gpt-5.5", Backend: BackendOpenAI, Label: "GPT-5.5", ReleaseDate: "2026-06-01"}, Live: true}, + {Model: Model{ID: "openai/gpt-5.4", Backend: BackendOpenAI, Label: "GPT-5.4", ReleaseDate: "2026-05-15"}, Live: true}, + {Model: Model{ID: "openai/gpt-realtime-2.1", Backend: BackendOpenAI, Label: "Realtime"}, Live: true}, + {Model: Model{ID: "openai/gpt-image-2", Backend: BackendOpenAI, Label: "Image"}, Live: true}, + {Model: Model{ID: "openai/gpt-audio-1.5", Backend: BackendOpenAI, Label: "Audio"}, Live: true}, + {Model: Model{ID: "openai/gpt-5.3-codex", Backend: BackendOpenAI, Label: "Codex"}, Live: true}, + {Model: Model{ID: "openai/gpt-5.3-chat-latest", Backend: BackendOpenAI, Label: "Chat latest"}, Live: true}, + {Model: Model{ID: "openai/gpt-5.5-pro", Backend: BackendOpenAI, Label: "Pro"}, Live: true}, + {Model: Model{ID: "openai/o4-mini", Backend: BackendOpenAI, Label: "O4 mini"}, Live: true}, + {Model: Model{ID: "openai/sora-2", Backend: BackendOpenAI, Label: "Sora"}, Live: true}, + }, nil + } + t.Cleanup(func() { resolveModelRows = prev }) + + adapters, err := ProbeAdapters(WhoamiOptions{Backend: string(BackendOpenAI), Models: true}, fakeProbe( + map[string]string{"OPENAI_API_KEY": "sk-test"}, + map[string]string{"codex": "/usr/local/bin/codex"}, + nil, + "/home/u", + )) + if err != nil { + t.Fatalf("ProbeAdapters: %v", err) + } + if len(adapters) != 1 { + t.Fatalf("adapter count = %d, want 1", len(adapters)) + } + got := adapters[0].Models + for _, want := range []string{"gpt-5.5", "gpt-5.4"} { + if !stringSliceContains(got, want) { + t.Fatalf("models = %v, want primary OpenAI model %q", got, want) + } + } + for _, hidden := range []string{"gpt-realtime-2.1", "gpt-image-2", "gpt-audio-1.5", "gpt-5.3-codex", "gpt-5.3-chat-latest", "gpt-5.5-pro", "o4-mini", "sora-2"} { + if stringSliceContains(got, hidden) { + t.Fatalf("models = %v, should hide noisy OpenAI model %q", got, hidden) + } + } +} + +func TestProbeAdaptersUsesCodexDebugModelsOnceRegardlessOfAPIKey(t *testing.T) { + probe := fakeProbe(map[string]string{"OPENAI_API_KEY": "sk-test"}, map[string]string{"codex": "/usr/local/bin/codex"}, nil, "/home/u") + calls := 0 + probe.CodexModels = func(_ context.Context, binary string) ([]ModelDef, error) { + calls++ + if binary != "/usr/local/bin/codex" { + t.Fatalf("binary = %q", binary) + } + return []ModelDef{ + {ID: "gpt-5.6-sol", Name: "GPT-5.6-Sol", Priority: 1, DefaultEffort: api.EffortLow, SupportedEfforts: []api.Effort{api.EffortLow, api.EffortHigh, api.EffortUltra}}, + {ID: "gpt-5.6-luna", Name: "GPT-5.6-Luna", Priority: 3, DefaultEffort: api.EffortMedium, SupportedEfforts: []api.Effort{api.EffortLow, api.EffortHigh, api.EffortMax}}, + }, nil + } + adapters, err := ProbeAdapters(WhoamiOptions{Models: true}, probe) + if err != nil { + t.Fatalf("ProbeAdapters: %v", err) + } + if calls != 1 { + t.Fatalf("codex debug calls = %d, want 1", calls) + } + for _, backend := range []Backend{BackendCodexCLI, BackendCodexAgent, BackendCodexCmux} { + adapter := findAdapter(t, adapters, backend) + if len(adapter.Models) != 2 || adapter.Models[0] != "gpt-5.6-sol" { + t.Fatalf("%s models = %v", backend, adapter.Models) + } + if adapter.ModelDetails[0].DefaultEffort != api.EffortLow { + t.Fatalf("%s details = %+v", backend, adapter.ModelDetails[0]) + } + } +} + +func TestFetchCodexModelsUsesDebugWithAPIKey(t *testing.T) { + probe := fakeProbe(map[string]string{"OPENAI_API_KEY": "sk-test"}, map[string]string{"codex": "/usr/local/bin/codex"}, nil, "/home/u") + calls := 0 + probe.CodexModels = func(_ context.Context, binary string) ([]ModelDef, error) { + calls++ + if binary != "/usr/local/bin/codex" { + t.Fatalf("binary = %q", binary) + } + return []ModelDef{{ID: "gpt-5.6-sol"}}, nil + } + got := fetchCodexModels([]Backend{BackendCodexAgent}, probe) + if got.err != nil || len(got.models) != 1 || calls != 1 { + t.Fatalf("fetch = %+v calls = %d, want one discovered model", got, calls) + } +} + +func TestProbeAdaptersFallsBackToRegistryWhenCodexDebugFails(t *testing.T) { + probe := fakeProbe(nil, map[string]string{"codex": "/usr/local/bin/codex"}, nil, "/home/u") + probe.CodexModels = func(context.Context, string) ([]ModelDef, error) { + return nil, errors.New("old codex") + } + adapters, err := ProbeAdapters(WhoamiOptions{Backend: string(BackendCodexCLI), Models: true}, probe) + if err != nil { + t.Fatalf("ProbeAdapters: %v", err) + } + if len(adapters) != 1 || !stringSliceContains(adapters[0].Models, "gpt-5.6-sol") { + t.Fatalf("registry fallback models = %+v", adapters) + } +} + +func findAdapter(t *testing.T, adapters []AdapterStatus, backend Backend) AdapterStatus { + t.Helper() + for _, adapter := range adapters { + if adapter.Backend == string(backend) { + return adapter + } + } + t.Fatalf("adapter %s not found", backend) + return AdapterStatus{} +} + +func TestSetModelsFiltersAndSortsByReleaseDate(t *testing.T) { + st := AdapterStatus{Backend: string(BackendAnthropic)} + setModels(&st, []ModelDef{ + {ID: "claude-sonnet-5", Backend: BackendAnthropic, ReleaseDate: "2026-06-01"}, + {ID: "claude-sonnet-4-6", Backend: BackendAnthropic, ReleaseDate: "2026-05-01"}, + {ID: "claude-sonnet-4-5", Backend: BackendAnthropic, ReleaseDate: "2026-04-01"}, + {ID: "claude-sonnet-4-4", Backend: BackendAnthropic, ReleaseDate: "2026-03-01"}, + {ID: "claude-haiku-4-5", Backend: BackendAnthropic, ReleaseDate: "2025-10-15"}, + {ID: "claude-3-5-sonnet-20241022", Backend: BackendAnthropic, ReleaseDate: "2024-10-22"}, + }, false) + + want := []string{"claude-sonnet-5", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5"} + if st.ModelCount != len(want) { + t.Fatalf("ModelCount = %d, want %d (%+v)", st.ModelCount, len(want), st) + } + for i, w := range want { + if st.Models[i] != w { + t.Errorf("Models[%d] = %q, want %q", i, st.Models[i], w) + } + } +} + +func TestSetModelsKeepsGeminiProFamily(t *testing.T) { + st := AdapterStatus{Backend: string(BackendGemini)} + setModels(&st, []ModelDef{ + {ID: "gemini-3.5-flash", Backend: BackendGemini, ReleaseDate: "2026-06-10"}, + {ID: "gemini-3.0-pro", Backend: BackendGemini}, + {ID: "gemini-2.5-pro", Backend: BackendGemini, ReleaseDate: "2025-06-17"}, + {ID: "gemini-2.0-flash", Backend: BackendGemini, ReleaseDate: "2025-02-05"}, + }, false) + + want := []string{"gemini-3.5-flash", "gemini-3.0-pro", "gemini-2.5-pro"} + if st.ModelCount != len(want) { + t.Fatalf("ModelCount = %d, want %d (%+v)", st.ModelCount, len(want), st) + } + for i, w := range want { + if st.Models[i] != w { + t.Errorf("Models[%d] = %q, want %q", i, st.Models[i], w) + } + } +} + +func TestSetModelsHidesCodexCodeVariantsForCLI(t *testing.T) { + st := AdapterStatus{Backend: string(BackendCodexCLI)} + setModels(&st, []ModelDef{ + {ID: "gpt-5-codex", Backend: BackendCodexCLI, ReleaseDate: "2025-08-07"}, + }, false) + + if st.ModelCount != 0 || len(st.Models) != 0 { + t.Fatalf("codex code variant should be hidden for CLI: %+v", st) + } +} + +func stringSliceContains(items []string, want string) bool { + for _, item := range items { + if item == want { + return true + } + } + return false +} diff --git a/pkg/ai/agent.go b/pkg/ai/agent.go index 6b62e00..7179ef7 100644 --- a/pkg/ai/agent.go +++ b/pkg/ai/agent.go @@ -3,12 +3,13 @@ package ai import ( "context" "encoding/json" - "io" + "fmt" "sync" "time" "github.com/flanksource/captain/pkg/ai/pricing" "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" ) // Agent is a convenience wrapper over a Provider that offers the named-prompt / @@ -49,14 +50,15 @@ type PromptRequest struct { // PromptResponse is the result of one PromptRequest. type PromptResponse struct { - Request PromptRequest `json:"request,omitempty"` - Result string `json:"result"` - StructuredData any `json:"structured_data,omitempty"` - Costs Costs `json:"costs,omitempty"` - Model string `json:"model,omitempty"` - Error string `json:"error,omitempty"` - Duration time.Duration `json:"duration,omitempty"` - CacheHit bool `json:"cache_hit,omitempty"` + Request PromptRequest `json:"request,omitempty"` + Result string `json:"result"` + StructuredData any `json:"structured_data,omitempty"` + TerminalOutcome *TerminalOutcome `json:"terminal_outcome,omitempty"` + Costs Costs `json:"costs,omitempty"` + Model string `json:"model,omitempty"` + Error string `json:"error,omitempty"` + Duration time.Duration `json:"duration,omitempty"` + CacheHit bool `json:"cache_hit,omitempty"` } // IsOK reports whether the prompt succeeded. @@ -83,9 +85,17 @@ func (a *Agent) GetModel() string { return a.provider.GetModel() } // GetBackend returns the underlying backend. func (a *Agent) GetBackend() Backend { return a.provider.GetBackend() } -// ExecutePrompt runs one prompt and accrues its cost onto the agent. +// ExecutePrompt runs one prompt and accrues its cost onto the agent. It fails +// with ErrBudgetExceeded before executing once accumulated spend has reached the +// configured USD budget, so a batch of prompts cannot run unbounded. func (a *Agent) ExecutePrompt(ctx context.Context, req PromptRequest) (*PromptResponse, error) { start := time.Now() + if budget := a.cfg.Budget.Cost; budget > 0 { + if spent := a.TotalCost(); spent >= budget { + err := fmt.Errorf("%w: spent $%.4f of $%.4f budget", ErrBudgetExceeded, spent, budget) + return &PromptResponse{Request: req, Model: a.cfg.Model.Name, Error: err.Error()}, err + } + } resp, err := a.provider.Execute(ctx, Request{Prompt: api.Prompt{ User: req.Prompt, System: req.SystemPrompt, @@ -100,13 +110,14 @@ func (a *Agent) ExecutePrompt(ctx context.Context, req PromptRequest) (*PromptRe cost := a.accrue(resp) return &PromptResponse{ - Request: req, - Result: resp.Text, - StructuredData: resp.StructuredData, - Costs: Costs{cost}, - Model: resp.Model, - Duration: time.Since(start), - CacheHit: resp.CacheHit, + Request: req, + Result: resp.Text, + StructuredData: resp.StructuredData, + TerminalOutcome: resp.TerminalOutcome, + Costs: Costs{cost}, + Model: resp.Model, + Duration: time.Since(start), + CacheHit: resp.CacheHit, }, nil } @@ -159,10 +170,18 @@ func (a *Agent) GetCosts() Costs { return out } +// TotalCost returns the accumulated USD spend across every prompt run on this +// agent, preferring provider-reported cost (via Cost.Total). +func (a *Agent) TotalCost() float64 { + a.mu.Lock() + defer a.mu.Unlock() + return a.costs.Sum().Total() +} + // Close releases the underlying provider's resources, if any. func (a *Agent) Close() error { - if c, ok := a.provider.(io.Closer); ok { - return c.Close() + if closer, ok := api.ProviderAs[api.CloseableProvider](a.provider); ok { + return closer.Close() } return nil } @@ -174,6 +193,20 @@ func (a *Agent) accrue(resp *Response) Cost { if model == "" { model = a.cfg.Model.Name } + cost := PriceResponse(a.provider.GetBackend(), model, resp) + + a.mu.Lock() + a.costs = append(a.costs, cost) + a.mu.Unlock() + return cost +} + +// PriceResponse builds the Cost for a response: the disjoint token buckets, the +// provider-reported total (which Cost.Total() prefers), and a best-effort +// list-price breakdown for display. A pricing miss leaves the bucket costs zero +// but preserves the token counts and any provider-reported total, so cost is +// never silently dropped just because a model is absent from the registry. +func PriceResponse(backend Backend, model string, resp *Response) Cost { cost := Cost{ Model: model, InputTokens: resp.Usage.InputTokens, @@ -182,11 +215,12 @@ func (a *Agent) accrue(resp *Response) Cost { CacheReadTokens: resp.Usage.CacheReadTokens, CacheWriteTokens: resp.Usage.CacheWriteTokens, TotalTokens: resp.Usage.TotalTokens(), + ProviderCostUSD: resp.CostUSD, } // The pricing registry is keyed on OpenRouter-style ids (provider/model); // try the backend-prefixed id first, then the bare model. - for _, id := range pricingIDs(a.provider.GetBackend(), model) { - if res, err := pricing.CalculateCost(id, cost.InputTokens, cost.OutputTokens, resp.Usage.ReasoningTokens, resp.Usage.CacheReadTokens, resp.Usage.CacheWriteTokens); err == nil { + for _, id := range PricingIDs(backend, model) { + if res, err := pricing.CalculateCost(id, cost.InputTokens, cost.OutputTokens, cost.ReasoningTokens, cost.CacheReadTokens, cost.CacheWriteTokens); err == nil { cost.InputCost = res.InputCost cost.OutputCost = res.OutputCost cost.ReasoningCost = res.ReasoningCost @@ -195,24 +229,20 @@ func (a *Agent) accrue(resp *Response) Cost { break } } - - a.mu.Lock() - a.costs = append(a.costs, cost) - a.mu.Unlock() return cost } -// pricingIDs returns the candidate pricing keys for a model, most specific first. -func pricingIDs(backend Backend, model string) []string { - prefix := map[Backend]string{ - BackendAnthropic: "anthropic/", - BackendOpenAI: "openai/", - BackendGemini: "google/", - }[backend] - if prefix != "" { - return []string{prefix + model, model} +// PricingIDs returns the candidate OpenRouter-style pricing keys for a model, +// most specific first. Every backend maps to its underlying vendor prefix so +// list-price lookups resolve even for the CLI/agent/cmux backends (whose models +// are still OpenAI/Anthropic/Google under the hood); the bare model is included +// as a fallback. +func PricingIDs(backend Backend, model string) []string { + p, _, ok := registry.ProviderFor(backend) + if !ok { + return []string{model} } - return []string{model} + return p.PricingIDs(model) } func (r *PromptResponse) errorValue() error { diff --git a/pkg/ai/agent/runner.go b/pkg/ai/agent/runner.go index be503d9..4089244 100644 --- a/pkg/ai/agent/runner.go +++ b/pkg/ai/agent/runner.go @@ -242,11 +242,12 @@ func (r *Runner[T]) runLoop(ctx context.Context, hc *HookContext, result *Result maxIter = 1 } req := r.Request - var verifyErr error + var responseErr, verifyErr error loop, loopErr := ai.RunUntil(ctx, ai.LoopOptions{ Provider: r.Provider, MaxIterations: maxIter, + MaxCostUSD: r.Request.Budget.Cost, // enforce the USD budget across iterations OnEvent: func(iter int, ev ai.Event) { r.recordEvent(hc, ev) if r.OnEvent != nil { @@ -255,7 +256,10 @@ func (r *Runner[T]) runLoop(ctx context.Context, hc *HookContext, result *Result }, BuildRequest: func(iter int, prev *ai.LoopIteration) (ai.Request, bool) { if prev != nil { - r.updateResponse(hc, prev) + if err := r.updateResponse(hc, prev); err != nil { + responseErr = err + return ai.Request{}, false + } verdicts, retry, allValid, err := r.verify(hc) if err != nil { verifyErr = err @@ -281,6 +285,12 @@ func (r *Runner[T]) runLoop(ctx context.Context, hc *HookContext, result *Result if loopErr != nil { return loopErr } + if loop != nil && loop.StopReason == "max-cost" { + return fmt.Errorf("%w: spent $%.4f of $%.4f budget", ai.ErrBudgetExceeded, loop.TotalCost, r.Request.Budget.Cost) + } + if responseErr != nil { + return responseErr + } return verifyErr } @@ -321,13 +331,20 @@ func (r *Runner[T]) verify(hc *HookContext) (verdicts []VerifyResult, retry *ai. // updateResponse folds one completed iteration into the accumulating response: // its assembled text, structured data, usage, and session id. -func (r *Runner[T]) updateResponse(hc *HookContext, prev *ai.LoopIteration) { +func (r *Runner[T]) updateResponse(hc *HookContext, prev *ai.LoopIteration) error { ws := hc.Workspace() if prev.SessionID != "" { ws.SessionID = prev.SessionID } var text strings.Builder for _, ev := range prev.Events { + outcome, err := ai.TerminalOutcomeFromEvent(ev) + if err != nil { + return fmt.Errorf("agent: invalid terminal outcome: %w", err) + } + if outcome != nil { + hc.Response.TerminalOutcome = outcome + } switch ev.Kind { case ai.EventText: text.WriteString(ev.Text) @@ -339,6 +356,7 @@ func (r *Runner[T]) updateResponse(hc *HookContext, prev *ai.LoopIteration) { } hc.Response.Text = text.String() hc.Response.Usage = prev.Usage + return nil } // recordEvent updates the workspace from one streamed event: session id and the diff --git a/pkg/ai/agent/runner_test.go b/pkg/ai/agent/runner_test.go index 623864d..9404bf9 100644 --- a/pkg/ai/agent/runner_test.go +++ b/pkg/ai/agent/runner_test.go @@ -58,6 +58,47 @@ func TestRunner_CapturesSessionAndChangedFiles(t *testing.T) { assert.Equal(t, 1, prov.calls, "no verify hooks ⇒ exactly one iteration") } +func TestRunner_CapturesNativePlanOutcome(t *testing.T) { + prov := &fakeProvider{events: func(int) []ai.Event { + return []ai.Event{ + {Kind: ai.EventToolUse, Tool: "ExitPlanMode", Input: map[string]any{ + "plan": "1. Inspect\n2. Implement", + "planFilePath": "/repo/.claude/plans/example.md", + }}, + {Kind: ai.EventResult, Success: true}, + } + }} + + result, err := (&Runner[string]{ + Provider: prov, + Request: ai.Request{Prompt: api.Prompt{User: "plan"}}, + }).Run(context.Background()) + + require.NoError(t, err) + require.NotNil(t, result.Response.TerminalOutcome) + assert.Equal(t, ai.TerminalOutcomePlan, result.Response.TerminalOutcome.Kind) + require.NotNil(t, result.Response.TerminalOutcome.Plan) + assert.Equal(t, "1. Inspect\n2. Implement", result.Response.TerminalOutcome.Plan.Content) + assert.Equal(t, "/repo/.claude/plans/example.md", result.Response.TerminalOutcome.Plan.Path) +} + +func TestRunner_FailsOnMalformedNativePlanOutcome(t *testing.T) { + prov := &fakeProvider{events: func(int) []ai.Event { + return []ai.Event{ + {Kind: ai.EventToolUse, Tool: "ExitPlanMode", Input: map[string]any{"planFilePath": "/repo/plan.md"}}, + {Kind: ai.EventResult, Success: true}, + } + }} + + _, err := (&Runner[string]{ + Provider: prov, + Request: ai.Request{Prompt: api.Prompt{User: "plan"}}, + }).Run(context.Background()) + + require.ErrorContains(t, err, "ExitPlanMode") + require.ErrorContains(t, err, "plan is required") +} + func TestRunner_VerifyDrivesRerunThenStops(t *testing.T) { prov := &fakeProvider{events: func(int) []ai.Event { return []ai.Event{{Kind: ai.EventResult, Success: true}} diff --git a/pkg/ai/agent_close_ginkgo_test.go b/pkg/ai/agent_close_ginkgo_test.go new file mode 100644 index 0000000..6abe2f8 --- /dev/null +++ b/pkg/ai/agent_close_ginkgo_test.go @@ -0,0 +1,65 @@ +package ai_test + +import ( + "context" + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + captainai "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" +) + +type closeTrackingProvider struct { + closeCalls int + closeErr error +} + +func (p *closeTrackingProvider) Execute(context.Context, api.Spec) (*api.Response, error) { + return &api.Response{}, nil +} + +func (p *closeTrackingProvider) GetModel() string { return "test-model" } +func (p *closeTrackingProvider) GetBackend() api.Backend { return api.BackendCodexAgent } +func (p *closeTrackingProvider) Close() error { + p.closeCalls++ + return p.closeErr +} + +type providerWrapper struct { + inner api.Provider +} + +func (w providerWrapper) Execute(ctx context.Context, req api.Spec) (*api.Response, error) { + return w.inner.Execute(ctx, req) +} + +func (w providerWrapper) GetModel() string { return w.inner.GetModel() } +func (w providerWrapper) GetBackend() api.Backend { return w.inner.GetBackend() } +func (w providerWrapper) Unwrap() api.Provider { return w.inner } + +var _ = Describe("Agent provider lifecycle", func() { + It("closes a provider through nested wrappers", func() { + leaf := &closeTrackingProvider{} + agent := captainai.NewAgentWithProvider( + providerWrapper{inner: providerWrapper{inner: leaf}}, + captainai.Config{}, + ) + + Expect(agent.Close()).To(Succeed()) + Expect(leaf.closeCalls).To(Equal(1)) + }) + + It("returns the wrapped provider close error", func() { + closeErr := errors.New("stop supervised provider") + leaf := &closeTrackingProvider{closeErr: closeErr} + agent := captainai.NewAgentWithProvider(providerWrapper{inner: leaf}, captainai.Config{}) + + err := agent.Close() + + Expect(err).To(MatchError(closeErr)) + Expect(errors.Is(err, closeErr)).To(BeTrue()) + Expect(leaf.closeCalls).To(Equal(1)) + }) +}) diff --git a/pkg/ai/agent_test.go b/pkg/ai/agent_test.go index efb252d..8072949 100644 --- a/pkg/ai/agent_test.go +++ b/pkg/ai/agent_test.go @@ -14,9 +14,12 @@ import ( type mockProvider struct { model string text string + cost float64 + calls int err error closed bool lastReq Request + outcome *TerminalOutcome } func (m *mockProvider) GetModel() string { return m.model } @@ -24,17 +27,29 @@ func (m *mockProvider) GetBackend() Backend { return BackendAnthropic } func (m *mockProvider) Close() error { m.closed = true; return nil } func (m *mockProvider) Execute(_ context.Context, req Request) (*Response, error) { m.lastReq = req + m.calls++ if m.err != nil { return nil, m.err } return &Response{ - Text: m.text + ":" + req.Prompt.User, - Model: m.model, - Backend: BackendAnthropic, - Usage: Usage{InputTokens: 10, OutputTokens: 5}, + Text: m.text + ":" + req.Prompt.User, + Model: m.model, + Backend: BackendAnthropic, + Usage: Usage{InputTokens: 10, OutputTokens: 5}, + CostUSD: m.cost, + TerminalOutcome: m.outcome, }, nil } +func TestAgent_ExecutePromptCarriesTerminalOutcome(t *testing.T) { + outcome := &TerminalOutcome{Kind: TerminalOutcomePlan, Plan: &TerminalPlan{Content: "1. Inspect"}} + a := NewAgentWithProvider(&mockProvider{model: "m", outcome: outcome}, Config{Model: api.Model{Name: "m"}}) + + resp, err := a.ExecutePrompt(context.Background(), PromptRequest{Name: "plan", Prompt: "plan this"}) + require.NoError(t, err) + assert.Same(t, outcome, resp.TerminalOutcome) +} + func TestAgent_ExecutePromptAccruesCost(t *testing.T) { a := NewAgentWithProvider(&mockProvider{model: "test-model", text: "out"}, Config{Model: api.Model{Name: "test-model"}}) @@ -51,6 +66,29 @@ func TestAgent_ExecutePromptAccruesCost(t *testing.T) { assert.Equal(t, "test-model", costs[0].Model) } +func TestAgent_ExecutePromptEnforcesBudget(t *testing.T) { + mp := &mockProvider{model: "test-model", text: "out", cost: 0.08} + a := NewAgentWithProvider(mp, Config{ + Model: api.Model{Name: "test-model"}, + Budget: api.Budget{Cost: 0.10}, + }) + + // First two calls fit under the $0.10 budget (spend 0 → 0.08 → 0.16). + _, err := a.ExecutePrompt(context.Background(), PromptRequest{Name: "p1", Prompt: "a"}) + require.NoError(t, err) + _, err = a.ExecutePrompt(context.Background(), PromptRequest{Name: "p2", Prompt: "b"}) + require.NoError(t, err, "second call still under budget at pre-flight (spent 0.08)") + + // Third call trips the pre-flight (spent 0.16 ≥ 0.10) and must not execute. + callsBefore := mp.calls + resp, err := a.ExecutePrompt(context.Background(), PromptRequest{Name: "p3", Prompt: "c"}) + require.Error(t, err) + assert.ErrorIs(t, err, ErrBudgetExceeded) + assert.Equal(t, callsBefore, mp.calls, "provider must not be invoked once budget is exceeded") + assert.Equal(t, err.Error(), resp.Error) + assert.InDelta(t, 0.16, a.TotalCost(), 1e-9, "TotalCost prefers provider-reported cost") +} + func TestAgent_ExecutePromptForwardsSchemaJSON(t *testing.T) { mp := &mockProvider{model: "m", text: "out"} a := NewAgentWithProvider(mp, Config{Model: api.Model{Name: "m"}}) diff --git a/pkg/ai/assistanttags/envelope.go b/pkg/ai/assistanttags/envelope.go new file mode 100644 index 0000000..ea75404 --- /dev/null +++ b/pkg/ai/assistanttags/envelope.go @@ -0,0 +1,42 @@ +package assistanttags + +import ( + "encoding/json" + "strings" +) + +// resultEnvelope is the subset of gavel's result/plan final-result envelope +// needed to recognize a structured completion message. The full schema (plan, +// questions) is owned by flanksource/gavel and intentionally not modeled here — +// captain only reads these messages, and the dependency runs gavel → captain. +type resultEnvelope struct { + Summary string `json:"summary"` + EndStatus string `json:"endStatus"` +} + +var envelopeEndStatuses = map[string]bool{ + "completed": true, + "failed": true, + "ask": true, +} + +// EnvelopeSummary returns the human-readable summary of a gavel result-envelope +// final message, and ok=false when text is not a recognized envelope. The whole +// (trimmed) message must be a single JSON object carrying a non-empty summary and +// an endStatus of completed|failed|ask — strict enough that ordinary assistant +// prose or unrelated JSON is never mistaken for an envelope. +func EnvelopeSummary(text string) (string, bool) { + trimmed := strings.TrimSpace(text) + if !strings.HasPrefix(trimmed, "{") { + return "", false + } + var env resultEnvelope + if err := json.Unmarshal([]byte(trimmed), &env); err != nil { + return "", false + } + summary := strings.TrimSpace(env.Summary) + if summary == "" || !envelopeEndStatuses[env.EndStatus] { + return "", false + } + return summary, true +} diff --git a/pkg/ai/assistanttags/envelope_test.go b/pkg/ai/assistanttags/envelope_test.go new file mode 100644 index 0000000..21544f9 --- /dev/null +++ b/pkg/ai/assistanttags/envelope_test.go @@ -0,0 +1,35 @@ +package assistanttags + +import "testing" + +func TestEnvelopeSummary(t *testing.T) { + planEnvelopeJSON := `{"endStatus":"completed","plan":{"content":"","path":"/Users/moshe/.codex/plans/todo-review-banner-radio-buttons.md","status":"new"},"questions":[],"summary":"Authored and verified a complete read-only implementation plan.\n\n"}` + resultEnvelopeJSON := `{"summary":"Implemented the fix and ran the tests.","endStatus":"failed","questions":[]}` + + tests := []struct { + name string + input string + want string + wantOK bool + }{ + {"full plan envelope, summary trimmed", planEnvelopeJSON, "Authored and verified a complete read-only implementation plan.", true}, + {"result envelope without plan", resultEnvelopeJSON, "Implemented the fix and ran the tests.", true}, + {"ask envelope", `{"summary":"Blocked on a question.","endStatus":"ask","questions":[{"text":"which db?"}]}`, "Blocked on a question.", true}, + {"leading and trailing whitespace", "\n {\"summary\":\"done\",\"endStatus\":\"completed\"} \n", "done", true}, + {"unrelated json object missing endStatus", `{"summary":"looks like a summary but no status"}`, "", false}, + {"invalid endStatus value", `{"summary":"x","endStatus":"in_progress"}`, "", false}, + {"blank summary", `{"summary":" ","endStatus":"completed"}`, "", false}, + {"ordinary prose", "Here is my plan: first do X, then Y.", "", false}, + {"json array not object", `[{"summary":"x","endStatus":"completed"}]`, "", false}, + {"trailing garbage after object", `{"summary":"x","endStatus":"completed"} and more`, "", false}, + {"empty string", "", "", false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := EnvelopeSummary(tt.input) + if ok != tt.wantOK || got != tt.want { + t.Fatalf("EnvelopeSummary(%q) = (%q, %v), want (%q, %v)", tt.input, got, ok, tt.want, tt.wantOK) + } + }) + } +} diff --git a/pkg/ai/assistanttags/parser.go b/pkg/ai/assistanttags/parser.go new file mode 100644 index 0000000..441a0e2 --- /dev/null +++ b/pkg/ai/assistanttags/parser.go @@ -0,0 +1,114 @@ +// Package assistanttags parses structured envelopes emitted in assistant text. +package assistanttags + +import "strings" + +const ( + proposedPlanOpen = "" + proposedPlanClose = "" + memoryCitationOpen = "" + memoryCitationClose = "" + citationEntriesOpen = "" + citationEntriesClose = "" + rolloutIDsOpen = "" + rolloutIDsClose = "" +) + +// SegmentKind identifies one normalized portion of an assistant message. +type SegmentKind string + +const ( + SegmentText SegmentKind = "text" + SegmentPlan SegmentKind = "plan" + SegmentMemoryCitation SegmentKind = "memory_citation" +) + +// MemoryCitation is the structured payload carried by oai-mem-citation. +type MemoryCitation struct { + CitationEntries []string `json:"citation_entries,omitempty"` + RolloutIDs []string `json:"rollout_ids,omitempty"` +} + +// Segment is one ordered portion of a structured assistant message. +type Segment struct { + Kind SegmentKind + Text string + Citation *MemoryCitation +} + +// Parse splits the two assistant protocols observed in Claude/Codex history: +// a leading proposed_plan and a trailing oai-mem-citation. Unknown, embedded, +// or malformed XML remains ordinary text so source/code examples are safe. +func Parse(text string) []Segment { + remaining := strings.TrimSpace(text) + if remaining == "" { + return nil + } + + segments := make([]Segment, 0, 3) + if strings.HasPrefix(remaining, proposedPlanOpen) { + bodyStart := len(proposedPlanOpen) + if closeAt := strings.Index(remaining[bodyStart:], proposedPlanClose); closeAt >= 0 { + closeAt += bodyStart + body := strings.TrimSpace(remaining[bodyStart:closeAt]) + if body != "" { + segments = append(segments, Segment{Kind: SegmentPlan, Text: body}) + } + remaining = strings.TrimSpace(remaining[closeAt+len(proposedPlanClose):]) + } + } + + if remaining == "" { + return segments + } + + if start := strings.LastIndex(remaining, memoryCitationOpen); start >= 0 { + block := strings.TrimSpace(remaining[start:]) + if citation, ok := parseMemoryCitation(block); ok { + if prefix := strings.TrimSpace(remaining[:start]); prefix != "" { + segments = append(segments, Segment{Kind: SegmentText, Text: prefix}) + } + segments = append(segments, Segment{Kind: SegmentMemoryCitation, Citation: citation}) + return segments + } + } + + segments = append(segments, Segment{Kind: SegmentText, Text: remaining}) + return segments +} + +func parseMemoryCitation(block string) (*MemoryCitation, bool) { + if !strings.HasPrefix(block, memoryCitationOpen) || !strings.HasSuffix(block, memoryCitationClose) { + return nil, false + } + body := strings.TrimSpace(strings.TrimSuffix(strings.TrimPrefix(block, memoryCitationOpen), memoryCitationClose)) + entries, ok := sectionLines(body, citationEntriesOpen, citationEntriesClose) + if !ok { + return nil, false + } + rolloutIDs, ok := sectionLines(body, rolloutIDsOpen, rolloutIDsClose) + if !ok { + return nil, false + } + return &MemoryCitation{CitationEntries: entries, RolloutIDs: rolloutIDs}, true +} + +func sectionLines(body, open, close string) ([]string, bool) { + start := strings.Index(body, open) + if start < 0 { + return nil, false + } + start += len(open) + end := strings.Index(body[start:], close) + if end < 0 { + return nil, false + } + end += start + var lines []string + for _, line := range strings.Split(body[start:end], "\n") { + if line = strings.TrimSpace(line); line != "" { + lines = append(lines, line) + } + } + return lines, true +} diff --git a/pkg/ai/assistanttags/parser_test.go b/pkg/ai/assistanttags/parser_test.go new file mode 100644 index 0000000..a946417 --- /dev/null +++ b/pkg/ai/assistanttags/parser_test.go @@ -0,0 +1,75 @@ +package assistanttags + +import ( + "reflect" + "testing" +) + +func TestParseStructuredAssistantMessage(t *testing.T) { + input := ` +# Ship it + +- Add parser + + +Source used: https://example.com/change + + + +MEMORY.md:10-12|note=[parser seam] + + +019f3754-ecfa-7323-a76b-a0205ea30bbe + +` + + got := Parse(input) + if len(got) != 3 { + t.Fatalf("segments = %#v, want 3", got) + } + if got[0].Kind != SegmentPlan || got[0].Text != "# Ship it\n\n- Add parser" { + t.Fatalf("plan = %#v", got[0]) + } + if got[1].Kind != SegmentText || got[1].Text != "Source used: https://example.com/change" { + t.Fatalf("text = %#v", got[1]) + } + if got[2].Kind != SegmentMemoryCitation || got[2].Citation == nil { + t.Fatalf("citation = %#v", got[2]) + } + if !reflect.DeepEqual(got[2].Citation.CitationEntries, []string{"MEMORY.md:10-12|note=[parser seam]"}) { + t.Fatalf("citation entries = %#v", got[2].Citation.CitationEntries) + } + if !reflect.DeepEqual(got[2].Citation.RolloutIDs, []string{"019f3754-ecfa-7323-a76b-a0205ea30bbe"}) { + t.Fatalf("rollout ids = %#v", got[2].Citation.RolloutIDs) + } +} + +func TestParseLeavesUnknownAndMalformedXMLAsText(t *testing.T) { + tests := []string{ + "missing close", + "Before embedded", + "```xml\nexample\n```", + "legacy text", + "x", + } + for _, input := range tests { + got := Parse(input) + if len(got) != 1 || got[0].Kind != SegmentText || got[0].Text != input { + t.Errorf("Parse(%q) = %#v, want unchanged text", input, got) + } + } +} + +func TestParsePlanAndCitationWithoutAssistantText(t *testing.T) { + input := `Do it + + + + + +` + got := Parse(input) + if len(got) != 2 || got[0].Kind != SegmentPlan || got[1].Kind != SegmentMemoryCitation { + t.Fatalf("segments = %#v", got) + } +} diff --git a/pkg/ai/attachment_capabilities.go b/pkg/ai/attachment_capabilities.go new file mode 100644 index 0000000..0391b94 --- /dev/null +++ b/pkg/ai/attachment_capabilities.go @@ -0,0 +1,51 @@ +package ai + +import ( + "fmt" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" +) + +// Which attachment types an adapter can carry is a per-mode provider capability +// (registry.ModeCapabilities.MediaTypes), not a table maintained here. + +// AdapterInputMediaTypes returns the attachment types a backend can carry. +func AdapterInputMediaTypes(backend api.Backend) []string { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return []string{} + } + return p.AdapterMediaTypes(mode) +} + +func modelInputMediaTypes(backend api.Backend, model string) []string { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return []string{} + } + return p.MediaTypesFor(mode, model) +} + +func clampInputMediaTypes(backend api.Backend, modelTypes []string) []string { + return registry.ClampMediaTypes(AdapterInputMediaTypes(backend), modelTypes) +} + +func ValidateAttachmentCompatibility(models []api.Model, refs []api.AttachmentRef) error { + if len(refs) == 0 { + return nil + } + for _, model := range models { + backend, err := model.ResolveBackend() + if err != nil { + return fmt.Errorf("resolve attachment backend for %s: %w", model.Name, err) + } + accepted := modelInputMediaTypes(backend, model.Name) + for _, ref := range refs { + if !registry.MediaTypeAccepted(accepted, ref.MediaType) { + return fmt.Errorf("%s:%s does not accept %s attachments", backend, model.Name, ref.MediaType) + } + } + } + return nil +} diff --git a/pkg/ai/attachment_capabilities_ginkgo_test.go b/pkg/ai/attachment_capabilities_ginkgo_test.go new file mode 100644 index 0000000..feaca11 --- /dev/null +++ b/pkg/ai/attachment_capabilities_ginkgo_test.go @@ -0,0 +1,65 @@ +package ai_test + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" +) + +var _ = Describe("Attachment capabilities", func() { + It("clamps catalog modalities to what the backend adapter can execute", func() { + DeferCleanup(ai.ResetModelCatalog) + model := ai.Model{ + ID: "openai/vision-test", + Backend: ai.BackendOpenAI, + InputMediaTypes: []string{"image/*", "audio/*"}, + } + Expect(ai.SetModelCatalog([]ai.Model{model})).To(Succeed()) + info := ai.CatalogInfo([]string{"openai"}) + Expect(info).To(HaveLen(1)) + Expect(info[0].InputMediaTypes).To(Equal([]string{"image/*"})) + }) + + It("fails the entire multi-model request when one backend is incompatible", func() { + models := []api.Model{ + {Name: "vision-test", Backend: api.BackendCodexAgent}, + {Name: "gemini-3.1-pro", Backend: api.BackendGeminiCLI}, + } + refs := []api.AttachmentRef{{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"}} + + err := ai.ValidateAttachmentCompatibility(models, refs) + Expect(err).To(MatchError(ContainSubstring("gemini-cli:gemini-3.1-pro does not accept image/png attachments"))) + }) + + It("rejects image formats the Claude Agent SDK cannot encode", func() { + err := ai.ValidateAttachmentCompatibility( + []api.Model{{Name: "claude-sonnet-5", Backend: api.BackendClaudeAgent}}, + []api.AttachmentRef{{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/svg+xml"}}, + ) + Expect(err).To(MatchError(ContainSubstring("does not accept image/svg+xml attachments"))) + }) + + It("expands a model image wildcard to the Claude adapter formats", func() { + model := api.Model{Name: "claude-sonnet-5", Backend: api.BackendClaudeAgent} + ref := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"} + Expect(ai.ValidateAttachmentCompatibility([]api.Model{model}, []api.AttachmentRef{ref})).To(Succeed()) + }) + + DescribeTable("publishes the adapter capability matrix", + func(backend api.Backend, expected []string) { + Expect(ai.AdapterInputMediaTypes(backend)).To(Equal(expected)) + }, + Entry("Gemini API", api.BackendGemini, []string{"image/*", "audio/*", "video/*", "application/pdf"}), + Entry("Anthropic API", api.BackendAnthropic, []string{"image/*"}), + Entry("OpenAI API", api.BackendOpenAI, []string{"image/*"}), + Entry("DeepSeek API", api.BackendDeepSeek, []string{}), + Entry("Codex agent", api.BackendCodexAgent, []string{"image/*"}), + Entry("Codex CLI", api.BackendCodexCLI, []string{"image/*"}), + Entry("Claude Agent", api.BackendClaudeAgent, []string{"image/png", "image/jpeg", "image/gif", "image/webp", "application/pdf"}), + Entry("Claude CLI", api.BackendClaudeCLI, []string{}), + Entry("Gemini CLI", api.BackendGeminiCLI, []string{}), + Entry("cmux", api.BackendCodexCmux, []string{}), + ) +}) diff --git a/pkg/ai/attachment_capabilities_suite_test.go b/pkg/ai/attachment_capabilities_suite_test.go new file mode 100644 index 0000000..c874b0f --- /dev/null +++ b/pkg/ai/attachment_capabilities_suite_test.go @@ -0,0 +1,13 @@ +package ai_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAttachmentCapabilities(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "AI Attachment Capabilities Suite") +} diff --git a/pkg/ai/catalog.go b/pkg/ai/catalog.go index 05ccb1f..4a0c42f 100644 --- a/pkg/ai/catalog.go +++ b/pkg/ai/catalog.go @@ -5,24 +5,34 @@ import ( "sort" "strings" "sync" + + "github.com/flanksource/captain/pkg/api" ) // Model describes one entry in the chat model menu. captain owns the catalog // (it is keyed on Backend, which is captain data); clicky/aichat consumes it via -// type aliases. For API backends ID is the full "provider/model" id used by the -// genkit execution path (e.g. "anthropic/claude-sonnet-4-6"). For agent/CLI -// backends ID is the menu key (e.g. "claude-agent-sonnet") and AgentModel is the -// run-time slug when it differs from ID. +// type aliases. API backends carry a provider-prefixed ID for storage/display +// stability, while CLI/agent backends carry the exact provider model ID that the +// local backend receives (never a family alias or synthetic backend prefix). type Model struct { - ID string - Backend Backend - Label string // human-friendly menu label - Reasoning bool // model honours Effort - ContextWindow int // max context tokens, for a usage gauge's denominator - ReleaseDate string // YYYY-MM-DD release date, when known - // AgentModel is the model slug passed to the captain backend when it differs - // from ID (e.g. menu id "codex-gpt-5-codex" → backend model "gpt-5-codex"). - // Empty means use ID. Unused for API backends. + ID string + Backend Backend + Label string // human-friendly menu label + Reasoning bool // model honours Effort + Temperature bool // model honours the temperature sampling control + // AdaptiveThinking marks Anthropic models that use the adaptive thinking + // schema (thinking:{type:adaptive} + output_config.effort) instead of the + // legacy enabled schema. + AdaptiveThinking bool + ContextWindow int // max context tokens, for a usage gauge's denominator + ReleaseDate string // YYYY-MM-DD release date, when known + InputMediaTypes []string + SupportedEfforts []api.Effort + DefaultEffort api.Effort + Priority int + // AgentModel is retained for backward compatibility with older registered + // models. New catalog entries should use exact runtime IDs directly in ID and + // leave AgentModel empty. AgentModel string // Default marks the catalog default. Exactly one entry sets it; its ID equals // DefaultModelID. @@ -51,42 +61,10 @@ func (m Model) BareID() string { // The catalog entry with this ID sets Default: true. const DefaultModelID = "anthropic/claude-sonnet-5" -// defaultCatalog is the model menu: only the latest generally-available model -// per tier for each provider — no preview or superseded entries. API entries -// carry the genkit "provider/model" id (kept byte-identical so the web menu and -// stored-thread model ids stay stable); agent entries carry the captain backend -// explicitly so codex slugs (which look like gpt-*) are not misrouted. -// -// Provider currency (reviewed 2026-07-02): -// - Anthropic: Fable 5 (most capable), Opus 4.8, Sonnet 5, Haiku 4.5. Mythos 5 -// is Project Glasswing invite-only, so it is intentionally excluded. -// - OpenAI: GPT-5.5 (flagship) and GPT-5.4 mini. GPT-5.6 is preview-only. -// - Google: keep recent Pro and Flash families visible separately so a newer -// Flash model does not hide recent Pro entries. -var defaultCatalog = []Model{ - {ID: "anthropic/claude-fable-5", Backend: BackendAnthropic, Label: "Claude Fable 5", Reasoning: true, ContextWindow: 1000000, ReleaseDate: "2026-06-15"}, - {ID: "anthropic/claude-opus-4-8", Backend: BackendAnthropic, Label: "Claude Opus 4.8", Reasoning: true, ContextWindow: 1000000, ReleaseDate: "2026-04-15"}, - {ID: "anthropic/claude-sonnet-5", Backend: BackendAnthropic, Label: "Claude Sonnet 5", Reasoning: true, ContextWindow: 1000000, ReleaseDate: "2026-05-20", Default: true}, - {ID: "anthropic/claude-haiku-4-5", Backend: BackendAnthropic, Label: "Claude Haiku 4.5", Reasoning: true, ContextWindow: 200000, ReleaseDate: "2025-10-15"}, - {ID: "openai/gpt-5.5", Backend: BackendOpenAI, Label: "GPT-5.5", Reasoning: true, ContextWindow: 1000000, ReleaseDate: "2026-06-01"}, - {ID: "openai/gpt-5.4-mini", Backend: BackendOpenAI, Label: "GPT-5.4 mini", Reasoning: true, ContextWindow: 400000, ReleaseDate: "2026-05-15"}, - {ID: "googleai/gemini-2.5-pro", Backend: BackendGemini, Label: "Gemini 2.5 Pro", Reasoning: true, ContextWindow: 1048576, ReleaseDate: "2025-06-17"}, - {ID: "googleai/gemini-3.0-pro", Backend: BackendGemini, Label: "Gemini 3.0 Pro", Reasoning: true, ContextWindow: 1048576}, - {ID: "googleai/gemini-3.5-flash", Backend: BackendGemini, Label: "Gemini 3.5 Flash", Reasoning: true, ContextWindow: 1048576, ReleaseDate: "2026-06-10"}, - // DeepSeek is OpenAI-compatible; deepseek-chat is the non-thinking model and - // deepseek-reasoner is the thinking model (reasoning selected by model, not - // effort). IDs are stable aliases that always resolve to DeepSeek's latest. - {ID: "deepseek/deepseek-chat", Backend: BackendDeepSeek, Label: "DeepSeek Chat", ContextWindow: 131072}, - {ID: "deepseek/deepseek-reasoner", Backend: BackendDeepSeek, Label: "DeepSeek Reasoner", Reasoning: true, ContextWindow: 131072}, - - // Agent-framework models (captain pkg/ai StreamingProvider). These run a - // supervised local subprocess that owns its own tools. IDs are tier aliases - // (not version-pinned), so they always resolve to the installed CLI's latest. - {ID: "claude-agent-sonnet", Backend: BackendClaudeAgent, Label: "Claude Agent · Sonnet", Reasoning: true, ContextWindow: 1000000, ReleaseDate: "2026-05-20"}, - {ID: "claude-agent-opus", Backend: BackendClaudeAgent, Label: "Claude Agent · Opus", Reasoning: true, ContextWindow: 1000000, ReleaseDate: "2026-04-15"}, - {ID: "claude-agent-haiku", Backend: BackendClaudeAgent, Label: "Claude Agent · Haiku", Reasoning: true, ContextWindow: 200000, ReleaseDate: "2025-10-15"}, - {ID: "codex-gpt-5-codex", Backend: BackendCodexCLI, AgentModel: "gpt-5-codex", Label: "Codex · GPT-5", Reasoning: true, ContextWindow: 400000, ReleaseDate: "2025-08-07"}, -} +// defaultCatalog is projected from the internal exact model registry. The API +// rows keep provider prefixes for stable storage; CLI/agent rows keep exact +// backend model IDs without synthetic claude-agent/codex prefixes. +var defaultCatalog = registryCatalogModels() var ( modelRegistryMu sync.RWMutex @@ -191,6 +169,7 @@ func normalizeModel(model Model) (Model, error) { if model.Label == "" { model.Label = model.ID } + model.InputMediaTypes = clampInputMediaTypes(model.Backend, model.InputMediaTypes) if model.ReleaseDate != "" { normalized := normalizeReleaseDate(model.ReleaseDate) if normalized == "" { diff --git a/pkg/ai/catalog_info.go b/pkg/ai/catalog_info.go index 2854f60..7756503 100644 --- a/pkg/ai/catalog_info.go +++ b/pkg/ai/catalog_info.go @@ -3,46 +3,54 @@ package ai import ( "os/exec" "slices" + + "github.com/flanksource/captain/pkg/api/registry" ) // ModelInfo is the JSON shape served at GET /api/chat/models so a client model // selector can be data-driven. Configured reports whether the model is // selectable (its API provider has a key, or its agent backend is installed). type ModelInfo struct { - ID string `json:"id"` - Provider string `json:"provider"` - Label string `json:"label"` - Reasoning bool `json:"reasoning"` - Configured bool `json:"configured"` - ContextWindow int `json:"contextWindow"` + ID string `json:"id"` + Provider string `json:"provider"` + Label string `json:"label"` + Reasoning bool `json:"reasoning"` + Temperature bool `json:"temperature"` + Configured bool `json:"configured"` + ContextWindow int `json:"contextWindow"` + InputMediaTypes []string `json:"inputMediaTypes"` } // BackendToProvider maps a backend to the provider string used in API model ids // and the web menu's "provider" field. API backends use the genkit provider // namespace (anthropic/openai/googleai); agent/CLI backends use their backend // string verbatim. +// BackendToProvider is the provider key the webapp and the session store use. +// +// Its output is a frozen wire format: it is persisted to sessions.provider and +// prompt_runs.runtime.resolved.provider, and /api/chat/models matches on it. The +// API backends map to their catalog namespace ("googleai" for Gemini, not +// "google"); every other backend is its own key verbatim. func BackendToProvider(b Backend) string { - switch b { - case BackendAnthropic: - return "anthropic" - case BackendOpenAI: - return "openai" - case BackendGemini: - return "googleai" - case BackendDeepSeek: - return "deepseek" - default: + p, mode, ok := registry.ProviderFor(b) + if !ok || mode != registry.ModeAPI { return string(b) } + return p.CatalogPrefix } -// CatalogInfo returns the model menu annotated with selectability. API models -// are configured when their provider is among configuredProviders (the keys the -// caller resolved); agent/CLI models are configured when their local backend -// binary is installed. Order mirrors the catalog so the client renders a stable, -// grouped menu. +// CatalogInfo returns the static model menu annotated with selectability. API +// models are configured when their provider is among configuredProviders (the +// keys the caller resolved); agent/CLI models are configured when their local +// backend binary is installed. Order mirrors the catalog so the client renders a +// stable, grouped menu. func CatalogInfo(configuredProviders []string) []ModelInfo { - models := Catalog() + return catalogInfoFrom(Catalog(), configuredProviders) +} + +// catalogInfoFrom annotates an arbitrary model list with selectability, shared +// by CatalogInfo (static catalog) and LiveCatalogInfo (whoami-probed catalog). +func catalogInfoFrom(models []Model, configuredProviders []string) []ModelInfo { out := make([]ModelInfo, len(models)) for i, m := range models { configured := false @@ -52,25 +60,30 @@ func CatalogInfo(configuredProviders []string) []ModelInfo { configured = slices.Contains(configuredProviders, BackendToProvider(m.Backend)) } out[i] = ModelInfo{ - ID: m.ID, - Provider: BackendToProvider(m.Backend), - Label: m.Label, - Reasoning: m.Reasoning, - Configured: configured, - ContextWindow: m.ContextWindow, + ID: m.ID, + Provider: BackendToProvider(m.Backend), + Label: m.Label, + Reasoning: m.Reasoning, + Temperature: m.Temperature, + Configured: configured, + ContextWindow: m.ContextWindow, + InputMediaTypes: append([]string(nil), m.InputMediaTypes...), } } return out } // agentBackendAvailable reports whether an agent backend's local binary is -// installed (best effort): codex needs the `codex` binary; claude-agent / -// claude-cli need `tsx` on PATH. A turn still fails loud if the probe is wrong. +// installed (best effort): codex backends need the `codex` binary; claude-cli +// needs `claude`; claude-agent needs `tsx`. A turn still fails loud if the probe +// is wrong. func agentBackendAvailable(b Backend) bool { switch b { - case BackendCodexCLI, BackendCodexCmux: + case BackendCodexCLI, BackendCodexAgent, BackendCodexCmux: return binaryOnPath("codex") - case BackendClaudeAgent, BackendClaudeCLI, BackendClaudeCmux: + case BackendClaudeCLI, BackendClaudeCmux: + return binaryOnPath("claude") + case BackendClaudeAgent: return binaryOnPath("tsx") default: return false diff --git a/pkg/ai/catalog_resolve.go b/pkg/ai/catalog_resolve.go index 9eb53bf..549fc07 100644 --- a/pkg/ai/catalog_resolve.go +++ b/pkg/ai/catalog_resolve.go @@ -2,11 +2,15 @@ package ai import ( "context" + "crypto/pbkdf2" + "crypto/sha512" + "encoding/hex" "fmt" "sort" "strings" "github.com/flanksource/captain/pkg/ai/pricing" + "github.com/flanksource/captain/pkg/api" "github.com/flanksource/commons/logger" ) @@ -65,7 +69,10 @@ var apiBackends = []Backend{BackendAnthropic, BackendOpenAI, BackendGemini, Back // merged OpenRouter/static pricing and persisted to ~/.config/captain/models.json. // A fresh, fingerprint-matching cache is reused (unless opts.Refresh). func ResolveModels(ctx context.Context, opts ResolveOptions) ([]ResolvedModel, error) { - fp := resolveFingerprint(opts) + fp, err := resolveFingerprint(opts) + if err != nil { + return nil, err + } rows, ok := cachedRows(opts, fp) if !ok { @@ -134,13 +141,16 @@ func seedCatalog(backend Backend) ([]ResolvedModel, map[modelKey]int) { } // catalogBackendMatch reports whether a catalog model's backend satisfies the -// requested filter. claude-cli shares the claude-agent catalog entries. +// requested filter. cli/cmux backends share their agent catalog entries. func catalogBackendMatch(want, modelBackend Backend) bool { if want == "" { return true } - if want == BackendClaudeCLI { + switch want { + case BackendClaudeCLI, BackendClaudeCmux: want = BackendClaudeAgent + case BackendCodexCLI, BackendCodexCmux: + want = BackendCodexAgent } return modelBackend == want } @@ -152,7 +162,11 @@ func unionLive(ctx context.Context, backend Backend, rows *[]ResolvedModel, inde if backend != "" && b != backend { continue } - if GetAPIKeyFromEnv(b) == "" { + resolved, err := ResolveAPIKey(b) + if err != nil { + return err + } + if resolved.Token == "" { continue } live, err := liveModelFetcher(ctx, b) @@ -193,15 +207,35 @@ func filterResolved(rows []ResolvedModel, filter string) []ResolvedModel { return out } -func resolveFingerprint(opts ResolveOptions) string { +// resolveSchemaVersion invalidates the on-disk model cache when the meaning of a +// cached row changes rather than its inputs. The fingerprint otherwise covers +// only options and API keys, so a resolution change would keep serving rows +// resolved by the old rules until the TTL expired. +// +// Bump this when catalog ids, capability fields, or pricing resolution change. +// - v2: model identity unified on the provider descriptors; catalog pricing now +// resolves through the same prefixed-first path as billing, so cached Claude +// prices from the static fallback table are stale. +const resolveSchemaVersion = "v2" +const resolveFingerprintSalt = "captain/model-cache/api-key" + +func resolveFingerprint(opts ResolveOptions) (string, error) { var present []string for _, b := range apiBackends { - if GetAPIKeyFromEnv(b) != "" { - present = append(present, string(b)) + resolved, err := ResolveAPIKey(b) + if err != nil { + return "", err + } + if resolved.Token != "" { + fingerprint, err := pbkdf2.Key(sha512.New, resolved.Token, []byte(resolveFingerprintSalt), 4096, 8) + if err != nil { + return "", fmt.Errorf("fingerprint %s API key: %w", b, err) + } + present = append(present, string(b)+":"+hex.EncodeToString(fingerprint)) } } sort.Strings(present) - return fmt.Sprintf("b=%s|tok=%v|keys=%s", opts.Backend, opts.UseTokens, strings.Join(present, ",")) + return fmt.Sprintf("v=%s|b=%s|tok=%v|keys=%s", resolveSchemaVersion, opts.Backend, opts.UseTokens, strings.Join(present, ",")), nil } func bareModelID(id string) string { @@ -213,13 +247,15 @@ func bareModelID(id string) string { // AgentCatalogModels returns the model list for a CLI/agent backend from the // catalog — the key-free source of truth shared with the chat menu and shell -// completion. The returned ID is the run-time slug the captain provider expects: -// AgentModel when set (e.g. codex's "gpt-5-codex"), otherwise the catalog ID -// (e.g. "claude-agent-sonnet"). claude-cli shares the claude-agent entries. +// completion. Returned IDs are exact provider model IDs; legacy AgentModel is +// still honored for externally registered old entries. func AgentCatalogModels(b Backend) []ModelDef { want := b - if want == BackendClaudeCLI { + switch want { + case BackendClaudeCLI, BackendClaudeCmux: want = BackendClaudeAgent + case BackendCodexCLI, BackendCodexCmux: + want = BackendCodexAgent } out := []ModelDef{} @@ -235,41 +271,38 @@ func AgentCatalogModels(b Backend) []ModelDef { if label == "" { label = id } - out = append(out, ModelDef{ID: id, Name: label, Backend: b, ReleaseDate: m.ReleaseDate}) + out = append(out, ModelDef{ + ID: id, + Name: label, + Backend: b, + ReleaseDate: m.ReleaseDate, + CapabilitiesKnown: true, + Reasoning: m.Reasoning, + Temperature: m.Temperature, + SupportedEfforts: append([]api.Effort(nil), m.SupportedEfforts...), + DefaultEffort: m.DefaultEffort, + Priority: m.Priority, + }) } sort.SliceStable(out, func(i, j int) bool { return out[i].ID < out[j].ID }) return out } // lookupPricing tries the bare model id first, then the OpenRouter -// "provider/model" key the registry uses for the major API providers. +// "provider/model" key the registry uses. +// +// The prefix comes from the provider descriptor's PricingPrefix, which is +// separate from CatalogPrefix for exactly one reason: Gemini's catalog namespace +// is "googleai" while OpenRouter keys it under "google". Deriving one from the +// other makes every Gemini price silently resolve to nothing. This used to be +// three hand-written maps (PricingIDs, orPrefix, pricingModelID) that disagreed +// about whether CLI/agent backends get a prefix at all — they do; a codex-run +// model costs what the model costs. func lookupPricing(backend Backend, id string) (pricing.ModelInfo, bool) { - if info, ok := pricing.GetModelInfo(id); ok { - return info, true - } - prefix := orPrefix(backend) - if prefix == "" { - return pricing.ModelInfo{}, false - } - if info, ok := pricing.GetModelInfo(prefix + "/" + id); ok { - return info, true + for _, candidate := range PricingIDs(backend, id) { + if info, ok := pricing.GetModelInfo(candidate); ok { + return info, true + } } return pricing.ModelInfo{}, false } - -// orPrefix is the OpenRouter id prefix for a backend. Gemini's catalog prefix is -// "googleai" but OpenRouter keys it under "google" — get this right or pricing -// silently misses. -func orPrefix(backend Backend) string { - switch backend { - case BackendOpenAI: - return "openai" - case BackendAnthropic: - return "anthropic" - case BackendGemini: - return "google" - case BackendDeepSeek: - return "deepseek" - } - return "" -} diff --git a/pkg/ai/catalog_resolve_test.go b/pkg/ai/catalog_resolve_test.go index 9076bfd..95297d4 100644 --- a/pkg/ai/catalog_resolve_test.go +++ b/pkg/ai/catalog_resolve_test.go @@ -3,9 +3,43 @@ package ai import ( "context" "errors" + "path/filepath" + "strings" "testing" + + "github.com/flanksource/captain/pkg/credentials" ) +func TestResolveFingerprintChangesWhenVaultTokenRotates(t *testing.T) { + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) + t.Setenv("OPENAI_API_KEY", "") + vault, err := credentials.DefaultVault() + if err != nil { + t.Fatalf("DefaultVault: %v", err) + } + if err := vault.Set("openai", "first-secret-token"); err != nil { + t.Fatalf("Set first token: %v", err) + } + first, err := resolveFingerprint(ResolveOptions{Backend: BackendOpenAI, UseTokens: true}) + if err != nil { + t.Fatalf("first fingerprint: %v", err) + } + if err := vault.Set("openai", "second-secret-token"); err != nil { + t.Fatalf("Set second token: %v", err) + } + second, err := resolveFingerprint(ResolveOptions{Backend: BackendOpenAI, UseTokens: true}) + if err != nil { + t.Fatalf("second fingerprint: %v", err) + } + if first == second { + t.Fatal("token rotation must invalidate the model cache fingerprint") + } + if strings.Contains(first+second, "secret-token") { + t.Fatalf("fingerprints expose token material: %q %q", first, second) + } +} + // stubLiveFetcher installs a fake live-model fetcher for the test and restores // the real one on cleanup. It also isolates API-key env so only the backends // the test opts into are fetched, and points the model cache at a temp HOME. @@ -129,6 +163,35 @@ func TestResolveModels_LegacyHiddenUnlessFiltered(t *testing.T) { } } +func TestResolveModels_PreferredOpenAIVariantsRemainVisible(t *testing.T) { + stubLiveFetcher(t, func(b Backend) ([]ModelDef, error) { + if b != BackendOpenAI { + return nil, nil + } + return []ModelDef{ + {ID: "gpt-5.6-sol", Backend: BackendOpenAI}, + {ID: "gpt-5.6-terra", Backend: BackendOpenAI}, + {ID: "gpt-5.6-luna", Backend: BackendOpenAI}, + {ID: "gpt-5.5-pro", Backend: BackendOpenAI}, + }, nil + }) + t.Setenv("OPENAI_API_KEY", "k") + + rows, err := ResolveModels(context.Background(), ResolveOptions{Backend: BackendOpenAI, UseTokens: true}) + if err != nil { + t.Fatalf("ResolveModels: %v", err) + } + for _, id := range []string{"openai/gpt-5.6-sol", "openai/gpt-5.6-terra", "openai/gpt-5.6-luna"} { + row, ok := hasModelID(rows, id) + if !ok || !row.Live { + t.Errorf("preferred API model %q = %+v, present=%v", id, row, ok) + } + } + if _, ok := hasModelID(rows, "gpt-5.5-pro"); ok { + t.Fatal("non-preferred API variant should remain hidden") + } +} + func TestResolveModels_CacheWrittenAndReused(t *testing.T) { calls := 0 stubLiveFetcher(t, func(b Backend) ([]ModelDef, error) { diff --git a/pkg/ai/catalog_test.go b/pkg/ai/catalog_test.go index 1e1d2e5..cebb727 100644 --- a/pkg/ai/catalog_test.go +++ b/pkg/ai/catalog_test.go @@ -43,8 +43,8 @@ func TestBareID(t *testing.T) { }{ "api strips provider prefix": {Model{ID: "anthropic/claude-sonnet-4-5", Backend: BackendAnthropic}, "claude-sonnet-4-5"}, "gemini googleai prefix": {Model{ID: "googleai/gemini-2.5-pro", Backend: BackendGemini}, "gemini-2.5-pro"}, - "agent id verbatim": {Model{ID: "claude-agent-sonnet", Backend: BackendClaudeAgent}, "claude-agent-sonnet"}, - "codex slug id verbatim": {Model{ID: "codex-gpt-5-codex", Backend: BackendCodexCLI, AgentModel: "gpt-5-codex"}, "codex-gpt-5-codex"}, + "agent id verbatim": {Model{ID: "claude-sonnet-5", Backend: BackendClaudeAgent}, "claude-sonnet-5"}, + "codex slug id verbatim": {Model{ID: "gpt-5.5", Backend: BackendCodexCLI}, "gpt-5.5"}, } for name, tc := range cases { t.Run(name, func(t *testing.T) { @@ -148,6 +148,8 @@ func TestCurrentModelsByReleaseDateFiltersAndSorts(t *testing.T) { {ID: "claude-sonnet-4-6", Backend: BackendAnthropic, ReleaseDate: "2026-05-01"}, {ID: "claude-sonnet-4-5", Backend: BackendAnthropic, ReleaseDate: "2026-04-01"}, {ID: "claude-sonnet-4-4", Backend: BackendAnthropic, ReleaseDate: "2026-03-01"}, + {ID: "sonnet-5", Backend: BackendClaudeCmux, ReleaseDate: "2026-06-01"}, + {ID: "sonnet-4-6", Backend: BackendClaudeCmux, ReleaseDate: "2026-05-01"}, {ID: "claude-haiku-4-5", Backend: BackendAnthropic, ReleaseDate: "2025-10-15"}, {ID: "gemini-3.0-pro", Backend: BackendGemini}, {ID: "gpt-4o", Backend: BackendOpenAI, ReleaseDate: "2026-04-01"}, @@ -159,7 +161,6 @@ func TestCurrentModelsByReleaseDateFiltersAndSorts(t *testing.T) { "claude-sonnet-5", "claude-sonnet-4-6", "claude-sonnet-4-5", - "gpt-5-codex", "claude-haiku-4-5", "gemini-3.0-pro", } @@ -177,6 +178,10 @@ func TestModelFamilyPrefix(t *testing.T) { cases := map[string]string{ "anthropic/claude-haiku-4-5": "claude-haiku", "claude-agent-sonnet": "claude-agent-sonnet", + "claude-sonnet-5": "claude-sonnet", + "sonnet-5": "sonnet", + "sonnet-4-6": "sonnet", + "opus-4-8": "opus", "googleai/gemini-3.0-pro": "gemini-pro", "gemini-3.5-flash": "gemini-flash", "gemini-2.5-flash-lite": "gemini-flash-lite", diff --git a/pkg/ai/client.go b/pkg/ai/client.go index 2c7a6ee..dee1499 100644 --- a/pkg/ai/client.go +++ b/pkg/ai/client.go @@ -26,20 +26,50 @@ func RegisterProvider(backend Backend, factory ProviderFactory) { // name is unrecognized, the error is enriched with the closest known model names // ("did you mean …"). When cfg.Model resolves to more than one candidate (a // comma-separated Name or a Fallbacks list), a fallback provider is returned that -// tries each in order on a retryable failure. +// tries each in order on a fallback-eligible failure. func NewProvider(cfg Config) (Provider, error) { + resolved, err := ResolveModelSelectors(cfg.Model) + if err != nil { + return nil, err + } + cfg.Model = normalizeProviderModel(resolved) candidates := cfg.Model.Candidates() + for i := range candidates { + candidates[i] = normalizeProviderModel(candidates[i]) + } cfg.Model = candidates[0] if len(candidates) > 1 { return newFallbackProvider(cfg, candidates), nil } - p, err := api.NewProvider(cfg) + p, err := newResolvedProvider(cfg) if err != nil { return nil, suggestModelName(err, cfg.Model.Name) } return p, nil } +func newResolvedProvider(cfg Config) (Provider, error) { + p, err := api.NewProvider(cfg) + if err != nil { + return nil, err + } + return withModelErrorRecommendations(withEffortValidation(p, cfg.Model.Effort)), nil +} + +func normalizeProviderModel(model api.Model) api.Model { + backend := model.Backend + if backend == "" { + if inferred, err := api.InferBackend(model.Name); err == nil { + backend = inferred + } + } + if backend != "" { + model.Name = NormalizeModelForBackend(backend, model.Name) + model.Backend = backend + } + return model +} + // suggestModelName appends the closest catalog model ids to an unresolvable-model // error, so e.g. "claud-sonnet-4" points at "claude-sonnet-4". func suggestModelName(err error, model string) error { @@ -73,3 +103,7 @@ func suggestModelName(err error, model string) error { func GetAPIKeyFromEnv(backend Backend) string { return api.GetAPIKeyFromEnv(backend) } + +func ResolveAPIKey(backend Backend) (api.ResolvedAPIKey, error) { + return api.ResolveAPIKey(backend) +} diff --git a/pkg/ai/effort.go b/pkg/ai/effort.go index 3353372..eea1cfd 100644 --- a/pkg/ai/effort.go +++ b/pkg/ai/effort.go @@ -13,6 +13,8 @@ const ( EffortMedium = api.EffortMedium EffortHigh = api.EffortHigh EffortXHigh = api.EffortXHigh + EffortMax = api.EffortMax + EffortUltra = api.EffortUltra ) // AllEfforts lists the non-empty effort tiers in ascending order. diff --git a/pkg/ai/errors.go b/pkg/ai/errors.go index 9df99fd..2c6cfa3 100644 --- a/pkg/ai/errors.go +++ b/pkg/ai/errors.go @@ -3,6 +3,8 @@ package ai import ( "errors" "strings" + + "github.com/flanksource/captain/pkg/api/registry" ) var ( @@ -12,14 +14,27 @@ var ( ErrTimeout = errors.New("operation timed out") ErrSchemaValidation = errors.New("schema validation failed") ErrModelNotFound = errors.New("model not found in pricing registry") + ErrModelUnavailable = errors.New("model unavailable") ErrNoAPIKey = errors.New("API key not found") ) +// ClassifyError normalizes a provider failure into an ErrorClass. Structured +// signals (wrapped sentinels, net.Error, HTTP status) are read before prose. +func ClassifyError(err error) registry.ErrorClass { + return registry.ClassifyError(err) +} + // IsRetryable reports whether err is a transient/overload failure worth retrying // on the same model (see middleware.WithRetry) or falling back to another model -// (see the fallback provider). Provider errors are unstructured strings, so it -// pattern-matches the well-known overload/rate-limit signals; ErrTimeout is -// matched by identity as well. +// (see the fallback provider). +// +// Classification is shared (registry.ClassifyError) so every caller agrees on +// what "transient" means. It reads structure — wrapped sentinels, net.Error, an +// HTTP status — before falling back to bounded text matching. The previous +// implementation matched `strings.Contains(msg, "429")` against raw prose, so +// "1429 tokens" or "id=4290" read as a rate limit and triggered a pointless +// retry. Note a context-length failure is deliberately NOT retryable: the same +// request fails the same way, so retrying only burns tokens. func IsRetryable(err error) bool { if err == nil { return false @@ -27,10 +42,91 @@ func IsRetryable(err error) bool { if errors.Is(err, ErrTimeout) { return true } - msg := err.Error() - return strings.Contains(msg, "rate limit") || - strings.Contains(msg, "429") || - strings.Contains(msg, "503") || - strings.Contains(msg, "overloaded") || - strings.Contains(msg, "timeout") + return registry.ClassifyError(err).Retryable() +} + +// IsModelUnavailable reports provider-confirmed model selection failures. It is +// intentionally separate from IsRetryable: retry middleware must not repeat an +// invalid model on the same provider, while an explicit fallback model may +// still recover the request. +func IsModelUnavailable(err error) bool { + if err == nil { + return false + } + if errors.Is(err, ErrModelNotFound) || errors.Is(err, ErrSchemaValidation) { + return false + } + if errors.Is(err, ErrModelUnavailable) { + return true + } + msg := strings.ToLower(err.Error()) + for _, signal := range []string{ + "model is not supported", + "model is unsupported", + "unsupported model", + "unknown model", + "model not found", + "model was not found", + "model does not exist", + "model is unavailable", + "model_not_found", + "do not have access to model", + "don't have access to model", + } { + if strings.Contains(msg, signal) { + return true + } + } + if strings.Contains(msg, "invalid model") && + !strings.Contains(msg, "model output") && + !strings.Contains(msg, "model response") && + !strings.Contains(msg, "model schema") { + return true + } + return strings.Contains(msg, "model") && + (strings.Contains(msg, "not supported") || + strings.Contains(msg, "not found") || + strings.Contains(msg, "does not exist") || + strings.Contains(msg, "is unavailable")) +} + +// IsMissingAPIKey reports missing credentials without treating rejected or +// invalid credentials as recoverable. Cross-backend fallbacks can therefore +// skip an unconfigured API provider without hiding a bad key. +// +// It stays narrower than ErrorAuth on purpose: a rejected key is an auth failure +// but is not "unconfigured", and skipping past it would hide a real problem. +func IsMissingAPIKey(err error) bool { + if err == nil { + return false + } + if errors.Is(err, ErrNoAPIKey) { + return true + } + msg := strings.ToLower(err.Error()) + for _, signal := range []string{ + "missing api key", + "no api key", + "api key is not set", + "api_key is not set", + } { + if strings.Contains(msg, signal) { + return true + } + } + return false +} + +// IsContextLengthExceeded reports a request that overflowed the model's context +// window. It is intentionally excluded from IsRetryable and IsFallbackEligible's +// transient set: retrying an oversized request reproduces the failure exactly. +func IsContextLengthExceeded(err error) bool { + return registry.ClassifyError(err) == registry.ErrorContextLength +} + +// IsFallbackEligible reports failures for which trying a different explicitly +// configured model may help. This includes transient provider failures, model +// availability mismatches, missing credentials, and a missing local CLI. +func IsFallbackEligible(err error) bool { + return IsRetryable(err) || IsModelUnavailable(err) || IsMissingAPIKey(err) || errors.Is(err, ErrCLINotFound) } diff --git a/pkg/ai/errors_test.go b/pkg/ai/errors_test.go index 28b478d..7c17c43 100644 --- a/pkg/ai/errors_test.go +++ b/pkg/ai/errors_test.go @@ -30,3 +30,33 @@ func TestIsRetryable(t *testing.T) { }) } } + +func TestIsFallbackEligible(t *testing.T) { + cases := []struct { + name string + err error + want bool + }{ + {"retryable", errors.New("429 rate limit"), true}, + {"typed unavailable", fmt.Errorf("codex: %w", ErrModelUnavailable), true}, + {"codex chatgpt unsupported", errors.New("The 'gpt-5.6-line' model is not supported when using Codex with a ChatGPT account."), true}, + {"unknown model", errors.New("provider rejected unknown model gpt-future"), true}, + {"invalid model name", errors.New("invalid model name gpt-future"), true}, + {"typed missing key", fmt.Errorf("openai: %w", ErrNoAPIKey), true}, + {"missing key text", errors.New("genkit provider: no API key for backend openai"), true}, + {"missing cli", fmt.Errorf("start codex: %w", ErrCLINotFound), true}, + {"bad key remains terminal", errors.New("authentication failed: invalid API key"), false}, + {"pricing miss remains terminal", fmt.Errorf("price lookup: %w", ErrModelNotFound), false}, + {"reasoning effort remains terminal", errors.New("model gpt-5.5 does not support reasoning effort max"), false}, + {"malformed request", errors.New("invalid request: missing user"), false}, + {"schema validation", ErrSchemaValidation, false}, + {"invalid model output", errors.New("invalid model output: missing body"), false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := IsFallbackEligible(tc.err); got != tc.want { + t.Errorf("IsFallbackEligible(%v) = %v, want %v", tc.err, got, tc.want) + } + }) + } +} diff --git a/pkg/ai/fallback.go b/pkg/ai/fallback.go index 070d0f2..af29824 100644 --- a/pkg/ai/fallback.go +++ b/pkg/ai/fallback.go @@ -14,7 +14,7 @@ import ( var fallbackLog = logger.GetLogger("ai") // fallbackProvider tries an ordered list of candidate models, advancing to the -// next when the current one fails with a retryable error or its provider cannot be +// next when the current one fails with a fallback-eligible error or its provider cannot be // constructed. Providers are built lazily and cached. GetModel/GetBackend report // the active (last-selected) candidate. It implements StreamingProvider; a // streaming candidate is abandoned in favour of the next only while it has not yet @@ -33,7 +33,7 @@ func newFallbackProvider(base Config, candidates []api.Model) *fallbackProvider return &fallbackProvider{ base: base, candidates: candidates, - build: api.NewProvider, + build: newResolvedProvider, built: make([]Provider, len(candidates)), } } @@ -60,7 +60,7 @@ func (f *fallbackProvider) providerAt(i int) (Provider, error) { } p, err := f.build(f.cfgFor(i)) if err != nil { - return nil, err + return nil, suggestModelName(err, f.candidates[i].Name) } f.built[i] = p return p, nil @@ -89,11 +89,34 @@ func (f *fallbackProvider) GetBackend() Backend { return b } -// Execute runs the primary and, on a retryable failure or a construction error, +func (f *fallbackProvider) Unwrap() Provider { + f.mu.Lock() + defer f.mu.Unlock() + return f.built[f.active] +} + +func (f *fallbackProvider) Close() error { + f.mu.Lock() + built := append([]Provider(nil), f.built...) + f.mu.Unlock() + var errs []error + for _, provider := range built { + if closer, ok := api.ProviderAs[api.CloseableProvider](provider); ok { + errs = append(errs, closer.Close()) + } + } + return errors.Join(errs...) +} + +// Execute runs the primary and, on a fallback-eligible failure or a construction error, // each fallback in turn. A non-retryable runtime error stops immediately (another // model will not fix a malformed request). When nothing succeeds the primary's // error is returned, as the most actionable one. func (f *fallbackProvider) Execute(ctx context.Context, req Request) (*Response, error) { + if err := ValidateAttachmentCompatibility(f.candidates, req.Prompt.Attachments); err != nil { + return nil, err + } + log := LoggerFromContext(ctx, fallbackLog) var firstErr error record := func(err error) { if firstErr == nil { @@ -103,21 +126,23 @@ func (f *fallbackProvider) Execute(ctx context.Context, req Request) (*Response, for i := range f.candidates { p, err := f.providerAt(i) if err != nil { - fallbackLog.Warnf("fallback: model %s unavailable (%v)", f.candidates[i].Name, err) + logFallbackTransition(log, f.candidates, i, err) record(err) continue } f.setActive(i) - resp, err := p.Execute(ctx, req) + candidateReq := req + candidateReq.Effort = f.candidates[i].Effort + resp, err := p.Execute(ctx, candidateReq) if err == nil { return resp, nil } record(err) - if !IsRetryable(err) { + if !IsFallbackEligible(err) { return resp, err } if i < len(f.candidates)-1 { - fallbackLog.Warnf("fallback: model %s failed (%v); trying %s", + log.Warnf("fallback: model %s failed (%v); trying %s", f.candidates[i].Name, err, f.candidates[i+1].Name) } } @@ -130,9 +155,13 @@ func (f *fallbackProvider) Execute(ctx context.Context, req Request) (*Response, // once content has been forwarded the stream is committed and any later error is // surfaced to the caller. func (f *fallbackProvider) ExecuteStream(ctx context.Context, req Request) (<-chan Event, error) { + if err := ValidateAttachmentCompatibility(f.candidates, req.Prompt.Attachments); err != nil { + return nil, err + } out := make(chan Event) go func() { defer close(out) + log := LoggerFromContext(ctx, fallbackLog) var firstErr error record := func(err error) { if firstErr == nil { @@ -146,7 +175,7 @@ func (f *fallbackProvider) ExecuteStream(ctx context.Context, req Request) (<-ch p, err := f.providerAt(i) if err != nil { - fallbackLog.Warnf("fallback: model %s unavailable (%v)", f.candidates[i].Name, err) + logFallbackTransition(log, f.candidates, i, err) record(err) continue } @@ -157,14 +186,20 @@ func (f *fallbackProvider) ExecuteStream(ctx context.Context, req Request) (<-ch } f.setActive(i) - ch, err := sp.ExecuteStream(ctx, req) + candidateReq := req + candidateReq.Effort = f.candidates[i].Effort + ch, err := sp.ExecuteStream(ctx, candidateReq) if err != nil { record(err) - if IsRetryable(err) && !last { - fallbackLog.Warnf("fallback: model %s failed (%v); trying %s", + eligible := IsFallbackEligible(err) + if eligible && !last { + log.Warnf("fallback: model %s failed (%v); trying %s", f.candidates[i].Name, err, f.candidates[i+1].Name) continue } + if eligible && firstErr != nil { + err = firstErr + } emit(err) return } @@ -177,11 +212,15 @@ func (f *fallbackProvider) ExecuteStream(ctx context.Context, req Request) (<-ch streamErr = fmt.Errorf("model %s produced no output", f.candidates[i].Name) } record(streamErr) - if IsRetryable(streamErr) && !last { - fallbackLog.Warnf("fallback: model %s failed pre-output (%v); trying %s", + eligible := IsFallbackEligible(streamErr) + if eligible && !last { + log.Warnf("fallback: model %s failed pre-output (%v); trying %s", f.candidates[i].Name, streamErr, f.candidates[i+1].Name) continue } + if eligible && firstErr != nil { + streamErr = firstErr + } emit(streamErr) return } @@ -192,6 +231,14 @@ func (f *fallbackProvider) ExecuteStream(ctx context.Context, req Request) (<-ch return out, nil } +func logFallbackTransition(log logger.Logger, candidates []api.Model, i int, err error) { + if i < len(candidates)-1 { + log.Warnf("fallback: model %s unavailable (%v); trying %s", candidates[i].Name, err, candidates[i+1].Name) + return + } + log.Warnf("fallback: model %s unavailable (%v)", candidates[i].Name, err) +} + // pump forwards a candidate's stream to out. Leading non-content events (e.g. // EventSystem session metadata) are buffered until the first content event, at // which point the buffer is flushed and the stream is "committed" — from then on diff --git a/pkg/ai/fallback_test.go b/pkg/ai/fallback_test.go index 311d1aa..168700d 100644 --- a/pkg/ai/fallback_test.go +++ b/pkg/ai/fallback_test.go @@ -3,9 +3,11 @@ package ai import ( "context" "errors" + "fmt" "testing" "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/commons/logger" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -93,6 +95,26 @@ func TestFallback_Execute_RetryableAdvances(t *testing.T) { assert.Equal(t, 1, fallback.execCalls) } +func TestFallback_Execute_UnsupportedModelAdvancesAndWarns(t *testing.T) { + primary := &fakeProv{model: "gpt-5.6-line", execErr: errors.New("The 'gpt-5.6-line' model is not supported when using Codex with a ChatGPT account.")} + fallback := &fakeProv{model: "gpt-5.6-sol", execResp: &Response{Text: "ok", Model: "gpt-5.6-sol"}} + fp := newTestFallback([]string{"gpt-5.6-line", "gpt-5.6-sol"}, map[string]*fakeProv{ + "gpt-5.6-line": primary, + "gpt-5.6-sol": fallback, + }) + logs := logger.NewBufferedLogger(10) + ctx := ContextWithLogger(context.Background(), logs) + + resp, err := fp.Execute(ctx, Request{}) + require.NoError(t, err) + assert.Equal(t, "ok", resp.Text) + assert.Equal(t, 1, fallback.execCalls) + warnings := logs.GetLogsByLevel(logger.Warn) + require.Len(t, warnings, 1) + assert.Contains(t, warnings[0].Message, "gpt-5.6-line") + assert.Contains(t, warnings[0].Message, "trying gpt-5.6-sol") +} + func TestFallback_Execute_NonRetryableStops(t *testing.T) { primary := &fakeProv{model: "primary", execErr: errors.New("invalid request: missing user")} fallback := &fakeProv{model: "fallback", execResp: &Response{Text: "ok"}} @@ -115,6 +137,31 @@ func TestFallback_Execute_BuildFailureAdvances(t *testing.T) { assert.Equal(t, 1, fallback.execCalls) } +func TestFallback_Execute_MissingKeyBuildFailureAdvancesAndWarns(t *testing.T) { + fallback := &fakeProv{model: "fallback", execResp: &Response{Text: "ok"}} + models := []api.Model{{Name: "primary", Backend: BackendOpenAI}, {Name: "fallback", Backend: BackendCodexAgent}} + fp := &fallbackProvider{ + candidates: models, + built: make([]Provider, len(models)), + build: func(cfg Config) (Provider, error) { + if cfg.Model.Name == "primary" { + return nil, fmt.Errorf("%w: no API key for openai", ErrNoAPIKey) + } + return fallback, nil + }, + } + logs := logger.NewBufferedLogger(10) + ctx := ContextWithLogger(context.Background(), logs) + + resp, err := fp.Execute(ctx, Request{}) + require.NoError(t, err) + assert.Equal(t, "ok", resp.Text) + warnings := logs.GetLogsByLevel(logger.Warn) + require.Len(t, warnings, 1) + assert.Contains(t, warnings[0].Message, "no API key for openai") + assert.Contains(t, warnings[0].Message, "trying fallback") +} + func TestFallback_Execute_AllFailReturnsPrimaryError(t *testing.T) { primary := &fakeProv{model: "primary", execErr: errors.New("primary 429 overloaded")} fallback := &fakeProv{model: "fallback", execErr: errors.New("fallback 503 overloaded")} @@ -151,6 +198,49 @@ func TestFallback_Stream_PreContentAdvances(t *testing.T) { assert.Equal(t, 1, fallback.streamCalls) } +func TestFallback_Stream_UnsupportedModelPreContentAdvances(t *testing.T) { + primary := &fakeProv{model: "gpt-5.6-line", streamEvents: []Event{ + {Kind: EventSystem, SessionID: "discarded"}, + {Kind: EventError, Error: "unknown model gpt-5.6-line"}, + }} + fallback := &fakeProv{model: "gpt-5.6-sol", streamEvents: []Event{ + {Kind: EventText, Text: "fallback output"}, + {Kind: EventResult, Success: true}, + }} + fp := newTestFallback([]string{"gpt-5.6-line", "gpt-5.6-sol"}, map[string]*fakeProv{ + "gpt-5.6-line": primary, + "gpt-5.6-sol": fallback, + }) + + ch, err := fp.ExecuteStream(context.Background(), Request{}) + require.NoError(t, err) + events := drain(ch) + require.Len(t, events, 2) + assert.Equal(t, "fallback output", events[0].Text) + assert.Equal(t, 1, fallback.streamCalls) +} + +func TestFallback_Stream_AllEligibleFailuresReturnPrimaryError(t *testing.T) { + primary := &fakeProv{model: "gpt-5.6-line", streamEvents: []Event{ + {Kind: EventError, Error: "unknown model gpt-5.6-line; did you mean gpt-5.6-luna?"}, + }} + fallback := &fakeProv{model: "fallback", streamEvents: []Event{ + {Kind: EventError, Error: "503 overloaded"}, + }} + fp := newTestFallback([]string{"gpt-5.6-line", "fallback"}, map[string]*fakeProv{ + "gpt-5.6-line": primary, + "fallback": fallback, + }) + + ch, err := fp.ExecuteStream(context.Background(), Request{}) + require.NoError(t, err) + events := drain(ch) + require.Len(t, events, 1) + assert.Equal(t, EventError, events[0].Kind) + assert.Contains(t, events[0].Error, "unknown model gpt-5.6-line") + assert.NotContains(t, events[0].Error, "503 overloaded") +} + func TestFallback_Stream_CommittedErrorSurfaces(t *testing.T) { primary := &fakeProv{model: "primary", streamEvents: []Event{ {Kind: EventText, Text: "partial"}, diff --git a/pkg/ai/fixture/cost.go b/pkg/ai/fixture/cost.go index 3090639..eacc2ab 100644 --- a/pkg/ai/fixture/cost.go +++ b/pkg/ai/fixture/cost.go @@ -7,13 +7,12 @@ import ( "fmt" "github.com/flanksource/captain/pkg/ai/pricing" - "github.com/flanksource/captain/pkg/claude" ) // resolveCost returns the per-iteration cost for a row. If the underlying // claude CLI didn't report a cost (CostMean == 0) we estimate from the token -// counts: Claude models use captain's deterministic pricing table, everything -// else falls back to the OpenRouter-backed pricing registry. +// counts via the pricing registry, which prices Claude and non-Claude models +// alike from the generated catalog (OpenRouter overlaid on top). // Returns (cost, estimated) — estimated is true when the value is from the // fallback path. func resolveCost(model string, a *aggregate) (float64, bool) { @@ -23,15 +22,6 @@ func resolveCost(model string, a *aggregate) (float64, bool) { if a.Input == 0 && a.Output == 0 && a.CacheRead == 0 && a.CacheWrite == 0 { return 0, false } - if claude.ClassifyModel(model) != claude.ModelFamilyUnknown { - cost := claude.CalculateCost(&claude.Usage{ - InputTokens: a.Input, - OutputTokens: a.Output, - CacheReadInputTokens: a.CacheRead, - CacheCreationInputTokens: a.CacheWrite, - }, model) - return cost, true - } res, err := pricing.CalculateCost(model, a.Input, a.Output, 0, a.CacheRead, a.CacheWrite) if err != nil { return 0, false diff --git a/pkg/ai/generation_config.go b/pkg/ai/generation_config.go new file mode 100644 index 0000000..6b83e47 --- /dev/null +++ b/pkg/ai/generation_config.go @@ -0,0 +1,17 @@ +package ai + +import "github.com/flanksource/captain/pkg/api/registry" + +// EffortConfig builds the provider-native generation config for one request. +// The per-provider translation lives on the provider descriptor +// (registry.Provider.GenerationConfig), so it sits next to the capability data +// that gates it rather than in a backend switch here. +// +// Returns nil when there is nothing to send. +func EffortConfig(backend Backend, model string, effort Effort, maxTokens int, temperature *float64) map[string]any { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return nil + } + return p.GenerationConfig(mode, model, effort, maxTokens, temperature) +} diff --git a/pkg/ai/generation_config_test.go b/pkg/ai/generation_config_test.go new file mode 100644 index 0000000..57d1450 --- /dev/null +++ b/pkg/ai/generation_config_test.go @@ -0,0 +1,118 @@ +package ai + +import ( + "reflect" + "testing" +) + +func f64(v float64) *float64 { return &v } + +func TestEffortConfig(t *testing.T) { + cases := []struct { + name string + backend Backend + model string + effort Effort + maxTokens int + temperature *float64 + want map[string]any + }{ + { + name: "anthropic no effort emits max_tokens only", + backend: BackendAnthropic, model: "claude-sonnet-5", effort: EffortNone, + want: map[string]any{"max_tokens": 4096}, + }, + { + name: "anthropic enabled model uses budget_tokens schema", + backend: BackendAnthropic, model: "claude-sonnet-4-6", effort: EffortHigh, + want: map[string]any{ + "max_tokens": 24576 + 4096, + "thinking": map[string]any{"type": "enabled", "budget_tokens": 24576}, + }, + }, + { + name: "anthropic adaptive model uses output_config schema", + backend: BackendAnthropic, model: "claude-sonnet-5", effort: EffortHigh, + want: map[string]any{ + "max_tokens": 24576 + 4096, + "thinking": map[string]any{"type": "adaptive"}, + "output_config": map[string]any{"effort": "high"}, + }, + }, + { + // Regression: opus-4-8 must use adaptive, not enabled, or the API 400s. + name: "anthropic opus-4-8 is adaptive", + backend: BackendAnthropic, model: "claude-opus-4-8", effort: EffortHigh, + want: map[string]any{ + "max_tokens": 24576 + 4096, + "thinking": map[string]any{"type": "adaptive"}, + "output_config": map[string]any{"effort": "high"}, + }, + }, + { + name: "anthropic xhigh passes source-supported output effort through", + backend: BackendAnthropic, model: "anthropic/claude-fable-5", effort: EffortXHigh, + want: map[string]any{ + "max_tokens": 32768 + 4096, + "thinking": map[string]any{"type": "adaptive"}, + "output_config": map[string]any{"effort": "xhigh"}, + }, + }, + { + name: "anthropic honours explicit max tokens as base", + backend: BackendAnthropic, model: "claude-sonnet-4-6", effort: EffortMedium, maxTokens: 1000, + want: map[string]any{ + "max_tokens": 8192 + 1000, + "thinking": map[string]any{"type": "enabled", "budget_tokens": 8192}, + }, + }, + { + name: "anthropic unknown model omits thinking (no reasoning capability)", + backend: BackendAnthropic, model: "claude-3-5-sonnet-20241022", effort: EffortHigh, + want: map[string]any{"max_tokens": 4096}, + }, + { + name: "openai reasoning model sets reasoning_effort", + backend: BackendOpenAI, model: "gpt-5.5", effort: EffortHigh, + want: map[string]any{"reasoning_effort": "high"}, + }, + { + name: "openai gpt-5.6 passes max effort through", + backend: BackendOpenAI, model: "gpt-5.6", effort: EffortMax, + want: map[string]any{"reasoning_effort": "max"}, + }, + { + name: "openai no effort omits config", + backend: BackendOpenAI, model: "gpt-5.5", effort: EffortNone, + want: nil, + }, + { + name: "gemini low effort sets thinkingBudget", + backend: BackendGemini, model: "gemini-3.5-flash", effort: EffortLow, + want: map[string]any{"thinkingConfig": map[string]any{"thinkingBudget": 2048}}, + }, + { + name: "deepseek omits effort config", + backend: BackendDeepSeek, model: "deepseek-v4-pro", effort: EffortHigh, + want: nil, + }, + { + name: "temperature included for a temperature-capable model", + backend: BackendGemini, model: "gemini-3.5-flash", effort: EffortNone, temperature: f64(0.5), + want: map[string]any{"temperature": 0.5}, + }, + { + name: "temperature dropped for a temperature-incapable model", + backend: BackendAnthropic, model: "claude-sonnet-5", effort: EffortNone, temperature: f64(0.7), + want: map[string]any{"max_tokens": 4096}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := EffortConfig(tc.backend, tc.model, tc.effort, tc.maxTokens, tc.temperature) + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("EffortConfig() = %#v, want %#v", got, tc.want) + } + }) + } +} diff --git a/pkg/ai/history/codex_dedupe_ginkgo_test.go b/pkg/ai/history/codex_dedupe_ginkgo_test.go new file mode 100644 index 0000000..36bd61a --- /dev/null +++ b/pkg/ai/history/codex_dedupe_ginkgo_test.go @@ -0,0 +1,120 @@ +package history + +import ( + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func dedupeUses(lines ...string) []ToolUse { + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(strings.Join(lines, "\n"))) + Expect(err).ToNot(HaveOccurred()) + return uses +} + +func toolNames(uses []ToolUse) []string { + out := make([]string, 0, len(uses)) + for _, use := range uses { + out = append(out, use.Tool) + } + return out +} + +const dedupeSessionMeta = `{"timestamp":"2026-01-21T12:43:41.000Z","type":"session_meta","payload":{"id":"sess-dedupe","cwd":"/repo"}}` + +// Codex logs every message twice -- once as an event_msg, once as a +// response_item -- so a row must merge with its twin but never with a genuine +// repeat of the same text. +var _ = Describe("codex twin dedupe", func() { + It("merges a skewed assistant twin, keeping the response_item content", func() { + uses := dedupeUses( + dedupeSessionMeta, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"event_msg","payload":{"type":"agent_message","message":"done"}}`, + `{"timestamp":"2026-01-21T12:43:41.396Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"done"}]}}`, + ) + + Expect(toolNames(uses)).To(Equal([]string{"Assistant"})) + Expect(uses[0].RecordType).To(Equal("response_item.message")) + }) + + It("merges a response_item-first user twin, keeping the first row", func() { + uses := dedupeUses( + dedupeSessionMeta, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}}`, + `{"timestamp":"2026-01-21T12:43:41.394Z","type":"event_msg","payload":{"type":"user_message","message":"hi"}}`, + ) + + Expect(toolNames(uses)).To(Equal([]string{"User"})) + Expect(uses[0].RecordType).To(Equal("response_item.message")) + }) + + // The family set, not adjacency, is what saves this: with no intervening + // record the two messages are back to back, so a rule keyed on "the previous + // row came from a different family" would fold all four records into one. + It("keeps a repeated message logged as two full twin pairs back to back", func() { + uses := dedupeUses( + dedupeSessionMeta, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]}}`, + `{"timestamp":"2026-01-21T12:43:41.394Z","type":"event_msg","payload":{"type":"user_message","message":"continue"}}`, + `{"timestamp":"2026-01-22T08:02:23.240Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]}}`, + `{"timestamp":"2026-01-22T08:02:23.241Z","type":"event_msg","payload":{"type":"user_message","message":"continue"}}`, + ) + + Expect(toolNames(uses)).To(Equal([]string{"User", "User"})) + Expect(uses[0].Timestamp.UTC().Format(nanoLayout)).To(Equal("2026-01-21T12:43:41.393Z")) + Expect(uses[1].Timestamp.UTC().Format(nanoLayout)).To(Equal("2026-01-22T08:02:23.24Z")) + }) + + // Drawn from rollout-2026-01-21T12-32-48-019be01d: the same prompt re-sent + // 19 hours later. Twin skew is unbounded, so no time window may be used to + // tell these apart -- only the intervening record and the family set. + It("keeps a prompt re-sent after an aborted turn", func() { + uses := dedupeUses( + dedupeSessionMeta, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"Group the transactions"}]}}`, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"event_msg","payload":{"type":"user_message","message":"Group the transactions"}}`, + `{"timestamp":"2026-01-21T12:43:44.797Z","type":"event_msg","payload":{"type":"turn_aborted","turn_id":"turn-1"}}`, + `{"timestamp":"2026-01-22T08:02:23.240Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"Group the transactions"}]}}`, + `{"timestamp":"2026-01-22T08:02:23.241Z","type":"event_msg","payload":{"type":"user_message","message":"Group the transactions"}}`, + ) + + Expect(toolNames(uses)).To(Equal([]string{"User", "TurnAborted", "User"})) + }) + + It("keeps an identical assistant verdict repeated across turns", func() { + uses := dedupeUses( + dedupeSessionMeta, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"event_msg","payload":{"type":"agent_message","message":"{\"outcome\":\"allow\"}"}}`, + `{"timestamp":"2026-01-21T12:43:41.394Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"{\"outcome\":\"allow\"}"}]}}`, + `{"timestamp":"2026-01-21T12:44:05.078Z","type":"event_msg","payload":{"type":"agent_message","message":"{\"outcome\":\"allow\"}"}}`, + `{"timestamp":"2026-01-21T12:44:05.079Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"{\"outcome\":\"allow\"}"}]}}`, + `{"timestamp":"2026-01-21T12:44:18.601Z","type":"event_msg","payload":{"type":"agent_message","message":"{\"outcome\":\"allow\"}"}}`, + `{"timestamp":"2026-01-21T12:44:18.602Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"{\"outcome\":\"allow\"}"}]}}`, + ) + + Expect(toolNames(uses)).To(Equal([]string{"Assistant", "Assistant", "Assistant"})) + }) + + It("merges a skewed plaintext reasoning twin", func() { + uses := dedupeUses( + dedupeSessionMeta, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"event_msg","payload":{"type":"agent_reasoning","text":"Checking the parser"}}`, + `{"timestamp":"2026-01-21T12:43:41.396Z","type":"response_item","payload":{"type":"reasoning","summary":[{"type":"summary_text","text":"Checking the parser"}]}}`, + ) + + Expect(toolNames(uses)).To(Equal([]string{"Reasoning"})) + Expect(uses[0].RecordType).To(Equal("response_item.reasoning")) + }) + + It("does not merge identical messages separated by other chat content", func() { + uses := dedupeUses( + dedupeSessionMeta, + `{"timestamp":"2026-01-21T12:43:41.393Z","type":"event_msg","payload":{"type":"user_message","message":""}}`, + `{"timestamp":"2026-01-21T12:43:42.000Z","type":"event_msg","payload":{"type":"agent_message","message":"working on it"}}`, + `{"timestamp":"2026-01-21T12:43:43.000Z","type":"event_msg","payload":{"type":"user_message","message":""}}`, + ) + + Expect(toolNames(uses)).To(Equal([]string{"User", "Assistant", "User"})) + }) +}) diff --git a/pkg/ai/history/codex_events.go b/pkg/ai/history/codex_events.go new file mode 100644 index 0000000..b5f4d88 --- /dev/null +++ b/pkg/ai/history/codex_events.go @@ -0,0 +1,268 @@ +package history + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/claude/tools" +) + +func codexEventTurnID(event CodexEvent) string { + if event.Payload.TurnID != "" { + return event.Payload.TurnID + } + if event.Payload.Metadata != nil { + return event.Payload.Metadata.TurnID + } + return "" +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + return "" +} + +func buildCodexEventUse(event CodexEvent, cwd, sessionID string) ToolUse { + input := make(map[string]any, len(event.Payload.Raw)+1) + for key, value := range event.Payload.Raw { + input[key] = value + } + input["event"] = event.Payload.Type + addCodexEventValue(input, "turn_id", event.Payload.TurnID) + addCodexEventValue(input, "message", event.Payload.Message) + addCodexEventValue(input, "phase", event.Payload.Phase) + addCodexEventValue(input, "started_at", event.Payload.StartedAt) + addCodexEventValue(input, "completed_at", event.Payload.CompletedAt) + addCodexEventValue(input, "duration_ms", event.Payload.DurationMS) + addCodexEventValue(input, "time_to_first_token_ms", event.Payload.TimeToFirstTokenMS) + addCodexEventValue(input, "last_agent_message", event.Payload.LastAgentMessage) + addCodexEventValue(input, "model_context_window", event.Payload.ModelContextWindow) + addCodexEventValue(input, "collaboration_mode_kind", event.Payload.CollaborationModeKind) + if event.Payload.Info != nil { + last := event.Payload.Info.LastTokenUsage + total := event.Payload.Info.TotalTokenUsage + input["total_tokens"] = last.TotalTokens + input["input_tokens"] = last.InputTokens + input["output_tokens"] = last.OutputTokens + input["cached_input_tokens"] = last.CachedInputTokens + input["cumulative_total_tokens"] = total.TotalTokens + input["cumulative_input_tokens"] = total.InputTokens + input["cumulative_output_tokens"] = total.OutputTokens + input["cumulative_cached_input_tokens"] = total.CachedInputTokens + if last.ReasoningOutputTokens != 0 { + input["reasoning_output_tokens"] = last.ReasoningOutputTokens + } + if event.Payload.Info.ModelContextWindow != 0 { + input["model_context_window"] = event.Payload.Info.ModelContextWindow + } + } + usage := eventTokenUsage(event) + return ToolUse{ + Tool: tools.EventToolName(event.Payload.Type), + Input: input, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + InputTokens: codexNonCachedInputTokens(usage), + OutputTokens: usage.OutputTokens, + CacheReadTokens: usage.CachedInputTokens, + TotalTokens: usage.TotalTokens, + ContextWindow: eventContextWindow(event), + RecordType: "event_msg." + event.Payload.Type, + } +} + +func addCodexEventValue(input map[string]any, key string, value any) { + switch typed := value.(type) { + case string: + if typed != "" { + input[key] = typed + } + case int: + if typed != 0 { + input[key] = typed + } + case int64: + if typed != 0 { + input[key] = typed + } + case float64: + if typed != 0 { + input[key] = typed + } + } +} + +func buildCodexTopLevelEventUse(event CodexEvent, eventType, cwd, sessionID string) ToolUse { + input := make(map[string]any, len(event.Payload.Raw)+1) + for key, value := range event.Payload.Raw { + input[key] = value + } + input["event"] = eventType + return ToolUse{ + Tool: tools.EventToolName(eventType), + Input: input, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + RecordType: event.Type, + } +} + +func eventTokenUsage(event CodexEvent) CodexTokenUsage { + if event.Payload.Info == nil { + return CodexTokenUsage{} + } + return event.Payload.Info.LastTokenUsage +} + +func eventContextWindow(event CodexEvent) int { + if event.Payload.Info == nil { + return 0 + } + return event.Payload.Info.ModelContextWindow +} + +func codexNonCachedInputTokens(usage CodexTokenUsage) int { + if usage.CachedInputTokens <= 0 { + return usage.InputTokens + } + input := usage.InputTokens - usage.CachedInputTokens + if input < 0 { + return usage.InputTokens + } + return input +} + +// codexChatSlot is an emitted chat row that a twin may still merge into. +// families is the set of record families already folded in: Codex logs each +// logical message at most once per family, so a family arriving twice means a +// genuine repeat rather than a twin. Tracking the accumulated set rather than +// just the previous family is what keeps an (event_msg, response_item) pair +// repeated across N turns -- the same assistant verdict re-sent verbatim -- +// as N rows instead of collapsing it to one. +type codexChatSlot struct { + index int + priority int + families map[string]struct{} +} + +// dedupeCodexToolUses folds each message's twin records into a single row. +// Codex logs the same logical message twice, once as an event_msg and once as +// a response_item, milliseconds to hours apart -- so the twins cannot be +// matched on timestamp, only on content plus adjacency. +// +// records holds one entry per source record in emission order. Only the +// immediately preceding record is a merge candidate, so any intervening chat +// content keeps two identical messages apart. +func dedupeCodexToolUses(records [][]ToolUse) []ToolUse { + var output []ToolUse + var previous map[string]*codexChatSlot + for _, record := range records { + current := make(map[string]*codexChatSlot, len(record)) + for _, use := range record { + key, ok := codexChatDedupeKey(use) + if !ok { + output = append(output, use) + continue + } + if slot := mergeCodexTwin(output, previous[key], use); slot != nil { + current[key] = slot + continue + } + output = append(output, use) + current[key] = &codexChatSlot{ + index: len(output) - 1, + priority: codexRecordPriority(use.RecordType), + families: map[string]struct{}{codexRecordFamily(use.RecordType): {}}, + } + } + previous = current + } + return output +} + +// mergeCodexTwin folds use into slot when slot holds the same message logged by +// a different record family, returning the slot merged into (nil when it did +// not). The higher-priority record wins the row, so response_item content and +// the model stamped on it survive whichever twin arrived first. +func mergeCodexTwin(output []ToolUse, slot *codexChatSlot, use ToolUse) *codexChatSlot { + if slot == nil { + return nil + } + family := codexRecordFamily(use.RecordType) + if _, exists := slot.families[family]; exists { + return nil + } + if priority := codexRecordPriority(use.RecordType); priority > slot.priority { + output[slot.index] = use + slot.priority = priority + } + slot.families[family] = struct{}{} + return slot +} + +// codexRecordFamily reduces a RecordType to the stream that logged it. +func codexRecordFamily(recordType string) string { + switch { + case strings.HasPrefix(recordType, "response_item."): + return "response_item" + case strings.HasPrefix(recordType, "item.completed"): + return "item.completed" + case strings.HasPrefix(recordType, "event_msg."): + return "event_msg" + } + return recordType +} + +func codexChatDedupeKey(use ToolUse) (string, bool) { + switch use.Tool { + case "System", "User", "Assistant", "Reasoning", "Plan", "MemoryCitation": + default: + return "", false + } + text := codexDedupeText(use) + if text == "" { + return "", false + } + // No timestamp: twins are written anywhere from the same millisecond to + // hours apart (a resumed session re-logs its user message), so the only + // reliable discriminators are content, adjacency, and record family. + return strings.Join([]string{use.SessionID, use.Tool, text}, "\x00"), true +} + +func codexDedupeText(use ToolUse) string { + if text, ok := use.Input["text"].(string); ok { + return strings.TrimSpace(text) + } + if content, ok := use.Input["content"].(string); ok { + return strings.TrimSpace(content) + } + if use.Tool == "MemoryCitation" { + return fmt.Sprint(use.Input["citation_entries"], "\x00", use.Input["rollout_ids"]) + } + return "" +} + +// codexRecordPriority ranks the families by authority: response_item carries +// the canonical content, so it wins the row over its event_msg twin. +func codexRecordPriority(recordType string) int { + switch codexRecordFamily(recordType) { + case "response_item": + return 3 + case "item.completed": + return 2 + case "event_msg": + return 1 + default: + return 0 + } +} diff --git a/pkg/ai/history/codex_messages.go b/pkg/ai/history/codex_messages.go new file mode 100644 index 0000000..ad488c0 --- /dev/null +++ b/pkg/ai/history/codex_messages.go @@ -0,0 +1,371 @@ +package history + +import ( + "strconv" + "strings" + + "github.com/flanksource/captain/pkg/ai/assistanttags" +) + +func extractResponseItem(event CodexEvent, pendingCall map[string]CodexEvent, cwd, sessionID string) []ToolUse { + switch event.Payload.Type { + case "function_call", "custom_tool_call": + pendingCall[event.Payload.CallID] = event + return nil + case "tool_search_call": + pendingCall[event.Payload.CallID] = event + return nil + case "function_call_output", "custom_tool_call_output": + callEvent, ok := pendingCall[event.Payload.CallID] + if !ok { + return nil + } + delete(pendingCall, event.Payload.CallID) + return []ToolUse{buildToolUse(callEvent, event, cwd, sessionID)} + case "tool_search_output": + callEvent, ok := pendingCall[event.Payload.CallID] + if ok { + delete(pendingCall, event.Payload.CallID) + } + return buildToolSearchUses(callEvent, event, cwd, sessionID) + case "message": + var tool string + var text string + switch event.Payload.Role { + case "assistant": + text = codexContentText(event.Payload.Content, "output_text", "text") + return extractCodexAssistantText(text, event, cwd, sessionID, "response_item.message") + case "user": + text = codexContentText(event.Payload.Content, "input_text", "text") + if shell, ok := parseCodexUserShellCommand(text); ok { + return []ToolUse{buildCodexUserShellCommandUse(shell, event, cwd, sessionID, "response_item.message.user_shell_command")} + } + var ok bool + if tool, ok = codexUserMessageTool(text); !ok { + return nil + } + default: + return nil + } + if text == "" { + return nil + } + return []ToolUse{{ + Tool: tool, + Input: map[string]any{"text": text}, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + RecordType: "response_item.message", + }} + } + return nil +} + +func extractLiveItemCompleted(event CodexEvent, cwd, sessionID string) []ToolUse { + if event.Item == nil || event.Item.Role != "" && event.Item.Role != "assistant" { + return nil + } + text := event.Item.Text + if text == "" { + text = codexContentText(event.Item.Content, "output_text", "text") + } + if text == "" { + return nil + } + return extractCodexAssistantText(text, event, cwd, sessionID, "item.completed") +} + +func extractLiveError(event CodexEvent, cwd, sessionID string) []ToolUse { + var message string + if event.Error != nil { + message = event.Error.Message + } + if message == "" { + message = event.Message + } + message = NormalizeCodexError(message) + if message == "" { + return nil + } + return []ToolUse{{ + Tool: "ApiError", + Input: map[string]any{"error": message}, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + }} +} + +func extractEventMsg(event CodexEvent, cwd, sessionID string) []ToolUse { + switch event.Payload.Type { + case "agent_reasoning": + if event.Payload.Text == "" { + return nil + } + return []ToolUse{{ + Tool: "Reasoning", + Input: map[string]any{"text": event.Payload.Text}, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + RecordType: "event_msg.agent_reasoning", + }} + case "agent_message": + if event.Payload.Message == "" { + return nil + } + return extractCodexAssistantText(event.Payload.Message, event, cwd, sessionID, "event_msg.agent_message") + case "user_message": + text := firstNonEmpty(event.Payload.Message, event.Payload.Text) + if text == "" { + return nil + } + if shell, ok := parseCodexUserShellCommand(text); ok { + return []ToolUse{buildCodexUserShellCommandUse(shell, event, cwd, sessionID, "event_msg.user_shell_command")} + } + tool, ok := codexUserMessageTool(text) + if !ok { + return nil + } + return []ToolUse{{ + Tool: tool, + Input: map[string]any{"text": text}, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + RecordType: "event_msg.user_message", + }} + } + if event.Payload.Type == "" { + return nil + } + return []ToolUse{buildCodexEventUse(event, cwd, sessionID)} +} + +func extractCodexAssistantText(text string, event CodexEvent, cwd, sessionID, recordType string) []ToolUse { + segments := assistanttags.Parse(text) + uses := make([]ToolUse, 0, len(segments)) + for _, segment := range segments { + use := ToolUse{ + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + RecordType: recordType, + } + switch segment.Kind { + case assistanttags.SegmentPlan: + use.Tool = "Plan" + use.Input = map[string]any{"content": segment.Text, "tag": "proposed_plan"} + case assistanttags.SegmentMemoryCitation: + if segment.Citation == nil { + continue + } + use.Tool = "MemoryCitation" + use.Input = map[string]any{ + "event": "memory_citation", + "source": "codex", + "citation_entries": segment.Citation.CitationEntries, + "rollout_ids": segment.Citation.RolloutIDs, + } + case assistanttags.SegmentText: + use.Tool = "Assistant" + body := segment.Text + if summary, ok := assistanttags.EnvelopeSummary(body); ok { + body = summary + } + use.Input = map[string]any{"text": body} + default: + continue + } + uses = append(uses, use) + } + return uses +} + +func buildToolSearchUses(callEvent, outputEvent CodexEvent, cwd, sessionID string) []ToolUse { + names := codexToolSearchNames(outputEvent.Payload.Tools) + if len(names) == 0 { + return nil + } + input := map[string]any{"event": "deferred_tools_delta", "addedNames": names} + for key, value := range extractArgumentsMap(callEvent.Payload.Arguments) { + input[key] = value + } + return []ToolUse{{ + Tool: "DeferredToolsDelta", + Input: input, + Timestamp: outputEvent.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: firstNonEmpty(codexEventTurnID(outputEvent), codexEventTurnID(callEvent)), + ToolUseID: outputEvent.Payload.CallID, + Source: "codex", + RecordType: "response_item.tool_search_output", + }} +} + +func codexToolSearchNames(groups []CodexToolSearchNamespace) []string { + var names []string + seen := map[string]struct{}{} + for _, group := range groups { + for _, tool := range group.Tools { + name := strings.TrimSpace(tool.Name) + if name == "" { + continue + } + if _, ok := seen[name]; ok { + continue + } + seen[name] = struct{}{} + names = append(names, name) + } + } + return names +} + +func codexContentText(content []CodexContent, accepted ...string) string { + accept := make(map[string]struct{}, len(accepted)) + for _, contentType := range accepted { + accept[contentType] = struct{}{} + } + var text string + for _, item := range content { + if _, ok := accept[item.Type]; ok && item.Text != "" { + text += item.Text + } + } + return text +} + +type codexUserShellCommand struct { + Command string + ExitCode int + DurationSeconds float64 + Output string +} + +func parseCodexUserShellCommand(text string) (codexUserShellCommand, bool) { + const wrapperOpen = "" + const wrapperClose = "" + text = strings.TrimSpace(text) + if !strings.HasPrefix(text, wrapperOpen) || !strings.HasSuffix(text, wrapperClose) { + return codexUserShellCommand{}, false + } + body := strings.TrimSpace(strings.TrimSuffix(strings.TrimPrefix(text, wrapperOpen), wrapperClose)) + command, ok := codexWrappedSection(body, "command") + if !ok || strings.TrimSpace(command) == "" { + return codexUserShellCommand{}, false + } + result, ok := codexWrappedSection(body, "result") + if !ok { + return codexUserShellCommand{}, false + } + lines := strings.Split(strings.TrimSpace(result), "\n") + if len(lines) < 3 { + return codexUserShellCommand{}, false + } + exitText, ok := strings.CutPrefix(strings.TrimSuffix(lines[0], "\r"), "Exit code:") + if !ok { + return codexUserShellCommand{}, false + } + exitCode, err := strconv.Atoi(strings.TrimSpace(exitText)) + if err != nil { + return codexUserShellCommand{}, false + } + durationText, ok := strings.CutPrefix(strings.TrimSuffix(lines[1], "\r"), "Duration:") + if !ok { + return codexUserShellCommand{}, false + } + durationText, ok = strings.CutSuffix(strings.TrimSpace(durationText), " seconds") + if !ok { + return codexUserShellCommand{}, false + } + durationSeconds, err := strconv.ParseFloat(strings.TrimSpace(durationText), 64) + if err != nil || strings.TrimSuffix(lines[2], "\r") != "Output:" { + return codexUserShellCommand{}, false + } + return codexUserShellCommand{ + Command: strings.TrimSpace(command), + ExitCode: exitCode, + DurationSeconds: durationSeconds, + Output: strings.TrimSpace(strings.Join(lines[3:], "\n")), + }, true +} + +func codexWrappedSection(text, name string) (string, bool) { + open := "<" + name + ">" + close := "" + start := strings.Index(text, open) + if start < 0 { + return "", false + } + start += len(open) + end := strings.Index(text[start:], close) + if end < 0 { + return "", false + } + return text[start : start+end], true +} + +func buildCodexUserShellCommandUse(cmd codexUserShellCommand, event CodexEvent, cwd, sessionID, recordType string) ToolUse { + status := "success" + if cmd.ExitCode != 0 { + status = "failed" + } + input := map[string]any{ + "event": "user_shell_command", + "command": cmd.Command, + "exit_code": cmd.ExitCode, + "duration_seconds": cmd.DurationSeconds, + "duration_ms": cmd.DurationSeconds * 1000, + "output": cmd.Output, + "stdout": cmd.Output, + "status": status, + } + return ToolUse{ + Tool: "UserShellCommand", + Input: input, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + Response: cmd.Output, + RecordType: recordType, + } +} + +func isCodexInternalUserText(text string) bool { + text = strings.TrimSpace(text) + return strings.HasPrefix(text, "") || + strings.HasPrefix(text, "") || + strings.HasPrefix(text, "") || + strings.HasPrefix(text, "") || + strings.HasPrefix(text, "") +} + +func codexUserMessageTool(text string) (string, bool) { + text = strings.TrimSpace(text) + if text == "" { + return "", false + } + if strings.HasPrefix(text, "# AGENTS.md") || + strings.HasPrefix(text, "") && strings.Contains(text, "# AGENTS.md instructions") { + return "System", true + } + if isCodexInternalUserText(text) { + return "", false + } + return "User", true +} diff --git a/pkg/ai/history/codex_normalize.go b/pkg/ai/history/codex_normalize.go new file mode 100644 index 0000000..410e26b --- /dev/null +++ b/pkg/ai/history/codex_normalize.go @@ -0,0 +1,241 @@ +package history + +import ( + "strings" + "time" + + "github.com/segmentio/encoding/json" +) + +// CodexToolCall is the transport-neutral input for Codex tool normalization. +// Rollout JSONL, codex exec JSON, and app-server notifications decode their +// native envelopes independently and converge on this shape. +type CodexToolCall struct { + Name string + Namespace string + Arguments json.RawMessage + Command string + Input map[string]any + Timestamp *time.Time + CWD string + SessionID string + TurnID string + ID string + Response string + Model string + ReasoningEffort string + RecordType string +} + +// NormalizeCodexToolCall maps a native Codex call into the canonical history +// tool shape consumed by both live providers and transcript rendering. +func NormalizeCodexToolCall(call CodexToolCall) ToolUse { + input := codexArgumentsMap(call.Arguments) + for key, value := range call.Input { + input[key] = value + } + + command := firstNonEmpty(call.Command, stringValue(input["command"]), stringValue(input["cmd"])) + if command == "" && call.Name == "" { + command = codexScalarArgument(call.Arguments) + delete(input, "arguments") + } + if command != "" { + delete(input, "cmd") + input["command"] = command + } + + tool := call.Name + switch call.Name { + case "update_plan": + tool = "TodoWrite" + if plan, ok := input["plan"]; ok { + input["todos"] = plan + } + case "request_user_input": + tool = "AskUserQuestion" + case "wait": + tool = "Wait" + case "spawn_agent": + tool = "Agent" + case "wait_agent": + tool = "CollabWaiting" + input["event"] = call.Name + case "close_agent": + tool = "CollabClose" + input["event"] = call.Name + case "send_input", "resume_agent": + tool = "CollabAgentInteraction" + input["event"] = call.Name + default: + if command != "" { + tool = "Bash" + } + } + if tool == "" { + tool = "Bash" + } + + use := ToolUse{ + Tool: tool, + Input: input, + Timestamp: call.Timestamp, + CWD: call.CWD, + SessionID: call.SessionID, + TurnID: call.TurnID, + ToolUseID: call.ID, + Source: "codex", + Model: call.Model, + ReasoningEffort: call.ReasoningEffort, + Namespace: call.Namespace, + Response: call.Response, + RecordType: call.RecordType, + } + if tool == "Agent" { + use.AgentType, _ = input["agent_type"].(string) + use.AgentDesc, _ = input["message"].(string) + input["subagent_type"] = use.AgentType + input["description"] = use.AgentDesc + input["prompt"] = use.AgentDesc + agentID, nickname := codexAgentOutput(call.Response) + use.AgentID = agentID + if use.AgentID != "" { + input["agent_id"] = use.AgentID + } + if nickname != "" { + input["nickname"] = nickname + } + } + return use +} + +func buildToolUse(callEvent, outputEvent CodexEvent, cwd, sessionID string) ToolUse { + timestamp := callEvent.Time() + if timestamp == nil { + timestamp = outputEvent.Time() + } + var input map[string]any + if callEvent.Payload.Input != "" { + input = map[string]any{"input": callEvent.Payload.Input} + } + return NormalizeCodexToolCall(CodexToolCall{ + Name: callEvent.Payload.Name, + Namespace: callEvent.Payload.Namespace, + Arguments: callEvent.Payload.Arguments, + Input: input, + Timestamp: timestamp, + CWD: cwd, + SessionID: sessionID, + TurnID: firstNonEmpty(codexEventTurnID(callEvent), codexEventTurnID(outputEvent)), + ID: callEvent.Payload.CallID, + Response: extractCommandOutput(CodexOutputText(outputEvent.Payload.Output)), + RecordType: "response_item." + callEvent.Payload.Type, + }) +} + +func codexArgumentsMap(arguments json.RawMessage) map[string]any { + raw := normalizeCodexArguments(arguments) + var input map[string]any + if raw != "" { + _ = json.Unmarshal([]byte(raw), &input) + } + if input != nil { + return input + } + + var value any + if len(arguments) > 0 && string(arguments) != "null" && json.Unmarshal(arguments, &value) == nil { + return map[string]any{"arguments": value} + } + return map[string]any{} +} + +func extractArgumentsMap(arguments json.RawMessage) map[string]any { + return codexArgumentsMap(arguments) +} + +func normalizeCodexArguments(arguments json.RawMessage) string { + if len(arguments) == 0 || string(arguments) == "null" { + return "" + } + var value string + if json.Unmarshal(arguments, &value) == nil { + return value + } + return string(arguments) +} + +func codexScalarArgument(arguments json.RawMessage) string { + raw := normalizeCodexArguments(arguments) + if raw == "" || strings.HasPrefix(strings.TrimSpace(raw), "{") { + return "" + } + return raw +} + +func stringValue(value any) string { + text, _ := value.(string) + return text +} + +func codexAgentOutput(response string) (agentID, nickname string) { + var payload struct { + AgentID string `json:"agent_id"` + Nickname string `json:"nickname"` + } + if json.Unmarshal([]byte(strings.TrimSpace(response)), &payload) == nil { + return payload.AgentID, payload.Nickname + } + return "", "" +} + +// NormalizeCodexError unwraps the JSON-encoded error envelope Codex sometimes +// stores inside a message field. +func NormalizeCodexError(raw string) string { + raw = strings.TrimSpace(raw) + if raw == "" || !strings.HasPrefix(raw, "{") { + return raw + } + var inner CodexErrorBlock + wrapper := struct { + Error *CodexErrorBlock `json:"error"` + Message string `json:"message"` + }{Error: &inner} + if json.Unmarshal([]byte(raw), &wrapper) != nil { + return raw + } + if wrapper.Error != nil && wrapper.Error.Message != "" { + return wrapper.Error.Message + } + if wrapper.Message != "" { + return wrapper.Message + } + return raw +} + +// CodexOutputText normalizes the scalar-string and ordered content-block forms +// Codex uses for function-call output records. +func CodexOutputText(raw json.RawMessage) string { + text := strings.TrimSpace(string(raw)) + if text == "" || text == "null" { + return "" + } + var scalar string + if json.Unmarshal(raw, &scalar) == nil { + return scalar + } + var content []CodexContent + if json.Unmarshal(raw, &content) == nil { + if combined := codexContentText(content, "input_text", "output_text", "text"); combined != "" { + return combined + } + } + return text +} + +func extractCommandOutput(raw string) string { + if _, after, ok := strings.Cut(raw, "Output:\n"); ok { + return after + } + return raw +} diff --git a/pkg/ai/history/codex_normalize_ginkgo_test.go b/pkg/ai/history/codex_normalize_ginkgo_test.go new file mode 100644 index 0000000..cc7b248 --- /dev/null +++ b/pkg/ai/history/codex_normalize_ginkgo_test.go @@ -0,0 +1,104 @@ +package history + +import ( + "strings" + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + . "github.com/onsi/gomega/gstruct" +) + +func TestCodexNormalization(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Codex Normalization Suite") +} + +var _ = Describe("NormalizeCodexToolCall", func() { + It("normalizes command calls from every Codex transport to Bash", func() { + uses := []ToolUse{ + NormalizeCodexToolCall(CodexToolCall{ + Name: "shell", + Arguments: []byte(`{"command":"pwd"}`), + ID: "call-1", + SessionID: "thread-1", + }), + NormalizeCodexToolCall(CodexToolCall{ + Command: "pwd", + ID: "call-1", + SessionID: "thread-1", + }), + } + + Expect(uses).To(ConsistOf( + MatchFields(IgnoreExtras, Fields{ + "Tool": Equal("Bash"), + "Input": Equal(map[string]any{"command": "pwd"}), + "ToolUseID": Equal("call-1"), + "SessionID": Equal("thread-1"), + "Source": Equal("codex"), + }), + MatchFields(IgnoreExtras, Fields{ + "Tool": Equal("Bash"), + "Input": Equal(map[string]any{"command": "pwd"}), + "ToolUseID": Equal("call-1"), + "SessionID": Equal("thread-1"), + "Source": Equal("codex"), + }), + )) + }) + + It("preserves named tool arguments without classifying them as shell commands", func() { + use := NormalizeCodexToolCall(CodexToolCall{ + Name: "search", + Namespace: "catalog", + Arguments: []byte(`{"query":"captain"}`), + ID: "tool-1", + }) + + Expect(use).To(MatchFields(IgnoreExtras, Fields{ + "Tool": Equal("search"), + "Input": Equal(map[string]any{"query": "captain"}), + "Namespace": Equal("catalog"), + "ToolUseID": Equal("tool-1"), + })) + }) + + It("preserves scalar arguments for named tools", func() { + use := NormalizeCodexToolCall(CodexToolCall{ + Name: "search", + Arguments: []byte(`"catalog"`), + }) + + Expect(use.Input).To(Equal(map[string]any{"arguments": "catalog"})) + }) + + It("extracts current custom tool calls from rollout history", func() { + stream := strings.Join([]string{ + `{"timestamp":"2026-07-16T10:00:00Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo"}}`, + `{"timestamp":"2026-07-16T10:00:01Z","type":"response_item","payload":{"type":"custom_tool_call","name":"exec","call_id":"call-patch","input":"const patch = \"*** Begin Patch\\n*** Update File: app.go\\n*** End Patch\"; tools.apply_patch(patch);"}}`, + `{"timestamp":"2026-07-16T10:00:02Z","type":"response_item","payload":{"type":"custom_tool_call_output","call_id":"call-patch","output":"Success"}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + Expect(err).ToNot(HaveOccurred()) + Expect(uses).To(ConsistOf(MatchFields(IgnoreExtras, Fields{ + "Tool": Equal("exec"), + "Input": HaveKeyWithValue("input", ContainSubstring("*** Update File: app.go")), + "SessionID": Equal("session-1"), + "ToolUseID": Equal("call-patch"), + "RecordType": Equal("response_item.custom_tool_call"), + }))) + }) + + DescribeTable("normalizes nested error envelopes", + func(input, expected string) { + Expect(NormalizeCodexError(input)).To(Equal(expected)) + }, + Entry("plain", "boom", "boom"), + Entry("empty", "", ""), + Entry("nested error", `{"type":"error","status":400,"error":{"type":"invalid_request_error","message":"bad model"}}`, "bad model"), + Entry("top-level message", `{"message":"top"}`, "top"), + Entry("non-json passthrough", "not { json", "not { json"), + ) +}) diff --git a/pkg/ai/history/codex_parser.go b/pkg/ai/history/codex_parser.go index 107fcaf..c4a4493 100644 --- a/pkg/ai/history/codex_parser.go +++ b/pkg/ai/history/codex_parser.go @@ -2,12 +2,14 @@ package history import ( "bufio" - "github.com/segmentio/encoding/json" "io" "os" "path/filepath" "strings" "time" + + "github.com/flanksource/captain/pkg/api" + "github.com/segmentio/encoding/json" ) func ParseCodexLine(line string) (CodexEvent, error) { @@ -25,8 +27,7 @@ func ExtractCodexToolUses(sessionFile string) ([]ToolUse, error) { return ExtractCodexToolUsesFromReader(file) } -// CodexSessionInfo is the first-line summary of a codex rollout file: -// session id, cwd, cli version, model provider, git branch, etc. +// CodexSessionInfo is the first-line summary of a codex rollout file. type CodexSessionInfo struct { ID string CWD string @@ -36,13 +37,12 @@ type CodexSessionInfo struct { GitBranch string GitCommit string StartedAt *time.Time - Model string // from the first turn_context payload - ReasoningEffort string // "low", "medium", "high" — per turn_context + Model string + ReasoningEffort string } -// ReadCodexSessionMeta parses only the leading session metadata needed for -// cheap history candidate filtering. It intentionally stops before scanning -// turns or response items. +// ReadCodexSessionMeta parses only the leading metadata needed for cheap +// history candidate filtering. func ReadCodexSessionMeta(sessionFile string) (*CodexSessionInfo, error) { file, err := os.Open(sessionFile) if err != nil { @@ -63,19 +63,8 @@ func ReadCodexSessionMeta(sessionFile string) (*CodexSessionInfo, error) { } switch event.Type { case "session_meta": - info := &CodexSessionInfo{ - ID: event.Payload.ID, - CWD: event.Payload.CWD, - CLIVersion: event.Payload.CLIVersion, - ModelProvider: event.Payload.ModelProvider, - Originator: event.Payload.Originator, - StartedAt: event.Time(), - } - if event.Payload.Git != nil { - info.GitBranch = event.Payload.Git.Branch - info.GitCommit = event.Payload.Git.CommitHash - } - return info, nil + info := codexSessionInfo(event) + return &info, nil case "thread.started": return &CodexSessionInfo{ID: event.ThreadID, StartedAt: event.Time()}, nil case "turn_context", "response_item", "event_msg", "item.completed", "turn.failed", "error": @@ -88,8 +77,7 @@ func ReadCodexSessionMeta(sessionFile string) (*CodexSessionInfo, error) { return nil, nil } -// ReadCodexSessionInfo parses just the leading `session_meta` event from a -// codex rollout file. Returns nil if the file has no session_meta header. +// ReadCodexSessionInfo parses the leading metadata and first turn context. func ReadCodexSessionInfo(sessionFile string) (*CodexSessionInfo, error) { file, err := os.Open(sessionFile) if err != nil { @@ -111,20 +99,9 @@ func ReadCodexSessionInfo(sessionFile string) (*CodexSessionInfo, error) { } switch event.Type { case "session_meta": - if info != nil { - continue - } - info = &CodexSessionInfo{ - ID: event.Payload.ID, - CWD: event.Payload.CWD, - CLIVersion: event.Payload.CLIVersion, - ModelProvider: event.Payload.ModelProvider, - Originator: event.Payload.Originator, - StartedAt: event.Time(), - } - if event.Payload.Git != nil { - info.GitBranch = event.Payload.Git.Branch - info.GitCommit = event.Payload.Git.CommitHash + if info == nil { + value := codexSessionInfo(event) + info = &value } case "thread.started": if info == nil { @@ -137,10 +114,10 @@ func ReadCodexSessionInfo(sessionFile string) (*CodexSessionInfo, error) { if info == nil { info = &CodexSessionInfo{StartedAt: event.Time()} } - if info.Model == "" && event.Payload.Model != "" { + if info.Model == "" { info.Model = event.Payload.Model } - if info.ReasoningEffort == "" && event.Payload.Effort != "" { + if info.ReasoningEffort == "" { info.ReasoningEffort = event.Payload.Effort } if info.Model != "" && info.ReasoningEffort != "" { @@ -154,338 +131,124 @@ func ReadCodexSessionInfo(sessionFile string) (*CodexSessionInfo, error) { return info, nil } -func ExtractCodexToolUsesFromReader(r io.Reader) ([]ToolUse, error) { - scanner := bufio.NewScanner(r) - scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024) +func codexSessionInfo(event CodexEvent) CodexSessionInfo { + info := CodexSessionInfo{ + ID: event.Payload.ID, + CWD: event.Payload.CWD, + CLIVersion: event.Payload.CLIVersion, + ModelProvider: event.Payload.ModelProvider, + Originator: event.Payload.Originator, + StartedAt: event.Time(), + } + if event.Payload.Git != nil { + info.GitBranch = event.Payload.Git.Branch + info.GitCommit = event.Payload.Git.CommitHash + } + return info +} - var ( - toolUses []ToolUse - sessionCWD string - sessionID string - currentModel string - currentEff string - pendingCall = make(map[string]CodexEvent) - ) +func IsCodexAutoReviewSession(sessionFile string) (bool, error) { + info, err := ReadCodexSessionInfo(sessionFile) + if err != nil || info == nil { + return false, err + } + return IsCodexAutoReviewModel(info.Model), nil +} + +func IsCodexAutoReviewModel(model string) bool { + return strings.EqualFold(strings.TrimSpace(model), api.CodexAutoReviewModel) +} +func ExtractCodexToolUsesFromReader(reader io.Reader) ([]ToolUse, error) { + scanner := bufio.NewScanner(reader) + scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024) + + // records holds the uses of one source record per entry, in emission order. + // Dedupe needs those boundaries to merge a message's twin records: the + // boundary cannot be recovered afterwards, because distinct adjacent records + // routinely share a RecordType and a millisecond. + var records [][]ToolUse + var sessionCWD, sessionID, currentTurn, currentModel, currentEffort string + pendingCall := make(map[string]CodexEvent) + var reasoning codexReasoningCollapser stamp := func(uses []ToolUse) []ToolUse { - for i := range uses { - if uses[i].Model == "" { - uses[i].Model = currentModel + for index := range uses { + if uses[index].TurnID == "" { + uses[index].TurnID = currentTurn + } + if uses[index].Model == "" { + uses[index].Model = currentModel } - if uses[i].ReasoningEffort == "" { - uses[i].ReasoningEffort = currentEff + if uses[index].ReasoningEffort == "" { + uses[index].ReasoningEffort = currentEffort } } return uses } + // A record that emits nothing -- an accumulating reasoning record, an + // unpaired function_call half, filtered internal user text -- must not count + // as intervening content, or it would break a twin pair apart. + add := func(uses []ToolUse) { + if len(uses) == 0 { + return + } + records = append(records, stamp(uses)) + } for scanner.Scan() { line := strings.TrimSpace(scanner.Text()) if line == "" { continue } - event, err := ParseCodexLine(line) if err != nil { log.Debugf("Error parsing codex line: %v", err) continue } - switch event.Type { case "session_meta": sessionCWD = event.Payload.CWD sessionID = event.Payload.ID - case "turn_context": - if event.Payload.Model != "" { - currentModel = event.Payload.Model - } - if event.Payload.Effort != "" { - currentEff = event.Payload.Effort + // Flush before the turn's model/effort roll over so the closing + // span is stamped with the turn it belongs to. turn_context is the + // only reliable turn boundary: older rollouts carry no turn_id at + // all, so keying the span on turn_id alone would fold a whole + // multi-turn session into one bogus span. + add(reasoning.flush()) + currentTurn = firstNonEmpty(event.Payload.TurnID, currentTurn) + currentModel = firstNonEmpty(event.Payload.Model, currentModel) + if IsCodexAutoReviewModel(currentModel) { + return nil, nil } - + currentEffort = firstNonEmpty(event.Payload.Effort, currentEffort) case "response_item": - toolUses = append(toolUses, stamp(extractResponseItem(event, pendingCall, sessionCWD, sessionID))...) - + if event.Payload.Type == "reasoning" { + add(reasoning.observe(event, currentTurn, sessionCWD, sessionID)) + continue + } + add(extractResponseItem(event, pendingCall, sessionCWD, sessionID)) case "event_msg": - toolUses = append(toolUses, stamp(extractEventMsg(event, sessionCWD, sessionID))...) - - // --- Newer dotted-name `codex exec --json` schema --- + currentTurn = firstNonEmpty(event.Payload.TurnID, currentTurn) + add(extractEventMsg(event, sessionCWD, sessionID)) + case "world_state": + add([]ToolUse{buildCodexTopLevelEventUse(event, "world_state", sessionCWD, sessionID)}) case "thread.started": - if event.ThreadID != "" { - sessionID = event.ThreadID - } - + sessionID = firstNonEmpty(event.ThreadID, sessionID) case "item.completed": - toolUses = append(toolUses, stamp(extractLiveItemCompleted(event, sessionCWD, sessionID))...) - + add(extractLiveItemCompleted(event, sessionCWD, sessionID)) case "turn.failed", "error": - toolUses = append(toolUses, stamp(extractLiveError(event, sessionCWD, sessionID))...) - - // turn.started, turn.completed, item.started, item.delta carry - // no tool-use information for our history view; intentionally - // dropped to keep parity with the legacy schema. + add(extractLiveError(event, sessionCWD, sessionID)) } } - + add(reasoning.flush()) for _, callEvent := range pendingCall { - toolUses = append(toolUses, stamp([]ToolUse{buildToolUse(callEvent, CodexEvent{}, sessionCWD, sessionID)})...) + add([]ToolUse{buildToolUse(callEvent, CodexEvent{}, sessionCWD, sessionID)}) } - if err := scanner.Err(); err != nil { return nil, err } - return toolUses, nil -} - -func extractResponseItem(event CodexEvent, pendingCall map[string]CodexEvent, cwd, sessionID string) []ToolUse { - switch event.Payload.Type { - case "function_call": - pendingCall[event.Payload.CallID] = event - return nil - - case "function_call_output": - callEvent, ok := pendingCall[event.Payload.CallID] - if !ok { - return nil - } - delete(pendingCall, event.Payload.CallID) - return []ToolUse{buildToolUse(callEvent, event, cwd, sessionID)} - - case "reasoning": - var summaries []CodexReasoningSummary - if len(event.Payload.Summary) > 0 { - _ = json.Unmarshal(event.Payload.Summary, &summaries) - } - var text string - for _, s := range summaries { - if s.Text != "" { - text = s.Text - } - } - if text == "" { - return nil - } - return []ToolUse{{ - Tool: "Reasoning", - Input: map[string]any{"text": text}, - Timestamp: event.Time(), - CWD: cwd, - SessionID: sessionID, - Source: "codex", - }} - - case "message": - if event.Payload.Role != "assistant" { - return nil - } - var text string - for _, c := range event.Payload.Content { - if c.Type == "output_text" && c.Text != "" { - text += c.Text - } - } - if text == "" { - return nil - } - return []ToolUse{{ - Tool: "Assistant", - Input: map[string]any{"text": text}, - Timestamp: event.Time(), - CWD: cwd, - SessionID: sessionID, - Source: "codex", - }} - } - return nil -} - -// extractLiveItemCompleted handles `item.completed` events from the newer -// `codex exec --json` schema. The item carries either a top-level Text field -// or a Content array shaped like response_item.message. -func extractLiveItemCompleted(event CodexEvent, cwd, sessionID string) []ToolUse { - if event.Item == nil { - return nil - } - text := event.Item.Text - if text == "" { - for _, c := range event.Item.Content { - if c.Type == "output_text" && c.Text != "" { - text += c.Text - } - } - } - if text == "" { - return nil - } - return []ToolUse{{ - Tool: "Assistant", - Input: map[string]any{"text": text}, - Timestamp: event.Time(), - CWD: cwd, - SessionID: sessionID, - Source: "codex", - }} -} - -// extractLiveError surfaces a `turn.failed` or top-level `error` event as an -// ApiError tool use, peeling one layer of stringified-JSON nesting that -// codex sometimes uses when relaying provider errors back to the user. -func extractLiveError(event CodexEvent, cwd, sessionID string) []ToolUse { - var msg string - if event.Error != nil { - msg = event.Error.Message - } - if msg == "" { - msg = event.Message - } - msg = unwrapCodexErrorMessage(msg) - if msg == "" { - return nil - } - return []ToolUse{{ - Tool: "ApiError", - Input: map[string]any{"error": msg}, - Timestamp: event.Time(), - CWD: cwd, - SessionID: sessionID, - Source: "codex", - }} -} - -func unwrapCodexErrorMessage(raw string) string { - raw = strings.TrimSpace(raw) - if raw == "" || !strings.HasPrefix(raw, "{") { - return raw - } - var inner CodexErrorBlock - wrapper := struct { - Error *CodexErrorBlock `json:"error"` - Message string `json:"message"` - }{Error: &inner} - if err := json.Unmarshal([]byte(raw), &wrapper); err != nil { - return raw - } - if wrapper.Error != nil && wrapper.Error.Message != "" { - return wrapper.Error.Message - } - if wrapper.Message != "" { - return wrapper.Message - } - return raw -} - -func extractEventMsg(event CodexEvent, cwd, sessionID string) []ToolUse { - switch event.Payload.Type { - case "agent_reasoning": - if event.Payload.Text == "" { - return nil - } - return []ToolUse{{ - Tool: "Reasoning", - Input: map[string]any{"text": event.Payload.Text}, - Timestamp: event.Time(), - CWD: cwd, - SessionID: sessionID, - Source: "codex", - }} - - case "agent_message": - if event.Payload.Message == "" { - return nil - } - return []ToolUse{{ - Tool: "Assistant", - Input: map[string]any{"text": event.Payload.Message}, - Timestamp: event.Time(), - CWD: cwd, - SessionID: sessionID, - Source: "codex", - }} - } - return nil -} - -func buildToolUse(callEvent, outputEvent CodexEvent, cwd, sessionID string) ToolUse { - var response string - if outputEvent.Payload.Output != "" { - response = extractCommandOutput(outputEvent.Payload.Output) - } - ts := callEvent.Time() - if ts == nil { - ts = outputEvent.Time() - } - - switch callEvent.Payload.Name { - case "update_plan": - args := extractArgumentsMap(callEvent.Payload.Arguments) - if plan, ok := args["plan"]; ok { - args["todos"] = plan - } - return ToolUse{ - Tool: "TodoWrite", - Input: args, - Timestamp: ts, - CWD: cwd, - SessionID: sessionID, - ToolUseID: callEvent.Payload.CallID, - Source: "codex", - Response: response, - } - case "request_user_input": - return ToolUse{ - Tool: "AskUserQuestion", - Input: extractArgumentsMap(callEvent.Payload.Arguments), - Timestamp: ts, - CWD: cwd, - SessionID: sessionID, - ToolUseID: callEvent.Payload.CallID, - Source: "codex", - Response: response, - } - } - - input := map[string]any{ - "command": extractCommand(callEvent.Payload.Arguments), - } - return ToolUse{ - Tool: "Bash", - Input: input, - Timestamp: ts, - CWD: cwd, - SessionID: sessionID, - ToolUseID: callEvent.Payload.CallID, - Source: "codex", - Response: response, - } -} - -func extractArgumentsMap(argsJSON string) map[string]any { - var args map[string]any - if argsJSON == "" || json.Unmarshal([]byte(argsJSON), &args) != nil { - return map[string]any{} - } - return args -} - -func extractCommand(argsJSON string) string { - if argsJSON == "" { - return "" - } - var args struct { - Cmd string `json:"cmd"` - } - if json.Unmarshal([]byte(argsJSON), &args) == nil && args.Cmd != "" { - return args.Cmd - } - return argsJSON -} - -func extractCommandOutput(raw string) string { - if _, after, ok := strings.Cut(raw, "Output:\n"); ok { - return after - } - return raw + return dedupeCodexToolUses(records), nil } func FindCodexSessionFiles() ([]string, error) { @@ -493,18 +256,13 @@ func FindCodexSessionFiles() ([]string, error) { if err != nil { return nil, err } - sessionsDir := filepath.Join(home, ".codex", "sessions") if _, err := os.Stat(sessionsDir); os.IsNotExist(err) { return nil, nil } - var files []string err = filepath.Walk(sessionsDir, func(path string, info os.FileInfo, err error) error { - if err != nil { - return nil - } - if !info.IsDir() && strings.HasSuffix(path, ".jsonl") { + if err == nil && !info.IsDir() && strings.HasSuffix(path, ".jsonl") { files = append(files, path) } return nil diff --git a/pkg/ai/history/codex_parser_test.go b/pkg/ai/history/codex_parser_test.go index 879bea2..f3b0f80 100644 --- a/pkg/ai/history/codex_parser_test.go +++ b/pkg/ai/history/codex_parser_test.go @@ -2,8 +2,12 @@ package history import ( "os" + "path/filepath" "strings" "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestExtractCodexToolUses_LiveDottedSchema(t *testing.T) { @@ -88,6 +92,141 @@ func TestExtractCodexToolUses_ModelAndEffortStamping(t *testing.T) { } } +func TestExtractCodexToolUses_IgnoresAutoReviewSession(t *testing.T) { + stream := strings.Join([]string{ + `{"timestamp":"2026-07-13T09:00:00Z","type":"session_meta","payload":{"id":"review-1","cwd":"/repo","source":{"subagent":{"other":"guardian"}}}}`, + `{"timestamp":"2026-07-13T09:00:01Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"review this command"}]}}`, + `{"timestamp":"2026-07-13T09:00:02Z","type":"turn_context","payload":{"turn_id":"turn-review","model":"codex-auto-review","effort":"low"}}`, + `{"timestamp":"2026-07-13T09:00:03Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"approved"}]}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + require.NoError(t, err) + assert.Empty(t, uses) +} + +func TestExtractCodexToolUses_DoesNotIgnoreMentionOfAutoReview(t *testing.T) { + stream := strings.Join([]string{ + `{"timestamp":"2026-07-13T09:00:00Z","type":"session_meta","payload":{"id":"user-1","cwd":"/repo"}}`, + `{"timestamp":"2026-07-13T09:00:01Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"Ignore codex-auto-review sessions"}]}}`, + `{"timestamp":"2026-07-13T09:00:02Z","type":"turn_context","payload":{"turn_id":"turn-user","model":"gpt-5.6-sol","effort":"high"}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + require.NoError(t, err) + require.Len(t, uses, 1) + assert.Equal(t, "User", uses[0].Tool) +} + +func TestIsCodexAutoReviewSession(t *testing.T) { + path := filepath.Join(t.TempDir(), "rollout.jsonl") + stream := strings.Join([]string{ + `{"timestamp":"2026-07-13T09:00:00Z","type":"session_meta","payload":{"id":"review-1","cwd":"/repo"}}`, + `{"timestamp":"2026-07-13T09:00:01Z","type":"turn_context","payload":{"model":"codex-auto-review","effort":"low"}}`, + }, "\n") + require.NoError(t, os.WriteFile(path, []byte(stream), 0o600)) + + ignored, err := IsCodexAutoReviewSession(path) + require.NoError(t, err) + assert.True(t, ignored) +} + +func TestExtractCodexToolUses_ParsesWaitCallWithContentOutput(t *testing.T) { + stream := strings.Join([]string{ + `{"timestamp":"2026-07-13T09:31:13.359Z","type":"session_meta","payload":{"id":"019f5a68-c638-7fe2-b0d0-c1d8a3c4de55","cwd":"/repo"}}`, + `{"timestamp":"2026-07-13T09:31:13.362Z","type":"turn_context","payload":{"turn_id":"019f5ad0-fc23-7b92-8b13-a9f50d903704","model":"gpt-5.6-sol","effort":"max"}}`, + `{"timestamp":"2026-07-13T09:41:33.611Z","type":"response_item","payload":{"type":"function_call","name":"wait","arguments":"{\"cell_id\":\"214\",\"yield_time_ms\":20000,\"max_tokens\":5000}","call_id":"call-wait"}}`, + `{"timestamp":"2026-07-13T09:41:41.017Z","type":"response_item","payload":{"type":"function_call_output","call_id":"call-wait","output":[{"type":"input_text","text":"Script completed\nWall time 7.4 seconds\nOutput:\n"},{"type":"input_text","text":"evaluation failed\n"},{"type":"input_text","text":"exit=1"}]}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + if err != nil { + t.Fatalf("ExtractCodexToolUsesFromReader: %v", err) + } + if len(uses) != 1 { + t.Fatalf("uses = %+v, want one Wait", uses) + } + use := uses[0] + if use.Tool != "Wait" || use.ToolUseID != "call-wait" { + t.Fatalf("use = %+v, want first-class Wait", use) + } + if use.Input["cell_id"] != "214" || use.Input["yield_time_ms"] != float64(20000) || use.Input["max_tokens"] != float64(5000) { + t.Fatalf("input = %#v, want structured wait arguments", use.Input) + } + if use.Response != "evaluation failed\nexit=1" { + t.Fatalf("response = %q", use.Response) + } + if use.TurnID != "019f5ad0-fc23-7b92-8b13-a9f50d903704" || use.Model != "gpt-5.6-sol" || use.ReasoningEffort != "max" { + t.Fatalf("metadata = %+v", use) + } + if use.RecordType != "response_item.function_call" { + t.Fatalf("record type = %q", use.RecordType) + } +} + +func TestExtractCodexToolUses_ParsesUserShellCommandMessage(t *testing.T) { + stream := strings.Join([]string{ + `{"timestamp":"2026-07-13T06:13:59Z","type":"session_meta","payload":{"id":"sess-shell","cwd":"/repo"}}`, + `{"timestamp":"2026-07-13T06:14:00Z","type":"turn_context","payload":{"turn_id":"turn-shell","model":"gpt-5.5-codex","effort":"high"}}`, + `{"timestamp":"2026-07-13T06:14:01Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"\n\ngavel proc restart\n\n\nExit code: 1\nDuration: 2.9909 seconds\nOutput:\n\u001b[36mdev | \u001b[0msignal: killed\n\u001b[36mdev | \u001b[0mKill sent but port 8088 is still bound\n\n"}]}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + if err != nil { + t.Fatalf("ExtractCodexToolUsesFromReader: %v", err) + } + if len(uses) != 1 { + t.Fatalf("uses = %+v, want exactly one UserShellCommand", uses) + } + use := uses[0] + if use.Tool != "UserShellCommand" || use.Input["event"] != "user_shell_command" { + t.Fatalf("use = %+v, want UserShellCommand event", use) + } + if use.Input["command"] != "gavel proc restart" || use.Input["exit_code"] != 1 || use.Input["status"] != "failed" { + t.Fatalf("command result = %+v", use.Input) + } + if use.Input["duration_seconds"] != 2.9909 || use.Input["duration_ms"] != 2990.9 { + t.Fatalf("duration = %v seconds / %v ms", use.Input["duration_seconds"], use.Input["duration_ms"]) + } + output, _ := use.Input["output"].(string) + wantOutput := "\x1b[36mdev | \x1b[0msignal: killed\n\x1b[36mdev | \x1b[0mKill sent but port 8088 is still bound" + if output != wantOutput { + t.Fatalf("output = %q, want exact preserved output %q", output, wantOutput) + } + if use.Response != output || use.TurnID != "turn-shell" || use.Model != "gpt-5.5-codex" || use.ReasoningEffort != "high" { + t.Fatalf("metadata/response = %+v", use) + } + if use.RecordType != "response_item.message.user_shell_command" { + t.Fatalf("record type = %q", use.RecordType) + } +} + +func TestExtractCodexToolUses_ParsesUserShellCommandWithEmptyOutput(t *testing.T) { + stream := `{"timestamp":"2026-07-13T06:14:01Z","type":"event_msg","payload":{"type":"user_message","message":"\n\nopen file.pdf\n\n\nExit code: 0\nDuration: 0.25 seconds\nOutput:\n\n"}}` + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + if err != nil { + t.Fatalf("ExtractCodexToolUsesFromReader: %v", err) + } + if len(uses) != 1 || uses[0].Tool != "UserShellCommand" { + t.Fatalf("uses = %+v, want one UserShellCommand", uses) + } + if uses[0].Input["output"] != "" || uses[0].Input["status"] != "success" || uses[0].RecordType != "event_msg.user_shell_command" { + t.Fatalf("use = %+v", uses[0]) + } +} + +func TestParseCodexUserShellCommandRejectsPartialWrapper(t *testing.T) { + for _, text := range []string{ + "please parse if you see one", + "echo hi", + "prefix echo hiExit code: 0\nDuration: 1 seconds\nOutput:\n", + } { + if _, ok := parseCodexUserShellCommand(text); ok { + t.Fatalf("parseCodexUserShellCommand(%q) unexpectedly succeeded", text) + } + } +} + func TestExtractCodexToolUses_MapsCodexControlCalls(t *testing.T) { stream := strings.Join([]string{ `{"timestamp":"2026-06-01T10:00:00Z","type":"session_meta","payload":{"id":"sess-1","cwd":"/p"}}`, @@ -115,6 +254,130 @@ func TestExtractCodexToolUses_MapsCodexControlCalls(t *testing.T) { } } +func TestExtractCodexToolUses_RolloutChatAndEvents(t *testing.T) { + stream := strings.Join([]string{ + `{"timestamp":"2026-07-08T11:19:57.028Z","type":"session_meta","payload":{"id":"sess-rollout","cwd":"/repo","cli_version":"0.143.0","model_provider":"openai"}}`, + `{"timestamp":"2026-07-08T11:19:57.028Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1","started_at":"2026-07-08T11:19:57.028Z","model_context_window":258400}}`, + `{"timestamp":"2026-07-08T11:19:58.000Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"\n /repo\n"}]}}`, + `{"timestamp":"2026-07-08T11:19:58.758Z","type":"turn_context","payload":{"model":"gpt-5.5","effort":"high"}}`, + // Twins are skewed, as Codex actually writes them: the user pair arrives + // response_item-first and the assistant pair event_msg-first. An earlier + // fixture gave both twins the same timestamp -- a shape Codex never + // emits -- which is exactly what hid the duplicate-row bug. + `{"timestamp":"2026-07-08T11:19:58.760Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}}`, + `{"timestamp":"2026-07-08T11:19:58.761Z","type":"event_msg","payload":{"type":"user_message","message":"hi"}}`, + `{"timestamp":"2026-07-08T11:20:00.403Z","type":"event_msg","payload":{"type":"agent_message","message":"Hi. What do you want to work on in ` + "`captain`" + `?"}}`, + `{"timestamp":"2026-07-08T11:20:00.404Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Hi. What do you want to work on in ` + "`captain`" + `?"}]}}`, + `{"timestamp":"2026-07-08T11:20:00.432Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":5,"output_tokens":2,"total_tokens":12},"model_context_window":258400}}}`, + `{"timestamp":"2026-07-08T11:20:00.435Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-1","duration_ms":3519,"time_to_first_token_ms":3123}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + if err != nil { + t.Fatalf("ExtractCodexToolUsesFromReader: %v", err) + } + + var got []string + var texts []string + for _, use := range uses { + got = append(got, use.Tool) + if text, _ := use.Input["text"].(string); text != "" { + texts = append(texts, text) + } + } + want := []string{"TaskStarted", "User", "Assistant", "TokenCount", "TaskComplete"} + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Fatalf("tools = %v, want %v; uses=%+v", got, want, uses) + } + if len(texts) != 2 || texts[0] != "hi" || texts[1] != "Hi. What do you want to work on in `captain`?" { + t.Fatalf("texts = %v", texts) + } + for _, use := range uses { + if use.Model != "" && use.Model != "gpt-5.5" { + t.Fatalf("unexpected model on %s: %q", use.Tool, use.Model) + } + } +} + +func TestExtractCodexToolUses_SplitsTaggedAssistantMessageAndDedupesSchemas(t *testing.T) { + message := ` +# Parser plan + +- Split wrappers + + +Source used: https://example.com/change + + + +MEMORY.md:10-12|note=[parser seam] + + +019f3754-ecfa-7323-a76b-a0205ea30bbe + +` + quoted := strings.ReplaceAll(strings.ReplaceAll(message, `\`, `\\`), `"`, `\"`) + quoted = strings.ReplaceAll(quoted, "\n", `\n`) + stream := strings.Join([]string{ + `{"timestamp":"2026-07-10T10:00:00Z","type":"session_meta","payload":{"id":"sess-plan","cwd":"/repo"}}`, + // 3ms skew, as Codex writes twins. The segments land 3 rows apart but + // only 1 source record apart, which is why the merge keys on the + // preceding record rather than a row-distance window. + `{"timestamp":"2026-07-10T10:00:01.000Z","type":"event_msg","payload":{"type":"agent_message","message":"` + quoted + `"}}`, + `{"timestamp":"2026-07-10T10:00:01.003Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"` + quoted + `"}]}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + if err != nil { + t.Fatalf("ExtractCodexToolUsesFromReader: %v", err) + } + if len(uses) != 3 { + t.Fatalf("uses = %+v, want Plan, Assistant, MemoryCitation", uses) + } + if uses[0].Tool != "Plan" || uses[0].Input["content"] != "# Parser plan\n\n- Split wrappers" { + t.Fatalf("plan = %+v", uses[0]) + } + if uses[1].Tool != "Assistant" || uses[1].Input["text"] != "Source used: https://example.com/change" { + t.Fatalf("assistant = %+v", uses[1]) + } + if uses[2].Tool != "MemoryCitation" || uses[2].Input["event"] != "memory_citation" { + t.Fatalf("citation = %+v", uses[2]) + } + if got, ok := uses[2].Input["rollout_ids"].([]string); !ok || len(got) != 1 || got[0] != "019f3754-ecfa-7323-a76b-a0205ea30bbe" { + t.Fatalf("rollout ids = %#v", uses[2].Input["rollout_ids"]) + } +} + +func TestExtractCodexToolUses_ClassifiesAgentsInstructionsAsSystem(t *testing.T) { + stream := strings.Join([]string{ + `{"timestamp":"2026-07-10T09:49:37.000Z","type":"session_meta","payload":{"id":"sess-system","cwd":"/repo"}}`, + `{"timestamp":"2026-07-10T09:49:37.100Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":" # AGENTS.md instructions for /repo\n\nAlways test."}]}}`, + `{"timestamp":"2026-07-10T09:49:37.101Z","type":"event_msg","payload":{"type":"user_message","message":"# AGENTS.md instructions for /repo\n\nAlways test."}}`, + `{"timestamp":"2026-07-10T09:49:37.200Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"Fix the parser"}]}}`, + `{"timestamp":"2026-07-10T09:49:37.201Z","type":"event_msg","payload":{"type":"user_message","message":"Fix the parser"}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + if err != nil { + t.Fatalf("ExtractCodexToolUsesFromReader: %v", err) + } + if len(uses) != 2 { + t.Fatalf("uses = %+v, want one System and one User", uses) + } + if uses[0].Tool != "System" || uses[1].Tool != "User" { + t.Fatalf("tools = %q, %q, want System, User", uses[0].Tool, uses[1].Tool) + } +} + +func TestCodexUserMessageTool_ClassifiesRecommendedPluginsAgentsEnvelopeAsSystem(t *testing.T) { + text := "system recommendations" + + "# AGENTS.md instructions for /repo\nAlways test." + tool, ok := codexUserMessageTool(text) + if !ok || tool != "System" { + t.Fatalf("codexUserMessageTool = %q, %v, want System, true", tool, ok) + } +} + func TestReadCodexSessionInfo_ModelAndEffort(t *testing.T) { stream := strings.Join([]string{ `{"timestamp":"2026-05-07T18:44:49.553Z","type":"session_meta","payload":{"id":"sess-1","cwd":"/p","cli_version":"0.128","model_provider":"openai","originator":"codex_exec","git":{"branch":"main","commit_hash":"abc"}}}`, @@ -165,6 +428,26 @@ func TestReadCodexSessionInfo_LiveThreadStarted(t *testing.T) { } } +func TestCodexSessionMetaAcceptsObjectSource(t *testing.T) { + stream := strings.Join([]string{ + `{"timestamp":"2026-06-19T07:05:51.061Z","type":"session_meta","payload":{"id":"019edeb3-449e-7af3-b300-7123f10944b2","cwd":"/repo","source":{"subagent":{"other":"guardian"}},"model_provider":"openai"}}`, + `{"timestamp":"2026-06-19T07:05:52.061Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"check the change"}]}}`, + }, "\n") + + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(stream)) + require.NoError(t, err) + require.Len(t, uses, 1) + assert.Equal(t, "019edeb3-449e-7af3-b300-7123f10944b2", uses[0].SessionID) + + dir := t.TempDir() + path := filepath.Join(dir, "rollout.jsonl") + require.NoError(t, os.WriteFile(path, []byte(stream), 0o600)) + info, err := ReadCodexSessionInfo(path) + require.NoError(t, err) + require.NotNil(t, info) + assert.Equal(t, "019edeb3-449e-7af3-b300-7123f10944b2", info.ID) +} + func TestReadCodexSessionMetaStopsAtHeader(t *testing.T) { stream := strings.Join([]string{ `{"timestamp":"2026-05-07T18:44:49.553Z","type":"session_meta","payload":{"id":"sess-1","cwd":"/p","cli_version":"0.128","model_provider":"openai"}}`, @@ -189,37 +472,3 @@ func TestReadCodexSessionMetaStopsAtHeader(t *testing.T) { t.Errorf("meta reader should not scan turn_context, got model/effort %q/%q", info.Model, info.ReasoningEffort) } } - -func TestUnwrapCodexErrorMessage(t *testing.T) { - tests := []struct { - name string - in string - want string - }{ - {"plain", "boom", "boom"}, - {"empty", "", ""}, - { - "nested error.message", - `{"type":"error","status":400,"error":{"type":"invalid_request_error","message":"bad model"}}`, - "bad model", - }, - { - "top-level message", - `{"message":"top"}`, - "top", - }, - { - "non-json passthrough", - "not { json", - "not { json", - }, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - got := unwrapCodexErrorMessage(tc.in) - if got != tc.want { - t.Errorf("unwrapCodexErrorMessage(%q) = %q, want %q", tc.in, got, tc.want) - } - }) - } -} diff --git a/pkg/ai/history/codex_reasoning.go b/pkg/ai/history/codex_reasoning.go new file mode 100644 index 0000000..ca987d5 --- /dev/null +++ b/pkg/ai/history/codex_reasoning.go @@ -0,0 +1,130 @@ +package history + +import ( + "fmt" + "time" + + "github.com/segmentio/encoding/json" +) + +// codexReasoningSummaryText returns the plaintext reasoning summary, if any. +// Codex ships the summary as an array of blocks; the last non-empty block wins. +func codexReasoningSummaryText(raw json.RawMessage) string { + if len(raw) == 0 { + return "" + } + var summaries []CodexReasoningSummary + if json.Unmarshal(raw, &summaries) != nil { + return "" + } + var text string + for _, summary := range summaries { + if summary.Text != "" { + text = summary.Text + } + } + return text +} + +// codexReasoningCollapser folds contentless reasoning records into one row per +// turn. +// +// Modern Codex emits reasoning as `summary:[]` plus `encrypted_content`: the +// text is unrecoverable, but the record still proves the session was alive at +// that instant. Dropping it understates the session's end time (and therefore +// last_activity_at) and inflates apparent idle gaps. Emitting one row per +// record would instead bury the transcript, so a turn's records collapse into a +// single row stamped with the LAST timestamp and carrying the span. +// +// Records that DO carry a plaintext summary keep their own row and never flush +// the span: plaintext and contentless reasoning mix within one turn. +type codexReasoningCollapser struct { + turnID string + cwd string + sessionID string + first time.Time + last time.Time + count int +} + +// observe consumes a response_item.reasoning event, returning any rows it +// completes. A turn change flushes the pending span before accumulating. +func (c *codexReasoningCollapser) observe(event CodexEvent, currentTurn, cwd, sessionID string) []ToolUse { + if text := codexReasoningSummaryText(event.Payload.Summary); text != "" { + return []ToolUse{{ + Tool: "Reasoning", + Input: map[string]any{"text": text}, + Timestamp: event.Time(), + CWD: cwd, + SessionID: sessionID, + TurnID: codexEventTurnID(event), + Source: "codex", + RecordType: "response_item.reasoning", + }} + } + + ts := event.Time() + if ts == nil { + // No parseable timestamp is no evidence of when the session was alive, + // which is the only thing a contentless record carries. + return nil + } + + var out []ToolUse + turn := firstNonEmpty(codexEventTurnID(event), currentTurn) + if c.count > 0 && c.turnID != turn { + out = c.flush() + } + if c.count == 0 { + c.turnID, c.cwd, c.sessionID = turn, cwd, sessionID + c.first, c.last = *ts, *ts + } + if ts.Before(c.first) { + c.first = *ts + } + if ts.After(c.last) { + c.last = *ts + } + c.count++ + return out +} + +// flush emits the pending span and resets. It returns nil when nothing is +// pending, so callers can flush unconditionally. +func (c *codexReasoningCollapser) flush() []ToolUse { + if c.count == 0 { + return nil + } + last := c.last + use := ToolUse{ + Tool: "Reasoning", + Input: map[string]any{ + "text": codexReasoningSpanText(c.first, c.last, c.count), + "first_at": c.first.UTC().Format(time.RFC3339Nano), + "last_at": c.last.UTC().Format(time.RFC3339Nano), + "count": c.count, + }, + Timestamp: &last, + CWD: c.cwd, + SessionID: c.sessionID, + TurnID: c.turnID, + Source: "codex", + RecordType: "response_item.reasoning", + } + *c = codexReasoningCollapser{} + return []ToolUse{use} +} + +// codexReasoningSpanText renders the span. It must stay deterministic: the +// result feeds codexChatDedupeKey. +func codexReasoningSpanText(first, last time.Time, count int) string { + if count == 1 { + return fmt.Sprintf("1 encrypted reasoning record at %s", first.UTC().Format(time.RFC3339)) + } + return fmt.Sprintf("%d encrypted reasoning records over %s (%s → %s)", + count, + last.Sub(first).Truncate(time.Second), + first.UTC().Format(time.RFC3339), + last.UTC().Format(time.RFC3339), + ) +} diff --git a/pkg/ai/history/codex_reasoning_ginkgo_test.go b/pkg/ai/history/codex_reasoning_ginkgo_test.go new file mode 100644 index 0000000..193e9da --- /dev/null +++ b/pkg/ai/history/codex_reasoning_ginkgo_test.go @@ -0,0 +1,139 @@ +package history + +import ( + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +const nanoLayout = time.RFC3339Nano + +// reasoningUses returns only the Reasoning rows produced from a rollout stream. +func reasoningUses(lines ...string) []ToolUse { + uses, err := ExtractCodexToolUsesFromReader(strings.NewReader(strings.Join(lines, "\n"))) + Expect(err).ToNot(HaveOccurred()) + var out []ToolUse + for _, use := range uses { + if use.Tool == "Reasoning" { + out = append(out, use) + } + } + return out +} + +const ( + reasoningSessionMeta = `{"timestamp":"2026-07-16T11:14:45.000Z","type":"session_meta","payload":{"id":"sess-enc","cwd":"/repo","cli_version":"0.160.0"}}` + reasoningTurnContext = `{"timestamp":"2026-07-16T11:14:55.426Z","type":"turn_context","payload":{"turn_id":"turn-1","model":"gpt-5.5","effort":"high"}}` +) + +var _ = Describe("codex reasoning collapse", func() { + // Modern Codex ships reasoning as summary:[] + encrypted_content. The text is + // unrecoverable, but the record still proves the session was alive at that + // instant, so its timestamp must reach extendRange via a collapsed row. + It("collapses contentless reasoning in a turn onto the last timestamp", func() { + uses := reasoningUses( + reasoningSessionMeta, + reasoningTurnContext, + `{"timestamp":"2026-07-16T11:15:04.024Z","type":"response_item","payload":{"type":"reasoning","id":"rs_1","summary":[],"encrypted_content":"gAAAA1","internal_chat_message_metadata_passthrough":{"turn_id":"turn-1"}}}`, + `{"timestamp":"2026-07-16T11:20:00.000Z","type":"response_item","payload":{"type":"function_call","name":"shell","arguments":"{\"command\":\"ls\"}","call_id":"call-1"}}`, + `{"timestamp":"2026-07-16T11:20:01.000Z","type":"response_item","payload":{"type":"function_call_output","call_id":"call-1","output":"ok"}}`, + `{"timestamp":"2026-07-16T11:32:58.744Z","type":"response_item","payload":{"type":"reasoning","id":"rs_2","summary":[],"encrypted_content":"gAAAA2","internal_chat_message_metadata_passthrough":{"turn_id":"turn-1"}}}`, + ) + + Expect(uses).To(HaveLen(1)) + use := uses[0] + Expect(use.Timestamp.UTC().Format(nanoLayout)).To(Equal("2026-07-16T11:32:58.744Z")) + Expect(use.Input["first_at"]).To(Equal("2026-07-16T11:15:04.024Z")) + Expect(use.Input["last_at"]).To(Equal("2026-07-16T11:32:58.744Z")) + Expect(use.Input["count"]).To(Equal(2)) + Expect(use.Input["text"]).To(ContainSubstring("2 encrypted reasoning records")) + Expect(use.TurnID).To(Equal("turn-1")) + Expect(use.Model).To(Equal("gpt-5.5")) + Expect(use.ReasoningEffort).To(Equal("high")) + Expect(use.SessionID).To(Equal("sess-enc")) + Expect(use.CWD).To(Equal("/repo")) + Expect(use.Source).To(Equal("codex")) + Expect(use.RecordType).To(Equal("response_item.reasoning")) + }) + + // Older rollouts carry real plaintext summaries; those must keep rendering + // one row each rather than being folded into a span. + It("keeps one row per plaintext summary", func() { + uses := reasoningUses( + reasoningSessionMeta, + reasoningTurnContext, + `{"timestamp":"2026-07-16T11:15:04.024Z","type":"response_item","payload":{"type":"reasoning","summary":[{"type":"summary_text","text":"Checking the parser"}]}}`, + ) + + Expect(uses).To(HaveLen(1)) + Expect(uses[0].Input["text"]).To(Equal("Checking the parser")) + Expect(uses[0].Timestamp.UTC().Format(nanoLayout)).To(Equal("2026-07-16T11:15:04.024Z")) + Expect(uses[0].Input).ToNot(HaveKey("first_at")) + Expect(uses[0].Input).ToNot(HaveKey("last_at")) + Expect(uses[0].Input).ToNot(HaveKey("count")) + }) + + // Regression pin: 47% of reasoning records carry no turn_id, and old-format + // turn_context carries none either. Keying the collapse on turn_id alone + // would fold a whole multi-turn session into one bogus span, so the flush + // must trigger on the turn_context event itself. + It("segments spans on turn_context when the rollout carries no turn ids", func() { + uses := reasoningUses( + `{"timestamp":"2026-01-28T15:32:20.000Z","type":"session_meta","payload":{"id":"sess-old","cwd":"/repo"}}`, + `{"timestamp":"2026-01-28T15:32:24.056Z","type":"turn_context","payload":{"cwd":"/repo","approval_policy":"on-request"}}`, + `{"timestamp":"2026-01-28T15:32:27.895Z","type":"response_item","payload":{"type":"reasoning","summary":[],"encrypted_content":"e1"}}`, + `{"timestamp":"2026-01-28T15:32:31.990Z","type":"turn_context","payload":{"cwd":"/repo","approval_policy":"on-request"}}`, + `{"timestamp":"2026-01-28T15:32:38.675Z","type":"response_item","payload":{"type":"reasoning","summary":[],"encrypted_content":"e2"}}`, + ) + + Expect(uses).To(HaveLen(2)) + Expect(uses[0].Timestamp.UTC().Format(nanoLayout)).To(Equal("2026-01-28T15:32:27.895Z")) + Expect(uses[0].Input["count"]).To(Equal(1)) + Expect(uses[0].Input["text"]).To(ContainSubstring("1 encrypted reasoning record ")) + Expect(uses[1].Timestamp.UTC().Format(nanoLayout)).To(Equal("2026-01-28T15:32:38.675Z")) + Expect(uses[1].Input["count"]).To(Equal(1)) + }) + + // Plaintext and contentless reasoning genuinely mix inside one file, so the + // collapse is per-turn rather than per-run: a plaintext row must not flush. + It("keeps plaintext rows and adds one collapsed row for a mixed turn", func() { + uses := reasoningUses( + reasoningSessionMeta, + reasoningTurnContext, + `{"timestamp":"2026-07-16T11:15:01.000Z","type":"response_item","payload":{"type":"reasoning","summary":[],"encrypted_content":"e1"}}`, + `{"timestamp":"2026-07-16T11:15:02.000Z","type":"response_item","payload":{"type":"reasoning","summary":[{"type":"summary_text","text":"first plain"}]}}`, + `{"timestamp":"2026-07-16T11:15:03.000Z","type":"response_item","payload":{"type":"reasoning","summary":[],"encrypted_content":"e2"}}`, + `{"timestamp":"2026-07-16T11:15:04.000Z","type":"response_item","payload":{"type":"reasoning","summary":[{"type":"summary_text","text":"second plain"}]}}`, + `{"timestamp":"2026-07-16T11:15:05.000Z","type":"response_item","payload":{"type":"reasoning","summary":[],"encrypted_content":"e3"}}`, + ) + + Expect(uses).To(HaveLen(3)) + Expect(uses[0].Input["text"]).To(Equal("first plain")) + Expect(uses[1].Input["text"]).To(Equal("second plain")) + Expect(uses[2].Input["count"]).To(Equal(3)) + Expect(uses[2].Input["first_at"]).To(Equal("2026-07-16T11:15:01Z")) + Expect(uses[2].Input["last_at"]).To(Equal("2026-07-16T11:15:05Z")) + }) + + It("flushes a pending span at EOF", func() { + uses := reasoningUses( + reasoningSessionMeta, + reasoningTurnContext, + `{"timestamp":"2026-07-16T11:15:04.024Z","type":"response_item","payload":{"type":"reasoning","summary":[],"encrypted_content":"e1"}}`, + ) + + Expect(uses).To(HaveLen(1)) + Expect(uses[0].Input["count"]).To(Equal(1)) + Expect(uses[0].Timestamp.UTC().Format(nanoLayout)).To(Equal("2026-07-16T11:15:04.024Z")) + }) + + It("drops reasoning records with no parseable timestamp", func() { + Expect(reasoningUses( + reasoningSessionMeta, + reasoningTurnContext, + `{"type":"response_item","payload":{"type":"reasoning","summary":[],"encrypted_content":"e1"}}`, + )).To(BeEmpty()) + }) +}) diff --git a/pkg/ai/history/codex_types.go b/pkg/ai/history/codex_types.go index 381adc4..3c91eb5 100644 --- a/pkg/ai/history/codex_types.go +++ b/pkg/ai/history/codex_types.go @@ -51,24 +51,34 @@ func (e CodexEvent) Time() *time.Time { } type CodexPayload struct { - Type string `json:"type"` + Type string `json:"type"` + Raw map[string]any `json:"-"` // session_meta - ID string `json:"id,omitempty"` - CWD string `json:"cwd,omitempty"` - CLIVersion string `json:"cli_version,omitempty"` - Source string `json:"source,omitempty"` - ModelProvider string `json:"model_provider,omitempty"` - Originator string `json:"originator,omitempty"` - Git *CodexGitMeta `json:"git,omitempty"` + ID string `json:"id,omitempty"` + CWD string `json:"cwd,omitempty"` + CLIVersion string `json:"cli_version,omitempty"` + // Source has appeared as both a scalar (for root sessions) and an object + // (for example {"subagent":{"other":"guardian"}}). Keep the provider + // shape opaque so a new source variant cannot invalidate the whole + // session_meta record and discard its ID. + Source json.RawMessage `json:"source,omitempty"` + ModelProvider string `json:"model_provider,omitempty"` + Originator string `json:"originator,omitempty"` + Git *CodexGitMeta `json:"git,omitempty"` // response_item: function_call - Name string `json:"name,omitempty"` - Arguments string `json:"arguments,omitempty"` - CallID string `json:"call_id,omitempty"` - - // response_item: function_call_output - Output string `json:"output,omitempty"` + Name string `json:"name,omitempty"` + Namespace string `json:"namespace,omitempty"` + Arguments json.RawMessage `json:"arguments,omitempty"` + CallID string `json:"call_id,omitempty"` + Input string `json:"input,omitempty"` + Metadata *CodexMetadata `json:"internal_chat_message_metadata_passthrough,omitempty"` + + // response_item: function_call_output. Codex emits either a JSON string or + // an ordered array of text content blocks depending on the tool transport. + Output json.RawMessage `json:"output,omitempty"` + Tools []CodexToolSearchNamespace `json:"tools,omitempty"` // response_item: reasoning carries an array of CodexReasoningSummary; // turn_context carries a plain string ("none", "auto", etc.). Use a @@ -80,8 +90,16 @@ type CodexPayload struct { Content []CodexContent `json:"content,omitempty"` // event_msg: agent_reasoning / agent_message / user_message - Text string `json:"text,omitempty"` - Message string `json:"message,omitempty"` + Text string `json:"text,omitempty"` + Message string `json:"message,omitempty"` + Phase string `json:"phase,omitempty"` + StartedAt any `json:"started_at,omitempty"` + CompletedAt int64 `json:"completed_at,omitempty"` + DurationMS int64 `json:"duration_ms,omitempty"` + TimeToFirstTokenMS int64 `json:"time_to_first_token_ms,omitempty"` + LastAgentMessage string `json:"last_agent_message,omitempty"` + ModelContextWindow int `json:"model_context_window,omitempty"` + CollaborationModeKind string `json:"collaboration_mode_kind,omitempty"` // event_msg: token_count Info *CodexTokenInfo `json:"info,omitempty"` @@ -96,12 +114,44 @@ type CodexPayload struct { Effort string `json:"effort,omitempty"` } +func (p *CodexPayload) UnmarshalJSON(data []byte) error { + type alias CodexPayload + var a alias + if err := json.Unmarshal(data, &a); err != nil { + return err + } + *p = CodexPayload(a) + var raw map[string]any + if err := json.Unmarshal(data, &raw); err == nil { + p.Raw = raw + } + return nil +} + type CodexGitMeta struct { CommitHash string `json:"commit_hash,omitempty"` Branch string `json:"branch,omitempty"` RepositoryURL string `json:"repository_url,omitempty"` } +type CodexMetadata struct { + TurnID string `json:"turn_id,omitempty"` +} + +type CodexToolSearchNamespace struct { + Type string `json:"type,omitempty"` + Name string `json:"name,omitempty"` + Description string `json:"description,omitempty"` + Tools []CodexToolSearchEntry `json:"tools,omitempty"` +} + +type CodexToolSearchEntry struct { + Type string `json:"type,omitempty"` + Name string `json:"name,omitempty"` + Description string `json:"description,omitempty"` + Parameters map[string]any `json:"parameters,omitempty"` +} + type CodexReasoningSummary struct { Type string `json:"type"` Text string `json:"text"` @@ -119,8 +169,9 @@ type CodexTokenInfo struct { } type CodexTokenUsage struct { - InputTokens int `json:"input_tokens"` - CachedInputTokens int `json:"cached_input_tokens"` - OutputTokens int `json:"output_tokens"` - TotalTokens int `json:"total_tokens"` + InputTokens int `json:"input_tokens"` + CachedInputTokens int `json:"cached_input_tokens"` + OutputTokens int `json:"output_tokens"` + ReasoningOutputTokens int `json:"reasoning_output_tokens"` + TotalTokens int `json:"total_tokens"` } diff --git a/pkg/ai/history/tools.go b/pkg/ai/history/tools.go index bbf02b8..5b8894c 100644 --- a/pkg/ai/history/tools.go +++ b/pkg/ai/history/tools.go @@ -34,16 +34,6 @@ type Option struct { Label string `json:"label"` } -type TodoWrite struct { - Todos []Todo `json:"todos"` -} - -type Todo struct { - ActiveForm string `json:"activeForm,omitempty"` - Content string `json:"content,omitempty"` - Status string `json:"status,omitempty"` -} - // --- Typed tool structs --- // Each delegates to ToolUse.Pretty() for unified rendering. diff --git a/pkg/ai/history/types.go b/pkg/ai/history/types.go index 6784a24..6a19612 100644 --- a/pkg/ai/history/types.go +++ b/pkg/ai/history/types.go @@ -8,11 +8,22 @@ type ToolUse struct { Timestamp *time.Time `json:"timestamp,omitempty"` CWD string `json:"cwd,omitempty"` SessionID string `json:"session_id,omitempty"` + TurnID string `json:"turn_id,omitempty"` ToolUseID string `json:"tool_use_id,omitempty"` Source string `json:"source,omitempty"` // "claude" or "codex" Model string `json:"model,omitempty"` ReasoningEffort string `json:"reasoning_effort,omitempty"` + Namespace string `json:"namespace,omitempty"` + InputTokens int `json:"input_tokens,omitempty"` + OutputTokens int `json:"output_tokens,omitempty"` + CacheReadTokens int `json:"cache_read_tokens,omitempty"` + TotalTokens int `json:"total_tokens,omitempty"` + ContextWindow int `json:"context_window,omitempty"` + AgentID string `json:"agent_id,omitempty"` + AgentType string `json:"agent_type,omitempty"` + AgentDesc string `json:"agent_desc,omitempty"` Response string `json:"response,omitempty"` + RecordType string `json:"-"` } type Filter struct { diff --git a/pkg/ai/internal/gen-model-registry/main.go b/pkg/ai/internal/gen-model-registry/main.go new file mode 100644 index 0000000..5316aeb --- /dev/null +++ b/pkg/ai/internal/gen-model-registry/main.go @@ -0,0 +1,496 @@ +package main + +import ( + "encoding/json" + "flag" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "sort" + "strings" + "time" + + captainai "github.com/flanksource/captain/pkg/ai" +) + +const modelsDevAPIURL = "https://models.dev/api.json" + +type modelsDevProvider struct { + Models map[string]modelsDevModel `json:"models"` +} + +type modelsDevModel struct { + ID string `json:"id"` + Name string `json:"name"` + Family string `json:"family"` + Reasoning bool `json:"reasoning"` + Temperature bool `json:"temperature"` + ReasoningOptions []modelsDevReasoningOption `json:"reasoning_options"` + ReleaseDate string `json:"release_date"` + Modalities modelsDevModalities `json:"modalities"` + Limit modelsDevLimit `json:"limit"` + Cost modelsDevCost `json:"cost"` +} + +// modelsDevCost is a model's published list price in USD per million tokens. +// The upstream block also carries tiered pricing (e.g. OpenAI's >200k-context +// rates) and per-modality audio rates; captain prices a single flat tier, so +// only the base four are read. +type modelsDevCost struct { + Input float64 `json:"input"` + Output float64 `json:"output"` + CacheRead float64 `json:"cache_read"` + CacheWrite float64 `json:"cache_write"` +} + +// modelsDevReasoningOption is one entry of a model's reasoning_options; only its +// type ("effort", "budget_tokens", "toggle", …) is needed to classify how the +// model controls thinking. +type modelsDevReasoningOption struct { + Type string `json:"type"` + Values []string `json:"values"` +} + +type modelsDevModalities struct { + Input []string `json:"input"` + Output []string `json:"output"` +} + +type modelsDevLimit struct { + Context int `json:"context"` +} + +// generatedModel mirrors pkg/ai.registryModel field-for-field (and its JSON +// tags), so the emitted JSON round-trips into the registry on load. +type generatedModel struct { + ID string `json:"id"` + Provider string `json:"provider"` + Family string `json:"family"` + Version string `json:"version"` + Label string `json:"label"` + ReleaseDate string `json:"releaseDate,omitempty"` + Reasoning bool `json:"reasoning,omitempty"` + Temperature bool `json:"temperature,omitempty"` + ContextWindow int `json:"contextWindow,omitempty"` + Cost *cost `json:"cost,omitempty"` + InputMediaTypes []string `json:"inputMediaTypes,omitempty"` + Preferred bool `json:"preferred,omitempty"` + AdaptiveThinking bool `json:"adaptiveThinking,omitempty"` + Availability []string `json:"availability,omitempty"` + SupportedEfforts []string `json:"supportedEfforts,omitempty"` + DefaultEffort string `json:"defaultEffort,omitempty"` + Priority int `json:"priority,omitempty"` + + // Aliases and SupersededBy are patch-only: models.dev has no concept of + // either. They carry the codename and retired-id knowledge that used to be + // hardcoded in pkg/ai, where the spec decoder could not see it. + Aliases []string `json:"aliases,omitempty"` + SupersededBy string `json:"supersededBy,omitempty"` +} + +// cost mirrors pkg/api/registry.ModelCost: USD per million tokens. It is a +// pointer on generatedModel so a model models.dev prices at nothing at all +// (and a patch-only entry that names no cost) omits the block entirely rather +// than claiming a free model. +type cost struct { + Input float64 `json:"input,omitempty"` + Output float64 `json:"output,omitempty"` + CacheRead float64 `json:"cacheRead,omitempty"` + CacheWrite float64 `json:"cacheWrite,omitempty"` +} + +// deriveCost carries models.dev's list price into the catalog so pricing has a +// per-model, per-version source of truth. Captain used to price Claude from a +// hand-written family table (any "opus" id → $15/$75), which silently drifted +// three-fold once Opus 4.5 cut prices; sourcing it here means `task +// models:update` corrects prices along with everything else. +func deriveCost(model modelsDevModel) *cost { + c := model.Cost + if c.Input == 0 && c.Output == 0 { + return nil + } + return &cost{Input: c.Input, Output: c.Output, CacheRead: c.CacheRead, CacheWrite: c.CacheWrite} +} + +func main() { + if err := run(); err != nil { + fmt.Fprintf(os.Stderr, "gen-model-registry: %v\n", err) + os.Exit(1) + } +} + +func run() error { + var output string + var source string + var patchPath string + flag.StringVar(&output, "output", "", "Path to write pkg/api/registry/models.json") + flag.StringVar(&source, "source", modelsDevAPIURL, "models.dev API URL, local file path, or - for stdin") + flag.StringVar(&patchPath, "patches", "", "Path to the per-model JSON merge-patch file") + flag.Parse() + if output == "" { + return fmt.Errorf("--output is required") + } + + data, err := readSource(source) + if err != nil { + return err + } + patches, err := readPatches(patchPath) + if err != nil { + return err + } + models, err := generateModels(data, patches) + if err != nil { + return err + } + src, err := renderJSON(models) + if err != nil { + return err + } + if dir := filepath.Dir(output); dir != "." { + if err := os.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("create output directory: %w", err) + } + } + if err := os.WriteFile(output, src, 0o644); err != nil { + return fmt.Errorf("write %s: %w", output, err) + } + return nil +} + +// readPatches loads the per-model merge-patch file: a JSON object keyed by model +// ID whose values carry only the fields to override on the fetched row (or a +// full row for models not present in the models.dev catalog). +func readPatches(path string) (map[string]json.RawMessage, error) { + if path == "" { + return nil, nil + } + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read patches %s: %w", path, err) + } + var patches map[string]json.RawMessage + if err := json.Unmarshal(data, &patches); err != nil { + return nil, fmt.Errorf("parse patches %s: %w", path, err) + } + return patches, nil +} + +func readSource(source string) ([]byte, error) { + switch { + case source == "-": + data, err := io.ReadAll(os.Stdin) + if err != nil { + return nil, fmt.Errorf("read stdin: %w", err) + } + return data, nil + case strings.HasPrefix(source, "http://") || strings.HasPrefix(source, "https://"): + return readURL(source) + default: + data, err := os.ReadFile(source) + if err != nil { + return nil, fmt.Errorf("read %s: %w", source, err) + } + return data, nil + } +} + +func readURL(source string) ([]byte, error) { + client := http.Client{Timeout: 30 * time.Second} + req, err := http.NewRequest(http.MethodGet, source, nil) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + req.Header.Set("User-Agent", "captain/gen-model-registry") + + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("fetch %s: %w", source, err) + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("fetch %s: unexpected status %s", source, resp.Status) + } + data, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read %s: %w", source, err) + } + return data, nil +} + +func generateModels(data []byte, patches map[string]json.RawMessage) ([]generatedModel, error) { + var providers map[string]modelsDevProvider + if err := json.Unmarshal(data, &providers); err != nil { + return nil, fmt.Errorf("parse models.dev catalog: %w", err) + } + out := map[string]generatedModel{} + for _, mapping := range []struct { + source string + provider string + }{ + {"anthropic", "anthropic"}, + {"openai", "openai"}, + {"google", "google"}, + {"deepseek", "deepseek"}, + } { + provider, ok := providers[mapping.source] + if !ok { + return nil, fmt.Errorf("models.dev catalog is missing provider %q", mapping.source) + } + for id, model := range provider.Models { + if model.ID != "" { + id = model.ID + } + _, explicitlyPatched := patches[id] + row, ok := generatedModelFromModelsDev(mapping.provider, id, model, explicitlyPatched) + if ok { + out[row.ID] = row + } + } + } + if err := applyPatches(out, patches); err != nil { + return nil, err + } + + models := make([]generatedModel, 0, len(out)) + for _, row := range out { + models = append(models, row) + } + sort.SliceStable(models, func(i, j int) bool { + if models[i].Provider == models[j].Provider { + if models[i].ReleaseDate == models[j].ReleaseDate { + return models[i].ID < models[j].ID + } + return models[i].ReleaseDate > models[j].ReleaseDate + } + return providerRank(models[i].Provider) < providerRank(models[j].Provider) + }) + return models, nil +} + +func generatedModelFromModelsDev(provider, id string, model modelsDevModel, explicitlyPatched bool) (generatedModel, bool) { + id = strings.TrimSpace(id) + if id == "" || !supportsTextIO(model) || (!explicitlyPatched && hiddenModel(provider, id)) { + return generatedModel{}, false + } + identity, ok := captainai.ParseModelIdentity(provider, id) + if !ok || identity.Provider != provider { + return generatedModel{}, false + } + if !supportedFamily(identity) { + return generatedModel{}, false + } + label := strings.TrimSpace(model.Name) + if label == "" { + label = titleModelID(id) + } + // Preferred is opt-in: models default to non-preferred and are surfaced in + // menus only when patches.json explicitly sets "preferred": true. + return generatedModel{ + ID: id, + Provider: provider, + Family: identity.Family, + Version: identity.Version, + Label: label, + ReleaseDate: model.ReleaseDate, + Reasoning: model.Reasoning, + Temperature: model.Temperature, + ContextWindow: model.Limit.Context, + Cost: deriveCost(model), + InputMediaTypes: deriveInputMediaTypes(model.Modalities.Input), + AdaptiveThinking: deriveAdaptiveThinking(provider, model), + SupportedEfforts: deriveSupportedEfforts(provider, model), + }, true +} + +func deriveInputMediaTypes(modalities []string) []string { + out := make([]string, 0, len(modalities)) + for _, modality := range modalities { + switch strings.ToLower(strings.TrimSpace(modality)) { + case "image": + out = append(out, "image/*") + case "audio": + out = append(out, "audio/*") + case "video": + out = append(out, "video/*") + case "pdf": + out = append(out, "application/pdf") + } + } + return out +} + +func deriveSupportedEfforts(provider string, model modelsDevModel) []string { + // Only retain source levels that Captain can submit through the selected + // backend. DeepSeek selects reasoning by model ID, while Anthropic's legacy + // budget-token models do not accept an effort value. Gemini's "minimal" + // level is intentionally dropped by the shared Effort validation below until + // Captain has a faithful token-budget mapping for it. + if provider == "deepseek" || (provider == "anthropic" && !deriveAdaptiveThinking(provider, model)) { + return nil + } + for _, opt := range model.ReasoningOptions { + if opt.Type != "effort" { + continue + } + out := make([]string, 0, len(opt.Values)) + seen := map[string]bool{} + for _, value := range opt.Values { + value = strings.ToLower(strings.TrimSpace(value)) + if value == "" || value == "none" || seen[value] || captainai.ValidateEffort(captainai.Effort(value)) != nil { + continue + } + seen[value] = true + out = append(out, value) + } + return out + } + return nil +} + +// deriveAdaptiveThinking reports whether an Anthropic model uses the adaptive +// thinking schema (thinking:{type:adaptive} + output_config.effort) rather than +// the legacy enabled schema. models.dev encodes this in reasoning_options: an +// effort control with no budget_tokens control means adaptive. It is +// Anthropic-only; other providers carry their own effort mechanisms. +func deriveAdaptiveThinking(provider string, model modelsDevModel) bool { + if provider != "anthropic" || !model.Reasoning { + return false + } + hasEffort, hasBudget := false, false + for _, opt := range model.ReasoningOptions { + switch opt.Type { + case "effort": + hasEffort = true + case "budget_tokens": + hasBudget = true + } + } + return hasEffort && !hasBudget +} + +// supportsTextIO keeps only models Captain can actually prompt: it must accept a +// text prompt and return text. Audio-only Live API entries (e.g. +// gemini-3.5-live-translate-preview) declare no text input and would otherwise +// land in the catalog as models no prompt can be sent to. +func supportsTextIO(model modelsDevModel) bool { + return declaresText(model.Modalities.Input) && declaresText(model.Modalities.Output) +} + +// declaresText reports whether a modality list names text. An empty list means +// models.dev did not record modalities, which predates the field and is treated +// as plain text. +func declaresText(modalities []string) bool { + if len(modalities) == 0 { + return true + } + for _, modality := range modalities { + if strings.EqualFold(strings.TrimSpace(modality), "text") { + return true + } + } + return false +} + +func hiddenModel(provider, id string) bool { + switch provider { + case "openai": + return captainai.IsLegacyModelIDForBackend(id, captainai.BackendOpenAI) + case "anthropic": + return captainai.IsLegacyModelIDForBackend(id, captainai.BackendAnthropic) + case "google": + lower := strings.ToLower(id) + return strings.Contains(lower, "embedding") || strings.Contains(lower, "image") || strings.Contains(lower, "tts") + default: + return false + } +} + +func supportedFamily(identity captainai.ModelIdentity) bool { + switch identity.Provider { + case "anthropic": + switch identity.Family { + case "opus", "sonnet", "haiku", "fable": + return true + } + case "openai": + return identity.Family == "gpt" + case "google": + return identity.Family == "gemini" + case "deepseek": + return identity.Family == "deepseek" + } + return false +} + +// applyPatches overlays the per-model merge patches onto the fetched catalog. +// Each patch overwrites only the fields it names (RFC 7386 merge semantics); a +// null patch value drops the model from the catalog entirely, and a patch whose +// ID is absent from the catalog is treated as a new full entry that must name +// provider, family, and label. +func applyPatches(out map[string]generatedModel, patches map[string]json.RawMessage) error { + ids := make([]string, 0, len(patches)) + for id := range patches { + ids = append(ids, id) + } + sort.Strings(ids) + for _, id := range ids { + raw := patches[id] + if strings.TrimSpace(string(raw)) == "null" { + delete(out, id) + continue + } + m, existed := out[id] + if m.ID == "" { + m.ID = id + } + if err := json.Unmarshal(raw, &m); err != nil { + return fmt.Errorf("patch %q: %w", id, err) + } + if !existed && (m.Provider == "" || m.Family == "" || m.Label == "") { + return fmt.Errorf("patch %q adds a new model but is missing provider, family, or label", id) + } + out[id] = m + } + return nil +} + +func renderJSON(models []generatedModel) ([]byte, error) { + data, err := json.MarshalIndent(models, "", " ") + if err != nil { + return nil, err + } + return append(data, '\n'), nil +} + +func providerRank(provider string) int { + switch provider { + case "anthropic": + return 0 + case "openai": + return 1 + case "google": + return 2 + case "deepseek": + return 3 + default: + return 100 + } +} + +func titleModelID(id string) string { + parts := strings.Split(strings.ReplaceAll(id, "_", "-"), "-") + for i, part := range parts { + switch strings.ToLower(part) { + case "gpt": + parts[i] = "GPT" + default: + if part != "" { + parts[i] = strings.ToUpper(part[:1]) + part[1:] + } + } + } + return strings.Join(parts, " ") +} diff --git a/pkg/ai/internal/gen-model-registry/main_test.go b/pkg/ai/internal/gen-model-registry/main_test.go new file mode 100644 index 0000000..6c3b2a8 --- /dev/null +++ b/pkg/ai/internal/gen-model-registry/main_test.go @@ -0,0 +1,142 @@ +package main + +import ( + "encoding/json" + "testing" +) + +func TestApplyPatchesPartialOverlay(t *testing.T) { + out := map[string]generatedModel{ + "claude-sonnet-5": { + ID: "claude-sonnet-5", Provider: "anthropic", Family: "sonnet", + Version: "5", Label: "Claude Sonnet 5", ReleaseDate: "2026-06-29", + Reasoning: true, ContextWindow: 1000000, Preferred: true, + }, + } + patches := map[string]json.RawMessage{ + "claude-sonnet-5": json.RawMessage(`{"reasoning": false, "adaptiveThinking": true}`), + } + if err := applyPatches(out, patches); err != nil { + t.Fatalf("applyPatches: %v", err) + } + got := out["claude-sonnet-5"] + if got.Reasoning != false { + t.Errorf("Reasoning = %v, want false (patched)", got.Reasoning) + } + if !got.AdaptiveThinking { + t.Errorf("AdaptiveThinking = %v, want true (patched)", got.AdaptiveThinking) + } + // Unspecified fields keep their fetched values. + if got.Preferred != true || got.ContextWindow != 1000000 || got.Label != "Claude Sonnet 5" { + t.Errorf("unpatched fields were altered: %+v", got) + } +} + +func TestApplyPatchesAddsNewEntry(t *testing.T) { + out := map[string]generatedModel{} + patches := map[string]json.RawMessage{ + "gpt-5.5": json.RawMessage(`{"provider": "openai", "family": "gpt", "version": "5.5", + "label": "GPT-5.5", "reasoning": true, "contextWindow": 1050000, "preferred": true}`), + } + if err := applyPatches(out, patches); err != nil { + t.Fatalf("applyPatches: %v", err) + } + got, ok := out["gpt-5.5"] + if !ok { + t.Fatal("gpt-5.5 was not added") + } + if got.ID != "gpt-5.5" { + t.Errorf("ID = %q, want gpt-5.5 (derived from patch key)", got.ID) + } + if got.Provider != "openai" || got.Family != "gpt" || !got.Reasoning { + t.Errorf("new entry has wrong fields: %+v", got) + } +} + +func TestApplyPatchesNullRemovesEntry(t *testing.T) { + out := map[string]generatedModel{ + "deepseek-chat": {ID: "deepseek-chat", Provider: "deepseek", Family: "deepseek"}, + "deepseek-v4-pro": {ID: "deepseek-v4-pro", Provider: "deepseek", Family: "deepseek"}, + "deepseek-reasoner": {ID: "deepseek-reasoner", Provider: "deepseek", Family: "deepseek"}, + } + patches := map[string]json.RawMessage{ + "deepseek-chat": json.RawMessage(`null`), + "deepseek-reasoner": json.RawMessage(`null`), + } + if err := applyPatches(out, patches); err != nil { + t.Fatalf("applyPatches: %v", err) + } + if _, ok := out["deepseek-chat"]; ok { + t.Error("deepseek-chat should have been removed by null patch") + } + if _, ok := out["deepseek-reasoner"]; ok { + t.Error("deepseek-reasoner should have been removed by null patch") + } + if _, ok := out["deepseek-v4-pro"]; !ok { + t.Error("deepseek-v4-pro should have been retained") + } +} + +func TestApplyPatchesNewEntryMissingFieldsErrors(t *testing.T) { + out := map[string]generatedModel{} + patches := map[string]json.RawMessage{ + "mystery-model": json.RawMessage(`{"reasoning": true}`), + } + if err := applyPatches(out, patches); err == nil { + t.Fatal("expected error for new entry missing provider/family/label, got nil") + } +} + +func TestSupportsTextIORejectsNonPromptableModalities(t *testing.T) { + for _, tc := range []struct { + name string + input []string + output []string + want bool + }{ + {"text only", []string{"text"}, []string{"text"}, true}, + {"multimodal input", []string{"text", "image", "audio"}, []string{"text"}, true}, + {"unrecorded modalities", nil, nil, true}, + {"audio-only input", []string{"audio"}, []string{"text"}, false}, + {"image-only output", []string{"text"}, []string{"image"}, false}, + } { + t.Run(tc.name, func(t *testing.T) { + model := modelsDevModel{Modalities: modelsDevModalities{Input: tc.input, Output: tc.output}} + if got := supportsTextIO(model); got != tc.want { + t.Errorf("supportsTextIO(in=%v out=%v) = %v, want %v", tc.input, tc.output, got, tc.want) + } + }) + } +} + +func TestDeriveSupportedEffortsDropsNone(t *testing.T) { + got := deriveSupportedEfforts("openai", modelsDevModel{ReasoningOptions: []modelsDevReasoningOption{{ + Type: "effort", + Values: []string{"none", "minimal", "low", "medium", "max", "max"}, + }}}) + want := []string{"low", "medium", "max"} + if len(got) != len(want) { + t.Fatalf("efforts = %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("efforts = %v, want %v", got, want) + } + } +} + +func TestDeriveSupportedEffortsKeepsOnlyExecutableProviderLevels(t *testing.T) { + effort := modelsDevReasoningOption{Type: "effort", Values: []string{"minimal", "low", "medium", "high", "xhigh", "max"}} + adaptiveAnthropic := modelsDevModel{Reasoning: true, ReasoningOptions: []modelsDevReasoningOption{effort}} + legacyAnthropic := modelsDevModel{Reasoning: true, ReasoningOptions: []modelsDevReasoningOption{{Type: "budget_tokens"}, effort}} + + if got := deriveSupportedEfforts("anthropic", adaptiveAnthropic); len(got) != 5 || got[4] != "max" { + t.Errorf("adaptive Anthropic efforts = %v, want low through max", got) + } + if got := deriveSupportedEfforts("anthropic", legacyAnthropic); got != nil { + t.Errorf("legacy Anthropic efforts = %v, want nil", got) + } + if got := deriveSupportedEfforts("deepseek", modelsDevModel{ReasoningOptions: []modelsDevReasoningOption{effort}}); got != nil { + t.Errorf("DeepSeek efforts = %v, want nil", got) + } +} diff --git a/pkg/ai/internal/gen-model-registry/patches.json b/pkg/ai/internal/gen-model-registry/patches.json new file mode 100644 index 0000000..cafcce8 --- /dev/null +++ b/pkg/ai/internal/gen-model-registry/patches.json @@ -0,0 +1,104 @@ +{ + "claude-opus-5": { + "preferred": true + }, + "claude-sonnet-5": { + "preferred": true + }, + "claude-fable-5": { + "preferred": true + }, + "claude-opus-4-8": { + "preferred": true + }, + "claude-sonnet-4-6": { + "preferred": false + }, + "claude-haiku-4-5": { + "preferred": true + }, + "gpt-5.6": { + "preferred": true, + "availability": [ + "api" + ] + }, + "gpt-5.6-sol": { + "preferred": true, + "availability": [ + "api", + "codex" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max", + "ultra" + ], + "priority": 1, + "aliases": [ + "sol" + ] + }, + "gpt-5.6-terra": { + "preferred": true, + "availability": [ + "api", + "codex" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max", + "ultra" + ], + "priority": 2, + "aliases": [ + "terra" + ] + }, + "gpt-5.6-luna": { + "preferred": true, + "availability": [ + "api", + "codex" + ], + "priority": 3, + "aliases": [ + "luna" + ] + }, + "gpt-5.5": { + "preferred": false + }, + "gemini-3.6-flash": { + "preferred": true + }, + "gemini-3.5-flash": { + "preferred": true + }, + "gemini-2.5-pro": { + "preferred": false + }, + "gemini-2.5-flash": { + "preferred": false + }, + "deepseek-v4-pro": { + "preferred": true + }, + "deepseek-v4-flash": { + "preferred": true + }, + "deepseek-reasoner": null, + "deepseek-chat": null, + "claude-sonnet-4-5": { + "supersededBy": "claude-sonnet-4-6" + }, + "claude-sonnet-4-5-20250929": { + "supersededBy": "claude-sonnet-4-6" + } +} diff --git a/pkg/ai/live_catalog.go b/pkg/ai/live_catalog.go new file mode 100644 index 0000000..984c6df --- /dev/null +++ b/pkg/ai/live_catalog.go @@ -0,0 +1,130 @@ +package ai + +import ( + "time" + + "github.com/flanksource/captain/pkg/api" +) + +// LiveCatalog is the model menu reconciled with what the host actually exposes +// right now (the cached whoami probe): live provider /v1/models and the codex +// debug catalog are merged over the static registry projection, keyed by menu +// ID. Static entries are never dropped — an API model with no key stays in the +// menu (rendered disabled by LiveCatalogInfo) so the picker still communicates +// what a key would unlock — but any model the probe describes with fresher data +// overrides its static counterpart. +func LiveCatalog() ([]Model, error) { + adapters, err := CachedAdapters(time.Now()) + if err != nil { + return nil, err + } + return mergeLiveCatalog(Catalog(), adapters), nil +} + +// LiveCatalogInfo annotates the live catalog with per-caller selectability, +// reusing the same rules as CatalogInfo: API models are configured when their +// provider key is present (configuredProviders), agent/CLI models when their +// local backend binary is installed. +func LiveCatalogInfo(configuredProviders []string) ([]ModelInfo, error) { + models, err := LiveCatalog() + if err != nil { + return nil, err + } + return catalogInfoFrom(models, configuredProviders), nil +} + +// mergeLiveCatalog upserts each probed model onto the static catalog. Ordering +// follows the static catalog, with live-only models appended in probe order so +// the menu stays stable across refreshes. +func mergeLiveCatalog(static []Model, adapters []AdapterStatus) []Model { + out := append([]Model(nil), static...) + pos := make(map[string]int, len(out)) + for i, m := range out { + pos[m.ID] = i + } + + for _, a := range adapters { + menuBackend, hasMenu := menuBackendFor(Backend(a.Backend)) + if !hasMenu { + continue + } + for _, md := range a.ModelDetails { + live := liveModel(menuBackend, md) + if idx, seen := pos[live.ID]; seen { + out[idx] = mergeModel(out[idx], live) + continue + } + pos[live.ID] = len(out) + out = append(out, live) + } + } + return out +} + +// menuBackendFor maps a probed backend to the backend the menu uses for that +// model, and whether the backend has a menu representation at all. The three +// claude execution backends collapse onto claude-agent and the three codex ones +// onto codex-agent — the menu offers one agent entry per model. Backends with no +// menu representation (gemini-cli, whose models already appear under the googleai +// API entry) return false and are skipped. Whether an id is provider-prefixed is +// decided later by the menu backend's Kind (see liveModel). +func menuBackendFor(b Backend) (Backend, bool) { + switch b { + case BackendAnthropic, BackendOpenAI, BackendGemini, BackendDeepSeek: + return b, true + case BackendClaudeAgent, BackendClaudeCLI, BackendClaudeCmux: + return BackendClaudeAgent, true + case BackendCodexAgent, BackendCodexCLI, BackendCodexCmux: + return BackendCodexAgent, true + default: + return "", false + } +} + +// liveModel builds a catalog Model from one probed model detail, in the menu's +// id convention (provider-prefixed for API backends, exact id for agent +// backends). ContextWindow/AdaptiveThinking are not carried by the live probe; +// a model already in the static catalog keeps them via mergeModel, and a +// live-only model leaves ContextWindow zero (the usage gauge degrades to no +// denominator rather than a fabricated one). +func liveModel(menuBackend Backend, md ModelDef) Model { + bare := bareProviderModelID(md.ID) + id := bare + if menuBackend.Kind() == "api" { + id = BackendToProvider(menuBackend) + "/" + bare + } + label := md.Name + if label == "" { + label = bare + } + return Model{ + ID: id, + Backend: menuBackend, + Label: label, + Reasoning: md.Reasoning, + Temperature: md.Temperature, + ReleaseDate: md.ReleaseDate, + SupportedEfforts: append([]api.Effort(nil), md.SupportedEfforts...), + DefaultEffort: md.DefaultEffort, + Priority: md.Priority, + } +} + +// mergeModel refreshes a static catalog entry with live probe data while keeping +// the richer static metadata (menu label, context window, adaptive-thinking +// flag) that the live probe does not carry. +func mergeModel(static, live Model) Model { + merged := static + if live.ReleaseDate != "" { + merged.ReleaseDate = live.ReleaseDate + } + if len(live.SupportedEfforts) > 0 { + merged.SupportedEfforts = live.SupportedEfforts + } + if live.DefaultEffort != api.EffortNone { + merged.DefaultEffort = live.DefaultEffort + } + merged.Reasoning = live.Reasoning + merged.Temperature = live.Temperature + return merged +} diff --git a/pkg/ai/live_catalog_test.go b/pkg/ai/live_catalog_test.go new file mode 100644 index 0000000..d7b710d --- /dev/null +++ b/pkg/ai/live_catalog_test.go @@ -0,0 +1,110 @@ +package ai + +import ( + "testing" + "time" + + "github.com/flanksource/captain/pkg/api" +) + +func findModel(t *testing.T, models []Model, id string) Model { + t.Helper() + for _, m := range models { + if m.ID == id { + return m + } + } + t.Fatalf("model %q not in catalog %+v", id, modelIDsFrom(models)) + return Model{} +} + +func TestMergeLiveCatalogUpsertsLiveAndPreservesStatic(t *testing.T) { + static := []Model{ + {ID: "anthropic/claude-sonnet-5", Backend: BackendAnthropic, Label: "Claude Sonnet 5", ContextWindow: 1_000_000, ReleaseDate: "2026-06-01"}, + {ID: "claude-opus-4-8", Backend: BackendClaudeAgent, Label: "Claude Agent · Opus 4.8", ContextWindow: 1_000_000}, + } + adapters := []AdapterStatus{ + {Backend: string(BackendCodexCLI), Type: "cli", ModelDetails: []ModelDef{ + {ID: "gpt-5.6-sol", Name: "GPT-5.6-Sol", Reasoning: true, ReleaseDate: "2026-07-09", SupportedEfforts: []api.Effort{api.EffortLow, api.EffortMax}}, + }}, + {Backend: string(BackendAnthropic), Type: "api", ModelDetails: []ModelDef{ + {ID: "claude-sonnet-5", Name: "Claude Sonnet 5", Reasoning: true, ReleaseDate: "2026-06-29"}, + }}, + // gemini-cli has no menu backend; its models must not leak into the menu. + {Backend: string(BackendGeminiCLI), Type: "cli", ModelDetails: []ModelDef{ + {ID: "gemini-3.5-flash", Name: "Gemini 3.5 Flash"}, + }}, + } + + merged := mergeLiveCatalog(static, adapters) + + // A live codex model becomes one codex-agent entry keyed by its exact id. + sol := findModel(t, merged, "gpt-5.6-sol") + if sol.Backend != BackendCodexAgent { + t.Errorf("sol backend = %q, want codex-agent", sol.Backend) + } + if !sol.Reasoning || sol.ReleaseDate != "2026-07-09" || len(sol.SupportedEfforts) != 2 { + t.Errorf("sol not projected from live probe: %+v", sol) + } + + // The API entry is upserted: static context window preserved, live release + // date wins. + sonnet := findModel(t, merged, "anthropic/claude-sonnet-5") + if sonnet.ContextWindow != 1_000_000 { + t.Errorf("sonnet context window = %d, want static 1000000 preserved", sonnet.ContextWindow) + } + if sonnet.ReleaseDate != "2026-06-29" { + t.Errorf("sonnet release date = %q, want live 2026-06-29", sonnet.ReleaseDate) + } + + // A static entry the probe did not rediscover is retained (shown disabled by + // LiveCatalogInfo when its provider is unconfigured). + findModel(t, merged, "claude-opus-4-8") + + for _, m := range merged { + if m.ID == "googleai/gemini-3.5-flash" || m.Backend == BackendGeminiCLI { + t.Errorf("gemini-cli model leaked into the menu: %+v", m) + } + } +} + +func TestLiveCatalogInfoAppliesPerProviderConfigured(t *testing.T) { + prevProbe := adapterProbe + prevCache, prevAt := adapterCache, adapterCacheAt + t.Cleanup(func() { + adapterProbe = prevProbe + adapterCache, adapterCacheAt = prevCache, prevAt + }) + adapterCache, adapterCacheAt = nil, time.Time{} + adapterProbe = func() ([]AdapterStatus, error) { + return []AdapterStatus{ + {Backend: string(BackendAnthropic), Type: "api", ModelDetails: []ModelDef{ + {ID: "claude-sonnet-5", Name: "Claude Sonnet 5", Reasoning: true}, + }}, + {Backend: string(BackendOpenAI), Type: "api", ModelDetails: []ModelDef{ + {ID: "gpt-5.5", Name: "GPT-5.5", Reasoning: true}, + }}, + }, nil + } + + infos, err := LiveCatalogInfo([]string{"anthropic"}) + if err != nil { + t.Fatalf("LiveCatalogInfo: %v", err) + } + byID := map[string]ModelInfo{} + for _, info := range infos { + byID[info.ID] = info + } + + anthropic, ok := byID["anthropic/claude-sonnet-5"] + if !ok || !anthropic.Configured { + t.Errorf("anthropic model should be configured when its provider key is present: %+v", anthropic) + } + if anthropic.Provider != "anthropic" { + t.Errorf("anthropic provider = %q", anthropic.Provider) + } + openai, ok := byID["openai/gpt-5.5"] + if !ok || openai.Configured { + t.Errorf("openai model should be unconfigured (no key) but still listed: %+v", openai) + } +} diff --git a/pkg/ai/loop.go b/pkg/ai/loop.go index 8f0ba6f..4f550be 100644 --- a/pkg/ai/loop.go +++ b/pkg/ai/loop.go @@ -131,6 +131,18 @@ func runOneIteration(ctx context.Context, opts LoopOptions, req Request, iter *L if ev.Usage != nil { iter.Usage = *ev.Usage } + // Providers that report no cost (codex-cli/cmux, gemini-cli) leave + // CostUSD at 0, which would make budget enforcement and rollups treat + // the run as free (finding C2). Price the usage from the registry so + // MaxCostUSD sees a real number for every backend. + if iter.CostUSD == 0 && ev.Usage != nil { + backend := opts.Provider.GetBackend() + model := ev.Model + if model == "" { + model = opts.Provider.GetModel() + } + iter.CostUSD = PriceResponse(backend, model, &Response{Backend: backend, Model: model, Usage: *ev.Usage}).Total() + } if ev.Error != "" && iter.Err == nil { iter.Err = fmt.Errorf("claude returned: %s", ev.Error) } diff --git a/pkg/ai/loop_test.go b/pkg/ai/loop_test.go index 31f0a50..00061b8 100644 --- a/pkg/ai/loop_test.go +++ b/pkg/ai/loop_test.go @@ -15,10 +15,22 @@ type fakeStreamingProvider struct { scripts [][]Event requests []Request err error + backend Backend + model string } -func (f *fakeStreamingProvider) GetModel() string { return "fake" } -func (f *fakeStreamingProvider) GetBackend() Backend { return Backend("fake") } +func (f *fakeStreamingProvider) GetModel() string { + if f.model != "" { + return f.model + } + return "fake" +} +func (f *fakeStreamingProvider) GetBackend() Backend { + if f.backend != "" { + return f.backend + } + return Backend("fake") +} func (f *fakeStreamingProvider) Execute(ctx context.Context, req Request) (*Response, error) { return nil, errors.New("not used") } @@ -120,6 +132,35 @@ func TestRunUntil_HitsMaxIterations(t *testing.T) { } } +// A provider that reports usage but no cost (codex-cli/cmux, gemini-cli) must +// still accrue a real cost so MaxCostUSD can enforce a budget for it (finding +// C2). The loop prices the usage from the registry. +func TestRunUntil_PricesUsageWhenProviderReportsNoCost(t *testing.T) { + p := &fakeStreamingProvider{ + backend: BackendAnthropic, + model: "claude-sonnet-4", + scripts: [][]Event{{ + {Kind: EventResult, Success: true, CostUSD: 0, Usage: &Usage{InputTokens: 1_000_000, OutputTokens: 1_000_000}}, + }}, + } + res, err := RunUntil(context.Background(), LoopOptions{ + Provider: p, + MaxIterations: 1, + BuildRequest: func(iter int, prev *LoopIteration) (Request, bool) { + if iter > 0 { + return Request{}, false + } + return Request{Prompt: api.Prompt{User: "loop"}}, true + }, + }) + if err != nil { + t.Fatalf("RunUntil err: %v", err) + } + if res.TotalCost <= 0 { + t.Fatalf("TotalCost = %v, want > 0 (loop must price usage the provider left uncosted)", res.TotalCost) + } +} + func TestRunUntil_HitsMaxCost(t *testing.T) { p := &fakeStreamingProvider{ scripts: [][]Event{ diff --git a/pkg/ai/middleware/caching.go b/pkg/ai/middleware/caching.go index e5e9820..484d3d2 100644 --- a/pkg/ai/middleware/caching.go +++ b/pkg/ai/middleware/caching.go @@ -16,13 +16,15 @@ type cachingProvider struct { func (c *cachingProvider) GetModel() string { return c.provider.GetModel() } func (c *cachingProvider) GetBackend() ai.Backend { return c.provider.GetBackend() } +func (c *cachingProvider) Unwrap() ai.Provider { return c.provider } func (c *cachingProvider) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { if ShouldBypassCache(ctx) { return c.provider.Execute(ctx, req) } - entry, err := c.cache.Get(req.Prompt.User, c.provider.GetModel()) + cacheIdentity := req.Prompt.CacheIdentity() + entry, err := c.cache.Get(cacheIdentity, c.provider.GetModel()) if err == nil && entry != nil && entry.Error == "" { log.Debugf("[%s/%s] cache hit", c.provider.GetBackend(), c.provider.GetModel()) return &ai.Response{ @@ -50,7 +52,7 @@ func (c *cachingProvider) Execute(ctx context.Context, req ai.Request) (*ai.Resp cacheEntry := &cache.Entry{ Model: c.provider.GetModel(), - Prompt: req.Prompt.User, + Prompt: cacheIdentity, DurationMS: duration.Milliseconds(), Provider: string(c.provider.GetBackend()), } diff --git a/pkg/ai/middleware/construct.go b/pkg/ai/middleware/construct.go index 81f4721..249740a 100644 --- a/pkg/ai/middleware/construct.go +++ b/pkg/ai/middleware/construct.go @@ -11,7 +11,7 @@ import ( // ai.NewProvider, because middleware imports ai (wrapping ai.Provider) and the // reverse import would be a cycle. func DefaultOptions(cfg ai.Config) []Option { - opts := []Option{WithLogging(), WithSchemaValidation()} + opts := []Option{WithLogging(), WithSchemaValidation(cfg)} if !cfg.NoCache && (cfg.CacheTTL > 0 || cfg.CacheDBPath != "") { opts = append(opts, WithCache(cache.Config{ DBPath: cfg.CacheDBPath, diff --git a/pkg/ai/middleware/cost.go b/pkg/ai/middleware/cost.go deleted file mode 100644 index 1ed7042..0000000 --- a/pkg/ai/middleware/cost.go +++ /dev/null @@ -1,67 +0,0 @@ -package middleware - -import ( - "context" - "fmt" - - "github.com/flanksource/captain/pkg/ai" - "github.com/flanksource/captain/pkg/ai/pricing" - "github.com/flanksource/captain/pkg/ai/session" -) - -type costProvider struct { - provider ai.Provider - session *session.Session - budgetUSD float64 -} - -func (c *costProvider) GetModel() string { return c.provider.GetModel() } -func (c *costProvider) GetBackend() ai.Backend { return c.provider.GetBackend() } - -func (c *costProvider) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { - if c.budgetUSD > 0 && c.session.TotalCost() >= c.budgetUSD { - return nil, fmt.Errorf("%w: spent $%.4f of $%.4f budget", - ai.ErrBudgetExceeded, c.session.TotalCost(), c.budgetUSD) - } - - resp, err := c.provider.Execute(ctx, req) - if err != nil { - return resp, err - } - - model := "anthropic/" + resp.Model - if resp.Backend == ai.BackendGemini || resp.Backend == ai.BackendGeminiCLI { - model = "google/" + resp.Model - } - - result, calcErr := pricing.CalculateCost(model, - resp.Usage.InputTokens, resp.Usage.OutputTokens, - resp.Usage.ReasoningTokens, resp.Usage.CacheReadTokens, resp.Usage.CacheWriteTokens) - - if calcErr != nil { - log.Debugf("Cost calculation failed for %s: %v", resp.Model, calcErr) - } else { - c.session.AddCost(ai.Cost{ - Model: resp.Model, - InputTokens: result.InputTokens, - OutputTokens: result.OutputTokens, - ReasoningTokens: result.ReasoningTokens, - CacheReadTokens: result.CacheReadTokens, - CacheWriteTokens: result.CacheWriteTokens, - TotalTokens: resp.Usage.TotalTokens(), - InputCost: result.InputCost, - OutputCost: result.OutputCost, - ReasoningCost: result.ReasoningCost, - CacheReadCost: result.CacheReadCost, - CacheWriteCost: result.CacheWriteCost, - }) - } - - return resp, nil -} - -func WithCostTracking(sess *session.Session, budgetUSD float64) Option { - return func(p ai.Provider) (ai.Provider, error) { - return &costProvider{provider: p, session: sess, budgetUSD: budgetUSD}, nil - } -} diff --git a/pkg/ai/middleware/logging.go b/pkg/ai/middleware/logging.go index e3a466b..81a6044 100644 --- a/pkg/ai/middleware/logging.go +++ b/pkg/ai/middleware/logging.go @@ -23,13 +23,16 @@ type loggingProvider struct { func (l *loggingProvider) GetModel() string { return l.provider.GetModel() } func (l *loggingProvider) GetBackend() ai.Backend { return l.provider.GetBackend() } +func (l *loggingProvider) Unwrap() ai.Provider { return l.provider } func (l *loggingProvider) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { start := time.Now() + backend, model := logRuntime(l.provider, req) + identity := runtimeLogIdentity(backend, model, req.Effort) dispatch := clicky.Text(""). Add(icons.AI). - Append(fmt.Sprintf(" %s/%s", l.provider.GetBackend(), l.provider.GetModel()), "text-purple-600 font-medium") + Append(" "+identity, "text-purple-600 font-medium") if req.Prompt.Source != "" { dispatch = dispatch.Append(fmt.Sprintf(" [%s]", req.Prompt.Source), "text-gray-500") } @@ -49,14 +52,14 @@ func (l *loggingProvider) Execute(ctx context.Context, req ai.Request) (*ai.Resp if err != nil { log.Errorf("%v", clicky.Text(""). Add(icons.Error). - Append(fmt.Sprintf(" %s/%s", l.provider.GetBackend(), l.provider.GetModel()), "text-red-600 font-medium"). + Append(" "+identity, "text-red-600 font-medium"). Append(fmt.Sprintf(" failed after %v: %v", duration.Round(time.Millisecond), err), "text-red-500")) return resp, err } log.Infof("%v", clicky.Text(""). Add(icons.Check). - Append(fmt.Sprintf(" %s/%s", l.provider.GetBackend(), l.provider.GetModel()), "text-green-600 font-medium"). + Append(" "+identity, "text-green-600 font-medium"). Append(fmt.Sprintf(" %v", duration.Round(time.Millisecond)), "text-gray-500"). Append(fmt.Sprintf(" (tokens: %d in / %d out)", resp.Usage.InputTokens, resp.Usage.OutputTokens), "text-gray-400")) @@ -85,10 +88,12 @@ func (l *loggingProvider) ExecuteStream(ctx context.Context, req ai.Request) (<- if !ok { return nil, fmt.Errorf("provider %s/%s does not support streaming", l.provider.GetBackend(), l.provider.GetModel()) } + backend, model := logRuntime(l.provider, req) + identity := runtimeLogIdentity(backend, model, req.Effort) dispatch := clicky.Text(""). Add(icons.AI). - Append(fmt.Sprintf(" %s/%s (stream)", l.provider.GetBackend(), l.provider.GetModel()), "text-purple-600 font-medium") + Append(" "+identity+" (stream)", "text-purple-600 font-medium") if req.Prompt.Source != "" { dispatch = dispatch.Append(fmt.Sprintf(" [%s]", req.Prompt.Source), "text-gray-500") } @@ -101,6 +106,40 @@ func (l *loggingProvider) ExecuteStream(ctx context.Context, req ai.Request) (<- return streamer.ExecuteStream(ctx, req) } +func logRuntime(provider ai.Provider, req ai.Request) (api.Backend, string) { + backend := req.Backend + if backend == "" { + backend = provider.GetBackend() + } + model := req.Name + if model == "" { + model = provider.GetModel() + } + return backend, model +} + +// runtimeLogIdentity renders the same compact selector notation accepted by +// Captain's model flags: mode:model[:effort]. The effort suffix is omitted when +// the request leaves effort at the backend/model default. +func runtimeLogIdentity(backend api.Backend, model string, effort api.Effort) string { + prefix := string(backend) + switch backend { + case api.BackendClaudeAgent, api.BackendCodexAgent: + prefix = "agent" + case api.BackendClaudeCLI, api.BackendCodexCLI, api.BackendGeminiCLI: + prefix = "cli" + case api.BackendClaudeCmux, api.BackendCodexCmux: + prefix = "cmux" + case api.BackendAnthropic, api.BackendGemini, api.BackendOpenAI, api.BackendDeepSeek: + prefix = "api" + } + identity := prefix + ":" + model + if effort != api.EffortNone { + identity += ":" + string(effort) + } + return identity +} + func WithLogging() Option { return func(p ai.Provider) (ai.Provider, error) { return &loggingProvider{provider: p}, nil diff --git a/pkg/ai/middleware/logging_test.go b/pkg/ai/middleware/logging_test.go index fedb257..76a5b09 100644 --- a/pkg/ai/middleware/logging_test.go +++ b/pkg/ai/middleware/logging_test.go @@ -28,6 +28,95 @@ func (s stubProvider) Execute(context.Context, ai.Request) (*ai.Response, error) func (s stubProvider) GetModel() string { return "test-model" } func (s stubProvider) GetBackend() ai.Backend { return ai.BackendAnthropic } +type streamingStubProvider struct{ stubProvider } + +func (s streamingStubProvider) ExecuteStream(context.Context, ai.Request) (<-chan ai.Event, error) { + events := make(chan ai.Event) + close(events) + return events, nil +} + +func TestRuntimeLogIdentity(t *testing.T) { + tests := []struct { + name string + backend api.Backend + model string + effort api.Effort + want string + }{ + {"agent without effort", api.BackendCodexAgent, "gpt-5.6-luna", api.EffortNone, "agent:gpt-5.6-luna"}, + {"agent with effort", api.BackendCodexAgent, "gpt-5.6-sol", api.EffortMax, "agent:gpt-5.6-sol:max"}, + {"cli", api.BackendClaudeCLI, "claude-sonnet-5", api.EffortHigh, "cli:claude-sonnet-5:high"}, + {"cmux", api.BackendCodexCmux, "gpt-5.6-terra", api.EffortUltra, "cmux:gpt-5.6-terra:ultra"}, + {"api", api.BackendOpenAI, "gpt-5.6", api.EffortLow, "api:gpt-5.6:low"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := runtimeLogIdentity(tt.backend, tt.model, tt.effort); got != tt.want { + t.Fatalf("runtimeLogIdentity() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestLoggingProvider_UsesEffortQualifiedRuntimeIdentity(t *testing.T) { + prev := logger.GetOutput() + t.Cleanup(func() { logger.SetOutput(prev) }) + var buf bytes.Buffer + logger.SetOutput(&buf) + logger.GetLogger("ai").SetLogLevel("info") + + p, err := WithLogging()(stubProvider{resp: &ai.Response{Model: "gpt-5.6-sol"}}) + if err != nil { + t.Fatalf("WithLogging: %v", err) + } + if _, err := p.Execute(context.Background(), ai.Request{ + Model: api.Model{Name: "gpt-5.6-sol", Backend: api.BackendCodexAgent, Effort: api.EffortMax}, + Prompt: api.Prompt{Source: "commit-grouping.prompt"}, + }); err != nil { + t.Fatalf("Execute: %v", err) + } + + got := ansi.ReplaceAllString(buf.String(), "") + for _, want := range []string{ + "✨ agent:gpt-5.6-sol:max [commit-grouping.prompt]", + "✓ agent:gpt-5.6-sol:max", + } { + if !strings.Contains(got, want) { + t.Errorf("logging middleware output missing %q\n--- output ---\n%s", want, got) + } + } +} + +func TestLoggingProvider_StreamOmitsEmptyEffort(t *testing.T) { + prev := logger.GetOutput() + t.Cleanup(func() { logger.SetOutput(prev) }) + var buf bytes.Buffer + logger.SetOutput(&buf) + logger.GetLogger("ai").SetLogLevel("info") + + p, err := WithLogging()(streamingStubProvider{stubProvider{resp: &ai.Response{}}}) + if err != nil { + t.Fatalf("WithLogging: %v", err) + } + streamer := p.(ai.StreamingProvider) + if _, err := streamer.ExecuteStream(context.Background(), ai.Request{ + Model: api.Model{Name: "gpt-5.6-luna", Backend: api.BackendCodexAgent}, + Prompt: api.Prompt{Source: "commit-grouping.prompt"}, + }); err != nil { + t.Fatalf("ExecuteStream: %v", err) + } + + got := ansi.ReplaceAllString(buf.String(), "") + want := "✨ agent:gpt-5.6-luna (stream) [commit-grouping.prompt]" + if !strings.Contains(got, want) { + t.Errorf("stream logging output missing %q\n--- output ---\n%s", want, got) + } + if strings.Contains(got, "luna:") { + t.Errorf("stream logging unexpectedly added an effort suffix\n--- output ---\n%s", got) + } +} + func TestSchemaInJSON(t *testing.T) { if got := schemaInJSON(api.Prompt{}); got != "" { t.Fatalf("text-mode (no schema): want empty, got %q", got) diff --git a/pkg/ai/middleware/retry.go b/pkg/ai/middleware/retry.go index 481bf8a..2dab7cb 100644 --- a/pkg/ai/middleware/retry.go +++ b/pkg/ai/middleware/retry.go @@ -29,6 +29,7 @@ type retryProvider struct { func (r *retryProvider) GetModel() string { return r.provider.GetModel() } func (r *retryProvider) GetBackend() ai.Backend { return r.provider.GetBackend() } +func (r *retryProvider) Unwrap() ai.Provider { return r.provider } func (r *retryProvider) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { var lastErr error diff --git a/pkg/ai/middleware/schema_repair.go b/pkg/ai/middleware/schema_repair.go new file mode 100644 index 0000000..568a499 --- /dev/null +++ b/pkg/ai/middleware/schema_repair.go @@ -0,0 +1,124 @@ +package middleware + +import ( + "context" + "embed" + "encoding/json" + "fmt" + "io" + "strings" + + "github.com/flanksource/captain/pkg/ai" + promptlib "github.com/flanksource/captain/pkg/ai/prompt" + "github.com/flanksource/captain/pkg/api" +) + +//go:embed schema_repair.prompt +var schemaRepairPrompts embed.FS + +const defaultSchemaRepairPrompt = "schema_repair.prompt" + +func (v *validatingProvider) executeRepair(ctx context.Context, parent ai.Request, schema json.RawMessage, verrs string, prev *ai.Response, attempt int) (*ai.Response, error) { + req, cfg, useParent, err := v.repairRequest(parent, schema, verrs, prev, attempt) + if err != nil { + return nil, err + } + if useParent { + return v.provider.Execute(ctx, req) + } + p, err := ai.NewProvider(cfg) + if err != nil { + return nil, err + } + if c, ok := p.(io.Closer); ok { + defer func() { _ = c.Close() }() + } + return p.Execute(ctx, req) +} + +func (v *validatingProvider) repairRequest(parent ai.Request, schema json.RawMessage, verrs string, prev *ai.Response, attempt int) (ai.Request, ai.Config, bool, error) { + tmpl, err := v.repairTemplate() + if err != nil { + return ai.Request{}, ai.Config{}, false, err + } + renderedReq, renderedCfg, err := tmpl.Render(map[string]any{ + "attempt": attempt, + "backend": string(v.provider.GetBackend()), + "model": v.provider.GetModel(), + "originalPrompt": parent.Prompt.User, + "previousResponse": responseJSON(prev), + "schema": string(schema), + "validationErrors": verrs, + }, nil) + if err != nil { + return ai.Request{}, ai.Config{}, false, fmt.Errorf("render schema repair prompt: %w", err) + } + + parentModel := parent.Model + if parentModel.Name == "" { + parentModel.Name = v.provider.GetModel() + } + if parentModel.Backend == "" { + parentModel.Backend = v.provider.GetBackend() + } + + model := parentModel + model = overlayModel(model, renderedCfg.Model) + model = overlayModel(model, renderedReq.Model) + model = overlayModel(model, v.cfg.SchemaRepair.Model) + + req := parent + req.Model = model + req.Prompt = renderedReq.Prompt + req.Prompt.Schema = parent.Prompt.Schema + req.Prompt.SchemaJSON = parent.Prompt.SchemaJSON + req.Prompt.SchemaStrictness = api.SchemaStrictnessDisabled + req.Workflow = nil + + cfg := v.cfg + cfg.Model = model + cfg.SchemaRepair = api.SchemaRepairConfig{} + if model.Backend != "" && model.Backend != v.provider.GetBackend() { + cfg.APIKey = "" + } + + return req, cfg, sameModelBackend(model, parentModel), nil +} + +func (v *validatingProvider) repairTemplate() (*promptlib.Template, error) { + if path := strings.TrimSpace(v.cfg.SchemaRepair.Prompt); path != "" { + return promptlib.LoadFile(path) + } + return promptlib.LoadFS(schemaRepairPrompts, defaultSchemaRepairPrompt) +} + +func overlayModel(base, overlay api.Model) api.Model { + if overlay.Name != "" { + base.Name = overlay.Name + } + if overlay.ID != "" { + base.ID = overlay.ID + } + if overlay.Backend != "" { + base.Backend = overlay.Backend + } + if overlay.Temperature != nil { + base.Temperature = overlay.Temperature + } + if overlay.Effort != "" { + base.Effort = overlay.Effort + } + if overlay.NoCache { + base.NoCache = true + } + if len(overlay.Fallbacks) > 0 { + base.Fallbacks = overlay.Fallbacks + } + return base +} + +func sameModelBackend(a, b api.Model) bool { + return strings.TrimSpace(a.Name) == strings.TrimSpace(b.Name) && + strings.TrimSpace(a.ID) == strings.TrimSpace(b.ID) && + a.Backend == b.Backend +} diff --git a/pkg/ai/middleware/schema_repair.prompt b/pkg/ai/middleware/schema_repair.prompt new file mode 100644 index 0000000..dacaf43 --- /dev/null +++ b/pkg/ai/middleware/schema_repair.prompt @@ -0,0 +1,19 @@ +{{role "system"}} +You repair JSON responses so they comply with the required JSON Schema. +Return only valid JSON. Do not include markdown fences, explanations, or prose. +{{role "user"}} +The previous response did not comply with the JSON Schema. + +Original prompt: +{{{originalPrompt}}} + +Previous response: +{{{previousResponse}}} + +Validation errors: +{{{validationErrors}}} + +JSON Schema: +{{{schema}}} + +Return a corrected JSON value that satisfies the schema. diff --git a/pkg/ai/middleware/validation.go b/pkg/ai/middleware/validation.go index 015c83e..508e9ef 100644 --- a/pkg/ai/middleware/validation.go +++ b/pkg/ai/middleware/validation.go @@ -3,41 +3,62 @@ package middleware import ( "context" "encoding/json" + "errors" "fmt" "strings" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" - "github.com/xeipuuv/gojsonschema" ) +// maxSchemaRetries bounds the number of fix-up re-asks the retry strictness makes +// when a structured response fails schema validation before giving up with a hard +// error. +const maxSchemaRetries = 2 + // validatingProvider validates a structured-output response against the request's // JSON schema and applies the prompt's SchemaStrictness policy: warning (log and // continue), error (fail), retry (re-ask the model once with the validation error, // then fail). It is a no-op unless the prompt sets both a schema and a strictness -// mode, so it is safe to keep in the default middleware stack. +// mode; Anthropic native structured-output backends default to retry validation +// because their provider-facing schema is intentionally stripped down. type validatingProvider struct { provider ai.Provider + cfg ai.Config } func (v *validatingProvider) GetModel() string { return v.provider.GetModel() } func (v *validatingProvider) GetBackend() ai.Backend { return v.provider.GetBackend() } +func (v *validatingProvider) Unwrap() ai.Provider { return v.provider } // WithSchemaValidation validates structured responses against the request schema // and enforces api.Prompt.SchemaStrictness. -func WithSchemaValidation() Option { +func WithSchemaValidation(cfg ...ai.Config) Option { return func(p ai.Provider) (ai.Provider, error) { - return &validatingProvider{provider: p}, nil + var c ai.Config + if len(cfg) > 0 { + c = cfg[0] + } + return &validatingProvider{provider: p, cfg: c}, nil } } func (v *validatingProvider) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { resp, err := v.provider.Execute(ctx, req) + strictness := v.effectiveStrictness(req) if err != nil { + // A provider that validates the model's structured output during generation + // (e.g. genkit's constrained output) rejects it before we ever see a + // response. Under "retry" that rejection is recoverable: re-ask the model + // with the schema errors instead of failing outright. + if strictness == api.SchemaStrictnessRetry && errors.Is(err, ai.ErrSchemaValidation) { + if schema, serr := ai.SchemaJSONFor(req.Prompt); serr == nil && len(schema) > 0 { + return v.retryWithFeedback(ctx, req, schema, err.Error(), nil) + } + } return resp, err } - strictness := req.Prompt.SchemaStrictness if strictness == api.SchemaStrictnessNone { return resp, nil } @@ -59,7 +80,7 @@ func (v *validatingProvider) Execute(ctx context.Context, req ai.Request) (*ai.R log.Warnf("schema validation failed (%s/%s): %s", v.provider.GetBackend(), v.provider.GetModel(), verrs) return resp, nil case api.SchemaStrictnessRetry: - return v.retry(ctx, req, resp, schema, verrs) + return v.retryWithFeedback(ctx, req, schema, verrs, resp) case api.SchemaStrictnessError: return resp, fmt.Errorf("%w: %s", ai.ErrSchemaValidation, verrs) default: @@ -67,31 +88,49 @@ func (v *validatingProvider) Execute(ctx context.Context, req ai.Request) (*ai.R } } -// retry re-asks the model once, feeding back the previous response and the schema -// validation errors, then re-validates. A still-invalid response is a hard error -// (retry → then error). -func (v *validatingProvider) retry(ctx context.Context, req ai.Request, prev *ai.Response, schema json.RawMessage, verrs string) (*ai.Response, error) { - correction := fmt.Sprintf( - "\n\nYour previous response did not conform to the required JSON schema.\n"+ - "Previous response:\n%s\n\nSchema validation errors:\n%s\n\n"+ - "Return a corrected JSON response that satisfies the schema.", - responseJSON(prev), verrs) +func (v *validatingProvider) effectiveStrictness(req ai.Request) api.SchemaStrictness { + switch req.Prompt.SchemaStrictness { + case api.SchemaStrictnessDisabled: + return api.SchemaStrictnessNone + case api.SchemaStrictnessNone: + if req.Prompt.HasSchema() && ai.UsesAnthropicSchemaSubset(v.provider.GetBackend()) { + return api.SchemaStrictnessRetry + } + return api.SchemaStrictnessNone + default: + return req.Prompt.SchemaStrictness + } +} - req2 := req - req2.Prompt.User = req.Prompt.User + correction +// retryWithFeedback re-asks the model up to maxSchemaRetries times, feeding back +// the schema validation errors (and the previous response when we have one) so it +// can correct itself. It serves both failure modes: a post-response validation +// failure (prev is the offending response) and a provider-side rejection during +// generation (prev is nil because the response never surfaced). A response that is +// still invalid after the last attempt is a hard error. +func (v *validatingProvider) retryWithFeedback(ctx context.Context, req ai.Request, schema json.RawMessage, verrs string, prev *ai.Response) (*ai.Response, error) { + for attempt := 1; attempt <= maxSchemaRetries; attempt++ { + resp, err := v.executeRepair(ctx, req, schema, verrs, prev, attempt) + if err != nil { + // The provider rejected the corrected output too; feed the new errors + // back and keep trying until the retry budget is exhausted. + if errors.Is(err, ai.ErrSchemaValidation) && attempt < maxSchemaRetries { + verrs, prev = err.Error(), nil + continue + } + return resp, err + } - resp2, err := v.provider.Execute(ctx, req2) - if err != nil { - return resp2, err - } - verrs2, err := validateResponse(schema, resp2) - if err != nil { - return resp2, err - } - if verrs2 == "" { - return resp2, nil + verrs, err = validateResponse(schema, resp) + if err != nil { + return resp, err + } + if verrs == "" { + return resp, nil + } + prev = resp } - return resp2, fmt.Errorf("%w: still invalid after retry: %s", ai.ErrSchemaValidation, verrs2) + return nil, fmt.Errorf("%w: still invalid after %d retries: %s", ai.ErrSchemaValidation, maxSchemaRetries, verrs) } // ExecuteStream tees the stream: it forwards every event unchanged, accumulates @@ -106,7 +145,7 @@ func (v *validatingProvider) ExecuteStream(ctx context.Context, req ai.Request) return nil, fmt.Errorf("provider %s/%s does not support streaming", v.provider.GetBackend(), v.provider.GetModel()) } - strictness := req.Prompt.SchemaStrictness + strictness := v.effectiveStrictness(req) schema, serr := ai.SchemaJSONFor(req.Prompt) if strictness == api.SchemaStrictnessNone || serr != nil || len(schema) == 0 { return streamer.ExecuteStream(ctx, req) // nothing to validate; forward as-is @@ -177,6 +216,11 @@ func responseJSON(resp *ai.Response) string { if raw, ok := resp.StructuredData.(json.RawMessage); ok && len(raw) > 0 { return string(raw) } + if resp.StructuredData != nil { + if data, err := json.Marshal(resp.StructuredData); err == nil && string(data) != "null" { + return string(data) + } + } return resp.Text } @@ -184,20 +228,5 @@ func responseJSON(resp *ai.Response) string { // returns a human-readable joined error string (empty when the response conforms), // plus a hard error only when the schema or document could not be loaded/parsed. func validateResponse(schema json.RawMessage, resp *ai.Response) (string, error) { - doc := responseJSON(resp) - if strings.TrimSpace(doc) == "" { - return "response carried no JSON to validate", nil - } - result, err := gojsonschema.Validate(gojsonschema.NewBytesLoader(schema), gojsonschema.NewStringLoader(doc)) - if err != nil { - return "", fmt.Errorf("%w: validation could not run: %v", ai.ErrSchemaValidation, err) - } - if result.Valid() { - return "", nil - } - msgs := make([]string, 0, len(result.Errors())) - for _, e := range result.Errors() { - msgs = append(msgs, e.String()) - } - return strings.Join(msgs, "; "), nil + return ai.ValidateStructuredJSON(schema, responseJSON(resp)) } diff --git a/pkg/ai/middleware/validation_test.go b/pkg/ai/middleware/validation_test.go index 0d2c83c..0ed1a09 100644 --- a/pkg/ai/middleware/validation_test.go +++ b/pkg/ai/middleware/validation_test.go @@ -4,6 +4,10 @@ import ( "context" "encoding/json" "errors" + "fmt" + "os" + "path/filepath" + "strings" "testing" "github.com/flanksource/captain/pkg/ai" @@ -24,6 +28,7 @@ const ( type scriptedProvider struct { responses []string calls int + backend ai.Backend } func (s *scriptedProvider) Execute(context.Context, ai.Request) (*ai.Response, error) { @@ -34,8 +39,49 @@ func (s *scriptedProvider) Execute(context.Context, ai.Request) (*ai.Response, e s.calls++ return &ai.Response{Text: s.responses[idx], Model: "test-model"}, nil } -func (s *scriptedProvider) GetModel() string { return "test-model" } -func (s *scriptedProvider) GetBackend() ai.Backend { return ai.BackendAnthropic } +func (s *scriptedProvider) GetModel() string { return "test-model" } +func (s *scriptedProvider) GetBackend() ai.Backend { + if s.backend != "" { + return s.backend + } + return ai.BackendOpenAI +} + +// outcome is one scripted provider result: a text response or an error. +type outcome struct { + text string + err error +} + +// erroringProvider scripts a sequence of (response|error) outcomes (repeating the +// last), records every request's user prompt, and counts calls — so a provider-side +// schema rejection and the feedback appended on the retry can both be asserted. +type erroringProvider struct { + outcomes []outcome + calls int + prompts []string + backend ai.Backend +} + +func (e *erroringProvider) Execute(_ context.Context, req ai.Request) (*ai.Response, error) { + e.prompts = append(e.prompts, req.Prompt.User) + idx := e.calls + if idx >= len(e.outcomes) { + idx = len(e.outcomes) - 1 + } + e.calls++ + if o := e.outcomes[idx]; o.err != nil { + return nil, o.err + } + return &ai.Response{Text: e.outcomes[idx].text, Model: "test-model"}, nil +} +func (e *erroringProvider) GetModel() string { return "test-model" } +func (e *erroringProvider) GetBackend() ai.Backend { + if e.backend != "" { + return e.backend + } + return ai.BackendOpenAI +} func capRequest(strictness api.SchemaStrictness) ai.Request { return ai.Request{Prompt: api.Prompt{ @@ -70,6 +116,39 @@ func TestValidation_NoStrictnessSkipsValidation(t *testing.T) { } } +func TestValidation_AnthropicNoStrictnessDefaultsToRetry(t *testing.T) { + inner := &scriptedProvider{backend: ai.BackendAnthropic, responses: []string{overCapJSON, validJSON}} + p, err := WithSchemaValidation()(inner) + if err != nil { + t.Fatalf("WithSchemaValidation: %v", err) + } + resp, err := p.Execute(context.Background(), capRequest(api.SchemaStrictnessNone)) + if err != nil { + t.Fatalf("anthropic default retry should recover, got %v", err) + } + if resp.Text != validJSON { + t.Fatalf("response = %q, want corrected JSON", resp.Text) + } + if inner.calls != 2 { + t.Fatalf("calls = %d, want initial + repair", inner.calls) + } +} + +func TestValidation_AnthropicExplicitNoneSkipsValidation(t *testing.T) { + inner := &scriptedProvider{backend: ai.BackendAnthropic, responses: []string{overCapJSON}} + p, err := WithSchemaValidation()(inner) + if err != nil { + t.Fatalf("WithSchemaValidation: %v", err) + } + resp, err := p.Execute(context.Background(), capRequest(api.SchemaStrictnessDisabled)) + if err != nil { + t.Fatalf("explicit none should pass through, got %v", err) + } + if resp.Text != overCapJSON || inner.calls != 1 { + t.Fatalf("pass-through = calls %d resp %q", inner.calls, resp.Text) + } +} + func TestValidation_WarningLogsAndReturns(t *testing.T) { resp, err, calls := execWith(t, api.SchemaStrictnessWarning, overCapJSON) if err != nil { @@ -115,12 +194,126 @@ func TestValidation_RetrySucceedsOnSecond(t *testing.T) { } func TestValidation_RetryThenErrorWhenStillInvalid(t *testing.T) { - // Both attempts violate → hard error after exactly one retry (retry → then error). - _, err, calls := execWith(t, api.SchemaStrictnessRetry, overCapJSON, overCapJSON) + // Every attempt violates → hard error after the retry budget is exhausted + // (1 initial call + maxSchemaRetries re-asks). + _, err, calls := execWith(t, api.SchemaStrictnessRetry, overCapJSON) if !errors.Is(err, ai.ErrSchemaValidation) { - t.Fatalf("want ErrSchemaValidation after retry, got %v", err) + t.Fatalf("want ErrSchemaValidation after retries, got %v", err) } - if calls != 2 { - t.Errorf("retry must re-ask exactly once; want 2 calls, got %d", calls) + if want := 1 + maxSchemaRetries; calls != want { + t.Errorf("retry must re-ask maxSchemaRetries times; want %d calls, got %d", want, calls) + } +} + +func TestValidation_RetryRecoversProviderSchemaError(t *testing.T) { + // A genkit-style hard rejection during generation on the first call, then a + // conforming response on the retry: the middleware re-asks and returns the + // valid response instead of surfacing the error. + schemaErr := fmt.Errorf("%w: - groups: Array must have at most 2 items", ai.ErrSchemaValidation) + inner := &erroringProvider{outcomes: []outcome{{err: schemaErr}, {text: validJSON}}} + p, err := WithSchemaValidation()(inner) + if err != nil { + t.Fatalf("WithSchemaValidation: %v", err) + } + resp, err := p.Execute(context.Background(), capRequest(api.SchemaStrictnessRetry)) + if err != nil { + t.Fatalf("provider schema error should recover on retry, got %v", err) + } + if resp.Text != validJSON { + t.Errorf("want the corrected response, got %q", resp.Text) + } + if inner.calls != 2 { + t.Errorf("want 1 initial + 1 retry = 2 calls, got %d", inner.calls) + } + // The re-ask must feed the schema errors back so the model can correct itself. + if len(inner.prompts) < 2 || !strings.Contains(inner.prompts[1], "Array must have at most 2 items") { + t.Errorf("retry prompt must include the validation errors, got %v", inner.prompts) + } +} + +func TestValidation_ProviderSchemaErrorPassesThroughWithoutRetryStrictness(t *testing.T) { + // Without retry strictness a provider schema error surfaces as-is; no re-ask. + schemaErr := fmt.Errorf("%w: bad", ai.ErrSchemaValidation) + inner := &erroringProvider{outcomes: []outcome{{err: schemaErr}}} + p, err := WithSchemaValidation()(inner) + if err != nil { + t.Fatalf("WithSchemaValidation: %v", err) + } + _, err = p.Execute(context.Background(), capRequest(api.SchemaStrictnessNone)) + if !errors.Is(err, ai.ErrSchemaValidation) { + t.Fatalf("want the provider error surfaced, got %v", err) + } + if inner.calls != 1 { + t.Errorf("none strictness must not retry; want 1 call, got %d", inner.calls) + } +} + +func TestValidation_ProviderSchemaErrorExhaustsRetries(t *testing.T) { + // A provider that keeps rejecting during generation exhausts the retry budget + // and fails hard with ErrSchemaValidation. + schemaErr := fmt.Errorf("%w: - groups: too many", ai.ErrSchemaValidation) + inner := &erroringProvider{outcomes: []outcome{{err: schemaErr}}} // repeats + p, err := WithSchemaValidation()(inner) + if err != nil { + t.Fatalf("WithSchemaValidation: %v", err) + } + _, err = p.Execute(context.Background(), capRequest(api.SchemaStrictnessRetry)) + if !errors.Is(err, ai.ErrSchemaValidation) { + t.Fatalf("want ErrSchemaValidation after exhausting retries, got %v", err) + } + if want := 1 + maxSchemaRetries; inner.calls != want { + t.Errorf("want %d calls, got %d", want, inner.calls) + } +} + +func TestValidation_RepairRequestUsesConfiguredModel(t *testing.T) { + v := &validatingProvider{ + provider: &scriptedProvider{backend: ai.BackendAnthropic}, + cfg: ai.Config{SchemaRepair: api.SchemaRepairConfig{ + Model: api.Model{Name: "gpt-5", Backend: ai.BackendOpenAI}, + }}, + } + req, cfg, useParent, err := v.repairRequest(capRequest(api.SchemaStrictnessRetry), json.RawMessage(capSchema), "bad", &ai.Response{Text: overCapJSON}, 1) + if err != nil { + t.Fatalf("repairRequest: %v", err) + } + if useParent { + t.Fatal("repair with configured model should not reuse parent provider") + } + if cfg.Model.Name != "gpt-5" || cfg.Model.Backend != ai.BackendOpenAI { + t.Fatalf("repair cfg model = %#v", cfg.Model) + } + if req.Prompt.SchemaStrictness != api.SchemaStrictnessDisabled { + t.Fatalf("repair prompt strictness = %q, want explicit none", req.Prompt.SchemaStrictness) + } + if !strings.Contains(req.Prompt.User, "Array must") && !strings.Contains(req.Prompt.User, "bad") { + t.Fatalf("repair prompt should include validation errors, got %q", req.Prompt.User) + } +} + +func TestValidation_RepairRequestUsesConfiguredPromptFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "repair.prompt") + if err := os.WriteFile(path, []byte(`{{role "user"}} +Fix this: {{validationErrors}} +Previous: {{previousResponse}} +`), 0o644); err != nil { + t.Fatalf("write repair prompt: %v", err) + } + v := &validatingProvider{ + provider: &scriptedProvider{backend: ai.BackendAnthropic}, + cfg: ai.Config{SchemaRepair: api.SchemaRepairConfig{Prompt: path}}, + } + req, _, useParent, err := v.repairRequest(capRequest(api.SchemaStrictnessRetry), json.RawMessage(capSchema), "custom-error", &ai.Response{Text: overCapJSON}, 1) + if err != nil { + t.Fatalf("repairRequest: %v", err) + } + if !useParent { + t.Fatal("custom prompt without model override should reuse parent provider") + } + if !strings.Contains(req.Prompt.User, "custom-error") || !strings.Contains(req.Prompt.User, overCapJSON) { + t.Fatalf("configured repair prompt did not render variables: %q", req.Prompt.User) + } + if req.Prompt.Source != path { + t.Fatalf("repair prompt source = %q, want %q", req.Prompt.Source, path) } } diff --git a/pkg/ai/model_effort.go b/pkg/ai/model_effort.go new file mode 100644 index 0000000..a7ad2d0 --- /dev/null +++ b/pkg/ai/model_effort.go @@ -0,0 +1,107 @@ +package ai + +import ( + "context" + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/commons/logger" +) + +var modelEffortLog = logger.GetLogger("ai") + +// ModelEfforts returns model-specific effort metadata when the embedded +// registry knows the exact backend/model combination. +func ModelEfforts(backend Backend, model string) (supported []api.Effort, defaultEffort api.Effort, ok bool) { + def, ok := RegistryModelDef(backend, model) + if !ok { + return nil, api.EffortNone, false + } + return append([]api.Effort(nil), def.SupportedEfforts...), def.DefaultEffort, true +} + +// ValidateModelEffort enforces model-aware effort levels without making the +// catalog exhaustive. Unknown models retain Captain's historical suggest-only +// behavior; known registry entries accept only source-backed executable tiers. +func ValidateModelEffort(backend Backend, model string, effort api.Effort) error { + if err := effort.Validate(); err != nil { + return err + } + if effort == api.EffortNone { + return nil + } + supported, _, known := ModelEfforts(backend, model) + if !known { + return nil + } + if len(supported) == 0 { + return fmt.Errorf("model %q on %s does not support a reasoning effort", model, backend) + } + for _, candidate := range supported { + if candidate == effort { + return nil + } + } + values := make([]string, 0, len(supported)) + for _, candidate := range supported { + values = append(values, string(candidate)) + } + return fmt.Errorf("model %q on %s does not support reasoning effort %q; want one of: %s", + model, backend, effort, strings.Join(values, ", ")) +} + +type effortValidatingProvider struct { + provider Provider + configuredEffort api.Effort +} + +func (p *effortValidatingProvider) GetModel() string { return p.provider.GetModel() } +func (p *effortValidatingProvider) GetBackend() Backend { return p.provider.GetBackend() } +func (p *effortValidatingProvider) Unwrap() Provider { return p.provider } +func (p *effortValidatingProvider) request(ctx context.Context, req Request) (Request, error) { + if req.Effort == api.EffortNone { + req.Effort = p.configuredEffort + } + if supported, _, known := ModelEfforts(p.GetBackend(), p.GetModel()); known && len(supported) == 0 && req.Effort != api.EffortNone { + LoggerFromContext(ctx, modelEffortLog).Warnf( + "model %q on %s does not support reasoning effort %q; continuing without effort", + p.GetModel(), p.GetBackend(), req.Effort, + ) + req.Effort = api.EffortNone + return req, nil + } + if err := ValidateModelEffort(p.GetBackend(), p.GetModel(), req.Effort); err != nil { + return Request{}, err + } + return req, nil +} + +func (p *effortValidatingProvider) Execute(ctx context.Context, req Request) (*Response, error) { + req, err := p.request(ctx, req) + if err != nil { + return nil, err + } + return p.provider.Execute(ctx, req) +} + +type effortValidatingStreamingProvider struct { + *effortValidatingProvider + streamer StreamingProvider +} + +func (p *effortValidatingStreamingProvider) ExecuteStream(ctx context.Context, req Request) (<-chan Event, error) { + req, err := p.request(ctx, req) + if err != nil { + return nil, err + } + return p.streamer.ExecuteStream(ctx, req) +} + +func withEffortValidation(provider Provider, configured api.Effort) Provider { + base := &effortValidatingProvider{provider: provider, configuredEffort: configured} + if streamer, ok := provider.(StreamingProvider); ok { + return &effortValidatingStreamingProvider{effortValidatingProvider: base, streamer: streamer} + } + return base +} diff --git a/pkg/ai/model_effort_ginkgo_test.go b/pkg/ai/model_effort_ginkgo_test.go new file mode 100644 index 0000000..7e41561 --- /dev/null +++ b/pkg/ai/model_effort_ginkgo_test.go @@ -0,0 +1,45 @@ +package ai + +import ( + "context" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/api" +) + +type effortCaptureProvider struct { + backend Backend + model string + request Request +} + +func (p *effortCaptureProvider) Execute(_ context.Context, req Request) (*Response, error) { + p.request = req + return &Response{}, nil +} + +func (p *effortCaptureProvider) GetModel() string { return p.model } +func (p *effortCaptureProvider) GetBackend() Backend { return p.backend } + +var _ = Describe("model effort runtime validation", func() { + It("drops configured effort when the model has no reasoning support", func() { + capture := &effortCaptureProvider{backend: BackendGemini, model: "gemini-2.5-flash"} + provider := withEffortValidation(capture, api.EffortHigh) + + _, err := provider.Execute(context.Background(), Request{}) + + Expect(err).NotTo(HaveOccurred()) + Expect(capture.request.Effort).To(Equal(api.EffortNone)) + }) + + It("still rejects an unsupported tier on a reasoning model", func() { + capture := &effortCaptureProvider{backend: BackendOpenAI, model: "gpt-5.5"} + provider := withEffortValidation(capture, api.EffortMax) + + _, err := provider.Execute(context.Background(), Request{}) + + Expect(err).To(MatchError(ContainSubstring("does not support reasoning effort"))) + }) +}) diff --git a/pkg/ai/model_effort_test.go b/pkg/ai/model_effort_test.go new file mode 100644 index 0000000..38608c9 --- /dev/null +++ b/pkg/ai/model_effort_test.go @@ -0,0 +1,55 @@ +package ai + +import ( + "strings" + "testing" + + "github.com/flanksource/captain/pkg/api" +) + +func TestValidateModelEffort(t *testing.T) { + tests := []struct { + name string + backend Backend + model string + effort api.Effort + wantErr string + }{ + {"api max", BackendOpenAI, "gpt-5.6", api.EffortMax, ""}, + {"api rejects ultra", BackendOpenAI, "gpt-5.6", api.EffortUltra, "does not support"}, + {"sol max", BackendCodexAgent, "sol", api.EffortMax, ""}, + {"sol ultra", BackendCodexAgent, "sol", api.EffortUltra, ""}, + {"terra ultra", BackendCodexCmux, "gpt-5.6-terra", api.EffortUltra, ""}, + {"luna supports max", BackendCodexCLI, "luna", api.EffortMax, ""}, + {"gpt55 rejects max", BackendCodexAgent, "gpt-5.5", api.EffortMax, "does not support"}, + {"unknown codex remains open", BackendCodexAgent, "gpt-future", api.EffortUltra, ""}, + {"adaptive Anthropic supports max", BackendAnthropic, "claude-sonnet-5", api.EffortMax, ""}, + {"DeepSeek rejects configured effort", BackendDeepSeek, "deepseek-v4-pro", api.EffortHigh, "does not support a reasoning effort"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateModelEffort(tt.backend, tt.model, tt.effort) + if tt.wantErr == "" && err != nil { + t.Fatalf("ValidateModelEffort: %v", err) + } + if tt.wantErr != "" && (err == nil || !strings.Contains(err.Error(), tt.wantErr)) { + t.Fatalf("error = %v, want %q", err, tt.wantErr) + } + }) + } +} + +func TestRegistryGPT56Availability(t *testing.T) { + if _, ok := RegistryModelDef(BackendOpenAI, "gpt-5.6"); !ok { + t.Fatal("gpt-5.6 should be available to OpenAI") + } + if _, ok := RegistryModelDef(BackendCodexAgent, "gpt-5.6"); ok { + t.Fatal("gpt-5.6 API base should not be available to Codex") + } + if _, ok := RegistryModelDef(BackendCodexAgent, "sol"); !ok { + t.Fatal("Sol should be available to Codex") + } + if _, ok := RegistryModelDef(BackendOpenAI, "sol"); !ok { + t.Fatal("Sol should be available to OpenAI") + } +} diff --git a/pkg/ai/model_error.go b/pkg/ai/model_error.go new file mode 100644 index 0000000..3ba67e7 --- /dev/null +++ b/pkg/ai/model_error.go @@ -0,0 +1,154 @@ +package ai + +import ( + "context" + "fmt" + "os/exec" + "strings" + "sync" + "time" + + "github.com/flanksource/captain/pkg/collections" +) + +const modelRecommendationLimit = 5 +const modelRecommendationTimeout = 2 * time.Second + +// modelAvailabilityResolver is replaceable in tests so model-error handling +// never needs a real provider account or installed CLI. +var modelAvailabilityResolver = availableModelsForBackend + +// availableModelsForBackend returns models executable by the selected backend. +// Codex's installed catalog is the strongest source for local Codex adapters; +// direct API backends use their authenticated model endpoint. The embedded +// registry keeps recommendations useful when discovery is unavailable. +func availableModelsForBackend(ctx context.Context, backend Backend) []ModelDef { + discoveryCtx, cancel := context.WithTimeout(ctx, modelRecommendationTimeout) + defer cancel() + if isCodexBackend(backend) { + if binary, err := exec.LookPath("codex"); err == nil { + if models, err := FetchCodexDebugModels(discoveryCtx, binary); err == nil && len(models) > 0 { + for i := range models { + models[i].Backend = backend + } + return CurrentCuratedModelsByReleaseDate(models) + } + } + } + if backend.Kind() == "api" && GetAPIKeyFromEnv(backend) != "" { + if models, err := ListModels(discoveryCtx, backend); err == nil && len(models) > 0 { + return CurrentModelsByReleaseDate(models) + } + } + return CurrentCuratedModelsByReleaseDate(RegistryModelDefs(backend)) +} + +func isCodexBackend(backend Backend) bool { + switch backend { + case BackendCodexCLI, BackendCodexAgent, BackendCodexCmux: + return true + default: + return false + } +} + +// recommendModelError preserves err while adding a backend-scoped replacement +// and a compact available-model list. It runs only after a provider has +// confirmed that the attempted model is unavailable. +func recommendModelError(ctx context.Context, backend Backend, attempted string, err error) error { + if !IsModelUnavailable(err) || strings.Contains(strings.ToLower(err.Error()), "available models for") { + return err + } + return recommendModelErrorFromModels(backend, attempted, err, modelAvailabilityResolver(ctx, backend)) +} + +func recommendModelErrorFromModels(backend Backend, attempted string, err error, models []ModelDef) error { + available := make([]string, 0, len(models)) + seen := map[string]bool{} + for _, model := range models { + id := strings.TrimSpace(model.ID) + if id == "" || strings.EqualFold(id, attempted) || seen[id] { + continue + } + seen[id] = true + available = append(available, id) + } + if len(available) == 0 { + return err + } + + closest := collections.FindSimilar(attempted, available, 1)[0] + shown := available + more := 0 + if len(shown) > modelRecommendationLimit { + more = len(shown) - modelRecommendationLimit + shown = shown[:modelRecommendationLimit] + } + list := strings.Join(shown, ", ") + if more > 0 { + list += fmt.Sprintf(" (+%d more)", more) + } + return fmt.Errorf("%w; did you mean %q? available models for %s: %s", err, closest, backend, list) +} + +type modelErrorProvider struct { + provider Provider + + modelsOnce sync.Once + models []ModelDef +} + +func (p *modelErrorProvider) GetModel() string { return p.provider.GetModel() } +func (p *modelErrorProvider) GetBackend() Backend { return p.provider.GetBackend() } +func (p *modelErrorProvider) Unwrap() Provider { return p.provider } + +func (p *modelErrorProvider) recommend(ctx context.Context, err error) error { + if !IsModelUnavailable(err) || strings.Contains(strings.ToLower(err.Error()), "available models for") { + return err + } + p.modelsOnce.Do(func() { + p.models = modelAvailabilityResolver(ctx, p.GetBackend()) + }) + return recommendModelErrorFromModels(p.GetBackend(), p.GetModel(), err, p.models) +} + +func (p *modelErrorProvider) Execute(ctx context.Context, req Request) (*Response, error) { + resp, err := p.provider.Execute(ctx, req) + if err != nil { + err = p.recommend(ctx, err) + } + return resp, err +} + +type modelErrorStreamingProvider struct { + *modelErrorProvider + streamer StreamingProvider +} + +func (p *modelErrorStreamingProvider) ExecuteStream(ctx context.Context, req Request) (<-chan Event, error) { + ch, err := p.streamer.ExecuteStream(ctx, req) + if err != nil { + return nil, p.recommend(ctx, err) + } + out := make(chan Event) + go func() { + defer close(out) + for ev := range ch { + if ev.Kind == EventError && ev.Error != "" { + ev.Error = p.recommend(ctx, fmt.Errorf("%s", ev.Error)).Error() + } + if !sendEvent(ctx, out, ev) { + return + } + } + }() + return out, nil +} + +func withModelErrorRecommendations(provider Provider) Provider { + base := &modelErrorProvider{provider: provider} + if streamer, ok := provider.(StreamingProvider); ok { + return &modelErrorStreamingProvider{modelErrorProvider: base, streamer: streamer} + } + return base +} diff --git a/pkg/ai/model_error_test.go b/pkg/ai/model_error_test.go new file mode 100644 index 0000000..0d3b7a8 --- /dev/null +++ b/pkg/ai/model_error_test.go @@ -0,0 +1,115 @@ +package ai + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/api" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type modelErrorTestProvider struct { + backend Backend + model string + err error + events []Event +} + +func (p *modelErrorTestProvider) GetModel() string { return p.model } +func (p *modelErrorTestProvider) GetBackend() Backend { return p.backend } +func (p *modelErrorTestProvider) Execute(context.Context, Request) (*Response, error) { + return nil, p.err +} +func (p *modelErrorTestProvider) ExecuteStream(context.Context, Request) (<-chan Event, error) { + if p.err != nil { + return nil, p.err + } + ch := make(chan Event, len(p.events)) + for _, ev := range p.events { + ch <- ev + } + close(ch) + return ch, nil +} + +func stubModelRecommendations(t *testing.T) { + t.Helper() + previous := modelAvailabilityResolver + modelAvailabilityResolver = func(context.Context, Backend) []ModelDef { + return []ModelDef{ + {ID: "gpt-5.6-sol"}, + {ID: "gpt-5.6-terra"}, + {ID: "gpt-5.6-luna"}, + {ID: "gpt-5.5"}, + {ID: "gpt-5.4"}, + {ID: "gpt-5.4-mini"}, + } + } + t.Cleanup(func() { modelAvailabilityResolver = previous }) +} + +func TestRecommendModelErrorUsesAvailableBackendModels(t *testing.T) { + stubModelRecommendations(t) + base := fmtModelUnavailable("The 'gpt-5.6-line' model is not supported when using Codex with a ChatGPT account.") + err := recommendModelError(context.Background(), BackendCodexAgent, "gpt-5.6-line", base) + + require.Error(t, err) + assert.Contains(t, err.Error(), `did you mean "gpt-5.6-luna"?`) + assert.Contains(t, err.Error(), "available models for codex-agent: gpt-5.6-sol, gpt-5.6-terra, gpt-5.6-luna, gpt-5.5, gpt-5.4 (+1 more)") + assert.ErrorIs(t, err, ErrModelUnavailable) +} + +func TestModelErrorProviderEnrichesBufferedFailure(t *testing.T) { + stubModelRecommendations(t) + base := &modelErrorTestProvider{ + backend: BackendCodexCLI, + model: "gpt-5.6-line", + err: errors.New("unsupported model gpt-5.6-line"), + } + p := withModelErrorRecommendations(base) + _, err := p.Execute(context.Background(), api.Spec{}) + require.Error(t, err) + assert.Contains(t, err.Error(), `did you mean "gpt-5.6-luna"?`) +} + +func TestModelErrorProviderEnrichesStreamingFailure(t *testing.T) { + previous := modelAvailabilityResolver + resolveCalls := 0 + modelAvailabilityResolver = func(context.Context, Backend) []ModelDef { + resolveCalls++ + return []ModelDef{{ID: "gpt-5.6-luna"}, {ID: "gpt-5.6-sol"}} + } + t.Cleanup(func() { modelAvailabilityResolver = previous }) + base := &modelErrorTestProvider{ + backend: BackendCodexAgent, + model: "gpt-5.6-line", + events: []Event{ + {Kind: EventError, Error: "unknown model gpt-5.6-line"}, + {Kind: EventError, Error: "unknown model gpt-5.6-line"}, + }, + } + p := withModelErrorRecommendations(base).(StreamingProvider) + ch, err := p.ExecuteStream(context.Background(), api.Spec{}) + require.NoError(t, err) + events := drain(ch) + require.Len(t, events, 2) + assert.Contains(t, events[0].Error, `did you mean "gpt-5.6-luna"?`) + assert.Contains(t, events[1].Error, `did you mean "gpt-5.6-luna"?`) + assert.Equal(t, 1, resolveCalls, "repeated terminal events should reuse one availability snapshot") +} + +func TestRecommendModelErrorLeavesOtherFailuresUnchanged(t *testing.T) { + stubModelRecommendations(t) + base := errors.New("authentication failed: invalid API key") + err := recommendModelError(context.Background(), BackendOpenAI, "gpt-5.6", base) + assert.Same(t, base, err) + assert.False(t, strings.Contains(err.Error(), "available models")) +} + +func fmtModelUnavailable(message string) error { + return fmt.Errorf("%w: %s", ErrModelUnavailable, message) +} diff --git a/pkg/ai/model_lists.go b/pkg/ai/model_lists.go index 62d6330..b62cdbd 100644 --- a/pkg/ai/model_lists.go +++ b/pkg/ai/model_lists.go @@ -4,12 +4,14 @@ import ( "sort" "strings" "time" + + "github.com/flanksource/captain/pkg/api/registry" ) const currentModelsPerFamily = 3 // legacyModelPrefixes hides model IDs that are either superseded by a newer -// generation or aren't chat completions (image/audio/embedding/moderation). +// generation or aren't primary text models. var legacyModelPrefixes = []string{ // OpenAI legacy "gpt-3", @@ -17,9 +19,16 @@ var legacyModelPrefixes = []string{ "gpt-5-", // API variants like mini/nano/codex/pro; CLI Codex is exempt by backend "o1", "o3", + "o4", "codex-mini", - // OpenAI non-chat endpoints + // OpenAI non-primary endpoints and aliases + "gpt-realtime", + "gpt-image", + "gpt-audio", + "sora", "dall-", + "image-", + "audio-", "whisper", "tts-", "text-embedding", @@ -27,6 +36,7 @@ var legacyModelPrefixes = []string{ "omni-moderation", "babbage", "davinci", + "chat-latest", "chatgpt-", "computer-use-preview", // Claude legacy @@ -37,6 +47,10 @@ var legacyModelPrefixes = []string{ "claude-sonnet-4-2", "claude-opus-4-0", "claude-opus-4-1", + "fable-", + "opus-", + "sonnet-", + "haiku-", // Gemini legacy "gemini-1", "gemini-2.0", @@ -49,31 +63,103 @@ var legacyModelPrefixes = []string{ // model id. Call IsLegacyModelIDForBackend when backend context is available. func IsLegacyModelID(id string) bool { idLower := strings.ToLower(bareModelID(strings.TrimPrefix(strings.TrimSpace(id), "models/"))) + if IsIgnoredOpenAIModelID(idLower) { + return true + } + for _, p := range legacyModelPrefixes { + if strings.HasPrefix(idLower, p) { + return true + } + } + return false +} + +// IsIgnoredOpenAIModelID reports whether an OpenAI model-list id should be +// hidden from ordinary model pickers. OpenAI exposes many non-primary surfaces +// (realtime, audio, image, Sora, code/chat aliases, dated and size variants); +// Captain's picker keeps stable primary GPT text ids and explicitly registered +// Codex runtime variants. Use this before remapping live OpenAI ids onto Codex +// backends. +func IsIgnoredOpenAIModelID(id string) bool { + idLower := strings.ToLower(bareModelID(strings.TrimPrefix(strings.TrimSpace(id), "models/"))) + if idLower == "" { + return false + } for _, p := range legacyModelPrefixes { if strings.HasPrefix(idLower, p) { return true } } + if strings.HasPrefix(idLower, "gpt-") { + return !isPrimaryGPTModelID(idLower) + } + for _, needle := range []string{"realtime", "image", "audio", "whisper", "tts", "sora", "codex", "code", "chat-latest", "chatgpt", "computer-use"} { + if strings.Contains(idLower, needle) { + return true + } + } return false } -// IsLegacyModelIDForBackend keeps API model menus clean while preserving local -// agent model slugs such as gpt-5-codex, which are current for Codex CLI even -// though the same id would be noisy in an OpenAI API model listing. +func isPrimaryGPTModelID(id string) bool { + version := strings.TrimPrefix(id, "gpt-") + if version == "" || version == id { + return false + } + sawDigit := false + for _, r := range version { + switch { + case r >= '0' && r <= '9': + sawDigit = true + case r == '.': + continue + default: + return false + } + } + return sawDigit +} + +// IsLegacyModelIDForBackend keeps model menus clean while retaining preferred +// registry models explicitly available to the selected backend. func IsLegacyModelIDForBackend(id string, backend Backend) bool { - if backend.Kind() == "cli" { + if isPreferredRegistryModelForBackend(backend, id) { return false } return IsLegacyModelID(id) } +func isPreferredRegistryModelForBackend(backend Backend, model string) bool { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return false + } + entry, found := p.Lookup(model) + if !found { + return false + } + _, available := p.Availability(mode, entry.ID) + return entry.Preferred && available +} + // CurrentModelsByReleaseDate returns a filtered copy sorted newest first, // retaining the newest few models per family prefix. Known catalog release // dates fill gaps left by provider list endpoints. func CurrentModelsByReleaseDate(models []ModelDef) []ModelDef { + return currentModelsByReleaseDate(models, true) +} + +// CurrentCuratedModelsByReleaseDate sorts and limits a trusted runtime catalog +// whose own visibility field has already removed hidden models. It deliberately +// skips the generic OpenAI variant blacklist used for raw provider listings. +func CurrentCuratedModelsByReleaseDate(models []ModelDef) []ModelDef { + return currentModelsByReleaseDate(models, false) +} + +func currentModelsByReleaseDate(models []ModelDef, filterLegacy bool) []ModelDef { out := make([]ModelDef, 0, len(models)) for _, m := range models { - if IsLegacyModelIDForBackend(m.ID, m.Backend) { + if filterLegacy && IsLegacyModelIDForBackend(m.ID, m.Backend) { continue } if m.ReleaseDate == "" { @@ -89,6 +175,15 @@ func CurrentModelsByReleaseDate(models []ModelDef) []ModelDef { // unknown dates last and id descending as the stable deterministic tie-breaker. func SortModelsByReleaseDateDesc(models []ModelDef) { sort.SliceStable(models, func(i, j int) bool { + if models[i].Priority != models[j].Priority && (models[i].Priority > 0 || models[j].Priority > 0) { + if models[i].Priority == 0 { + return false + } + if models[j].Priority == 0 { + return true + } + return models[i].Priority < models[j].Priority + } if ModelFamilyPrefix(models[i].ID) == ModelFamilyPrefix(models[j].ID) { if cmp := compareModelVersions(models[i].ID, models[j].ID); cmp != 0 { return cmp > 0 @@ -151,6 +246,8 @@ func ModelFamilyPrefix(id string) string { return strings.Join(parts[:3], "-") } return "claude-" + parts[1] + case "fable", "opus", "sonnet", "haiku": + return parts[0] case "gemini": for i := 1; i < len(parts); i++ { if isModelVersionToken(parts[i]) { diff --git a/pkg/ai/model_parse_conformance_test.go b/pkg/ai/model_parse_conformance_test.go new file mode 100644 index 0000000..a08098d --- /dev/null +++ b/pkg/ai/model_parse_conformance_test.go @@ -0,0 +1,184 @@ +package ai + +import ( + "fmt" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/api" +) + +// The parse conformance suite is the single behavioural contract for turning a +// user-written model string into a concrete model/backend/effort. It pins the +// two entry points that reach the parser: +// +// - specPath — `captain ai --model X` and spec/`.captain.yaml` decoding: +// api.Model.Expand (the compact grammar) then ResolveModelSelectors. +// See pkg/cli/provider_defaults.go and pkg/cli/ai.go. +// - frontmatterPath — prompt frontmatter: api.Model.ExpandCSV (no colon +// grammar at all) then ResolveModelSelectors. +// See pkg/cli/prompt_render.go. +// +// These two ran different parsers over the same syntax and disagreed about which +// models existed. They now share one parser (pkg/api/registry), so every case +// below asserts both paths agree — that agreement is the contract. + +func specPath(s string) (api.Model, error) { + m, err := api.Model{Name: s}.Expand() + if err != nil { + return api.Model{}, err + } + return ResolveModelSelectors(m) +} + +func frontmatterPath(s string) (api.Model, error) { + return ResolveModelSelectors(api.Model{Name: s}.ExpandCSV()) +} + +// parseCase is one model string and the model it must resolve to, on both entry +// points. wantErr instead asserts both fail with a message containing it. +type parseCase struct { + in string + want api.Model + wantErr string +} + +func model(name string, backend api.Backend, effort api.Effort) api.Model { + return api.Model{Name: name, Backend: backend, Effort: effort} +} + +// agreedCases are inputs both entry points already resolve identically. These +// are pure regression pins: the unified parser must not change any of them. +var agreedCases = []parseCase{ + // Family aliases resolve to the exact current registry model. + {in: "sonnet", want: model("claude-sonnet-5", api.BackendAnthropic, "")}, + {in: "opus", want: model("claude-opus-5", api.BackendAnthropic, "")}, + {in: "fable", want: model("claude-fable-5", api.BackendAnthropic, "")}, + + // A superseded exact id is rewritten to its successor. + {in: "claude-sonnet-4-5", want: model("claude-sonnet-4-6", api.BackendAnthropic, "")}, + + // The catalog prefix is stripped; "models/" likewise. + {in: "anthropic/claude-sonnet-5", want: model("claude-sonnet-5", api.BackendAnthropic, "")}, + {in: "googleai/gemini-3.5-flash", want: model("gemini-3.5-flash", api.BackendGemini, "")}, + {in: "models/gemini-3.5-flash", want: model("gemini-3.5-flash", api.BackendGemini, "")}, + + // Mode prefix and effort suffix. + {in: "cli:sonnet:high", want: model("claude-sonnet-5", api.BackendClaudeCLI, api.EffortHigh)}, + {in: "agent:sonnet", want: model("claude-sonnet-5", api.BackendClaudeAgent, "")}, + + // Codex codenames are catalog aliases, so they resolve wherever a model name + // is accepted. `--model agent:sol` used to error while the identical value in + // prompt frontmatter ran: the compact grammar could not see pkg/ai's alias + // table. This is the divergence the single parser exists to remove. + {in: "agent:sol", want: model("gpt-5.6-sol", api.BackendCodexAgent, "")}, + {in: "agent:sol:high", want: model("gpt-5.6-sol", api.BackendCodexAgent, api.EffortHigh)}, + {in: "api:sol", want: model("gpt-5.6-sol", api.BackendOpenAI, "")}, + // A bare codename resolves too. It used to fail on both paths only because + // the bare path went through an alias-blind claim — an accident of which + // parser ran, not a decision. + {in: "sol", want: model("gpt-5.6-sol", api.BackendOpenAI, "")}, + {in: "terra", want: model("gpt-5.6-terra", api.BackendOpenAI, "")}, + {in: "luna", want: model("gpt-5.6-luna", api.BackendOpenAI, "")}, + + // Bare CLI sentinels are asymmetric and must stay that way: "codex" resolves + // to the latest codex model on the CLI, "claude" stays a literal sentinel on + // the API backend. + {in: "codex", want: model("gpt-5.6-sol", api.BackendCodexCLI, "")}, + {in: "claude", want: model("claude", api.BackendAnthropic, "")}, + + // A sentinel with an explicit mode resolves too. This does NOT go through the + // agent-sentinel shortcut (that one only fires off the API mode), so it lands + // on the provider's emptyFamily — which must be a family name. An id there + // matched no catalog row and left "api:codex" as the literal "codex". + {in: "api:codex", want: model("gpt-5.6-sol", api.BackendOpenAI, "")}, + {in: "cli:codex", want: model("gpt-5.6-sol", api.BackendCodexCLI, "")}, + + // A multi-slash id resolves off its LAST segment and keeps its name verbatim, + // so OpenRouter-style proxied names survive. (api.InferBackend alone cannot do + // this; inferModelBackend's last-slash retry is what makes it work.) + {in: "openrouter/anthropic/claude-x", want: model("openrouter/anthropic/claude-x", api.BackendAnthropic, "")}, + + {in: "gemini-3.5-flash", want: model("gemini-3.5-flash", api.BackendGemini, "")}, + {in: "deepseek-chat", want: model("deepseek-chat", api.BackendDeepSeek, "")}, + {in: "gpt-5.6", want: model("gpt-5.6", api.BackendOpenAI, "")}, + {in: "o3", want: model("o3", api.BackendOpenAI, "")}, + + // Unknown models fail loud on both paths, telling the user how to recover. + {in: "totally-unknown", wantErr: "pass an explicit backend"}, + + // sora is a video model captain cannot run. It is claimed by no provider, so + // it fails loud rather than resolving to something unusable. + {in: "sora", wantErr: "pass an explicit backend"}, + {in: "sora-2", wantErr: "pass an explicit backend"}, + + // grok mode was removed from the codex CLI, so no provider claims it. Pinned + // as failing rather than deleted: silently re-claiming grok would route models + // to a backend captain no longer drives. + {in: "grok-2", wantErr: "pass an explicit backend"}, + {in: "grok-code-fast-1", wantErr: "pass an explicit backend"}, +} + +func expectModel(got api.Model, want api.Model) { + GinkgoHelper() + Expect(got.Name).To(Equal(want.Name), "model name") + Expect(got.Backend).To(Equal(want.Backend), "backend") + Expect(got.Effort).To(Equal(want.Effort), "effort") +} + +var _ = Describe("model parse conformance", func() { + Describe("inputs both entry points agree on", func() { + for _, tc := range agreedCases { + tc := tc + label := fmt.Sprintf("%q", tc.in) + if tc.wantErr != "" { + It(label+" fails loud on both paths", func() { + _, specErr := specPath(tc.in) + _, fmErr := frontmatterPath(tc.in) + Expect(specErr).To(MatchError(ContainSubstring(tc.wantErr))) + Expect(fmErr).To(MatchError(ContainSubstring(tc.wantErr))) + }) + continue + } + It(label+" resolves identically on both paths", func() { + got, err := specPath(tc.in) + Expect(err).NotTo(HaveOccurred()) + expectModel(got, tc.want) + + got, err = frontmatterPath(tc.in) + Expect(err).NotTo(HaveOccurred()) + expectModel(got, tc.want) + }) + } + }) + + Describe("comma-separated fallback chains", func() { + It("keeps the primary first and resolves each element independently", func() { + got, err := specPath("sonnet,cli:opus:high") + Expect(err).NotTo(HaveOccurred()) + expectModel(got, model("claude-sonnet-5", api.BackendAnthropic, "")) + Expect(got.Fallbacks).To(HaveLen(1)) + expectModel(got.Fallbacks[0], model("claude-opus-5", api.BackendClaudeCLI, api.EffortHigh)) + }) + }) + + Describe("wildcard selectors", func() { + It("is rejected for a single model", func() { + _, err := frontmatterPath("*:fable") + Expect(err).To(MatchError(ContainSubstring("only valid for --multi-models"))) + }) + + It("fans out to every backend of the claimed family", func() { + models, err := ResolveRuntimeSelectors([]string{"*:fable"}, api.Model{Name: "sonnet"}) + Expect(err).NotTo(HaveOccurred()) + backends := make([]api.Backend, 0, len(models)) + for _, m := range models { + Expect(m.Name).To(Equal("claude-fable-5")) + backends = append(backends, m.Backend) + } + Expect(backends).To(ContainElement(api.BackendAnthropic)) + Expect(backends).To(ContainElement(api.BackendClaudeAgent)) + }) + }) +}) diff --git a/pkg/ai/model_registry.go b/pkg/ai/model_registry.go new file mode 100644 index 0000000..d3ce302 --- /dev/null +++ b/pkg/ai/model_registry.go @@ -0,0 +1,206 @@ +package ai + +import ( + "strings" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" +) + +// Model identity — claiming a name for a provider, resolving aliases and +// superseded ids, and picking the exact catalog id — lives in pkg/api/registry. +// This file is only the projection of those rows onto pkg/ai's catalog types. +// +// It used to own a second copy of that knowledge (splitModelProvider, +// ParseModelIdentity, normalizeCodexVariantAlias, isSupersededRegistryExact), +// which disagreed with pkg/api's InferBackend about grok, sora, and codenames. + +// ModelIdentity is captain's parsed model key. +type ModelIdentity = registry.ModelIdentity + +// defaultCatalogModelID is the bare id behind DefaultModelID — the one catalog +// row that carries Default: true. +var defaultCatalogModelID = registry.StripProviderPrefix(DefaultModelID) + +// ParseModelIdentity parses a model token into its provider/family/version tuple. +// The token's own provider wins; defaultProvider only decides tokens that claim +// no family of their own. +func ParseModelIdentity(defaultProvider, model string) (ModelIdentity, bool) { + p, token, _, ok := registry.ProviderForToken(model) + if !ok { + if p, ok = registry.ProviderByName(defaultProvider); !ok { + return ModelIdentity{}, false + } + token = model + } + return p.ParseIdentity(token) +} + +func registryModelDef(m registry.KnownModel, backend Backend) ModelDef { + return ModelDef{ + ID: m.ID, + Name: m.Label, + Backend: backend, + ReleaseDate: m.ReleaseDate, + CapabilitiesKnown: true, + Reasoning: m.Reasoning, + Temperature: m.Temperature, + InputMediaTypes: clampInputMediaTypes(backend, m.InputMediaTypes), + SupportedEfforts: append([]api.Effort(nil), m.SupportedEfforts...), + DefaultEffort: m.DefaultEffort, + Priority: m.Priority, + } +} + +// RegistryModelDef returns the registry metadata for an exact model on a +// backend. The boolean is false when the model is known but unavailable there. +// +// The lookup is exact (aliases aside) and deliberately does NOT resolve version +// lines: "gpt-5.6" is an API-only base model, and resolving it here would answer +// with its codex-available sibling gpt-5.6-sol and report the base as available +// on Codex. Callers that want resolution call ResolveExactModelForBackend first. +func RegistryModelDef(backend Backend, model string) (ModelDef, bool) { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return ModelDef{}, false + } + entry, found := p.Lookup(model) + if !found { + return ModelDef{}, false + } + if _, available := p.Availability(mode, entry.ID); !available { + return ModelDef{}, false + } + return registryModelDef(entry, backend), true +} + +// RegistryModelAvailability distinguishes an unknown model from a registry +// model that is intentionally unavailable on the requested backend. +func RegistryModelAvailability(backend Backend, model string) (known, available bool) { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return false, false + } + return p.Availability(mode, model) +} + +// RegistryModelDefs returns exact, provider-native model IDs for a backend. CLI +// and cmux backends are projected from their parent provider's model registry. +func RegistryModelDefs(backend Backend) []ModelDef { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return nil + } + out := make([]ModelDef, 0) + for _, m := range p.Models() { + if !m.Preferred { + continue + } + if known, available := p.Availability(mode, m.ID); !known || !available { + continue + } + out = append(out, registryModelDef(m, backend)) + } + SortModelsByReleaseDateDesc(out) + return out +} + +func registryCatalogModels() []Model { + out := make([]Model, 0, len(registry.KnownModels())+5) + for _, p := range registry.Providers() { + apiBackend, err := p.BackendFor(registry.ModeAPI) + if err != nil { + continue + } + for _, m := range p.Models() { + if !m.Preferred { + continue + } + if known, available := p.Availability(registry.ModeAPI, m.ID); !known || !available { + continue + } + out = append(out, Model{ + ID: p.CatalogPrefix + "/" + m.ID, + Backend: apiBackend, + Label: m.Label, + Reasoning: m.Reasoning, + Temperature: m.Temperature, + AdaptiveThinking: m.AdaptiveThinking, + ContextWindow: m.ContextWindow, + ReleaseDate: m.ReleaseDate, + InputMediaTypes: clampInputMediaTypes(apiBackend, m.InputMediaTypes), + SupportedEfforts: append([]api.Effort(nil), m.SupportedEfforts...), + DefaultEffort: m.DefaultEffort, + Priority: m.Priority, + Default: p == registry.Anthropic && m.ID == defaultCatalogModelID, + }) + } + } + out = append(out, agentCatalogModels()...) + return out +} + +// agentCatalogModels projects the agent-mode rows, which carry the bare exact id +// (not a catalog-prefixed one) and an agent-flavoured label. +func agentCatalogModels() []Model { + out := make([]Model, 0) + for _, spec := range []struct { + provider *registry.Provider + label func(string) string + }{ + {registry.Anthropic, func(l string) string { return "Claude Agent · " + strings.TrimPrefix(l, "Claude ") }}, + {registry.OpenAI, func(l string) string { return "Codex Agent · " + l }}, + } { + backend, err := spec.provider.BackendFor(registry.ModeAgent) + if err != nil { + continue + } + for _, m := range spec.provider.Models() { + if !m.Preferred { + continue + } + if known, available := spec.provider.Availability(registry.ModeAgent, m.ID); !known || !available { + continue + } + out = append(out, Model{ + ID: m.ID, + Backend: backend, + Label: spec.label(m.Label), + Reasoning: m.Reasoning, + Temperature: m.Temperature, + AdaptiveThinking: m.AdaptiveThinking && spec.provider == registry.Anthropic, + ContextWindow: m.ContextWindow, + ReleaseDate: m.ReleaseDate, + InputMediaTypes: clampInputMediaTypes(backend, m.InputMediaTypes), + SupportedEfforts: append([]api.Effort(nil), m.SupportedEfforts...), + DefaultEffort: m.DefaultEffort, + Priority: m.Priority, + }) + } + } + return out +} + +// ResolveExactModelForBackend resolves a user/catalog model token into the exact +// model ID the selected backend should receive. It accepts old aliases for input +// compatibility but never returns an alias. +func ResolveExactModelForBackend(backend Backend, model string) (string, bool) { + p, mode, ok := registry.ProviderFor(backend) + if !ok { + return registry.StripProviderPrefix(model), false + } + return p.ResolveExact(mode, model) +} + +// ModelUsesAdaptiveThinking reports whether an Anthropic model uses adaptive +// thinking, per the registry's adaptiveThinking annotation. It accepts aliases, +// provider-prefixed, and dated model tokens by resolving them to their exact +// registry entry first. +func ModelUsesAdaptiveThinking(model string) bool { + exact, ok := registry.Anthropic.ResolveExact(registry.ModeAPI, model) + if !ok { + return false + } + m, ok := registry.Anthropic.Lookup(exact) + return ok && m.AdaptiveThinking +} diff --git a/pkg/ai/model_registry_test.go b/pkg/ai/model_registry_test.go new file mode 100644 index 0000000..869b3d1 --- /dev/null +++ b/pkg/ai/model_registry_test.go @@ -0,0 +1,49 @@ +package ai + +import "testing" + +func TestRegistryModelDefsIncludeFableCapabilities(t *testing.T) { + for _, backend := range []Backend{BackendClaudeAgent, BackendClaudeCLI, BackendClaudeCmux} { + defs := RegistryModelDefs(backend) + var fable *ModelDef + for i := range defs { + if defs[i].ID == "claude-fable-5" { + fable = &defs[i] + break + } + } + if fable == nil { + t.Fatalf("%s registry models omit claude-fable-5: %+v", backend, defs) + } + if !fable.CapabilitiesKnown || !fable.Reasoning || fable.Temperature { + t.Fatalf("%s fable capabilities = %+v", backend, *fable) + } + if len(fable.SupportedEfforts) != 5 || fable.SupportedEfforts[0] != "low" || fable.SupportedEfforts[4] != "max" { + t.Fatalf("%s fable efforts = %v", backend, fable.SupportedEfforts) + } + } +} + +func TestModelUsesAdaptiveThinking(t *testing.T) { + cases := []struct { + model string + want bool + }{ + {"claude-sonnet-5", true}, + {"anthropic/claude-fable-5", true}, + {"sonnet-5", true}, + {"claude-opus-4-8", true}, + {"claude-opus-4-7", true}, + {"claude-haiku-4-5", false}, + {"claude-sonnet-4-6", false}, + {"claude-3-5-sonnet-20241022", false}, + {"gpt-5.5", false}, + } + for _, tc := range cases { + t.Run(tc.model, func(t *testing.T) { + if got := ModelUsesAdaptiveThinking(tc.model); got != tc.want { + t.Errorf("ModelUsesAdaptiveThinking(%q) = %v, want %v", tc.model, got, tc.want) + } + }) + } +} diff --git a/pkg/ai/models_codex.go b/pkg/ai/models_codex.go new file mode 100644 index 0000000..67c42da --- /dev/null +++ b/pkg/ai/models_codex.go @@ -0,0 +1,96 @@ +package ai + +import ( + "context" + "encoding/json" + "fmt" + "os/exec" + "strings" + + "github.com/flanksource/captain/pkg/api" +) + +type codexDebugModelsResponse struct { + Models []codexDebugModel `json:"models"` +} + +type codexDebugModel struct { + Slug string `json:"slug"` + DisplayName string `json:"display_name"` + DefaultReasoningLevel api.Effort `json:"default_reasoning_level"` + SupportedReasoningLevels []codexDebugReasoningLevel `json:"supported_reasoning_levels"` + Visibility string `json:"visibility"` + Priority int `json:"priority"` +} + +type codexDebugReasoningLevel struct { + Effort api.Effort `json:"effort"` +} + +// FetchCodexDebugModels reads the model catalog exposed by the installed Codex +// binary. The command works from Codex's bundled catalog without an API key. +func FetchCodexDebugModels(ctx context.Context, binary string) ([]ModelDef, error) { + binary = strings.TrimSpace(binary) + if binary == "" { + return nil, fmt.Errorf("codex binary is required") + } + ctx, cancel := context.WithTimeout(ctx, remoteModelsTimeout) + defer cancel() + out, err := exec.CommandContext(ctx, binary, "debug", "models").Output() + if err != nil { + return nil, fmt.Errorf("codex debug models: %w", err) + } + return ParseCodexDebugModels(out) +} + +// ParseCodexDebugModels converts Codex's raw catalog into Captain model rows. +func ParseCodexDebugModels(data []byte) ([]ModelDef, error) { + var response codexDebugModelsResponse + if err := json.Unmarshal(data, &response); err != nil { + return nil, fmt.Errorf("codex debug models decode: %w", err) + } + models := make([]ModelDef, 0, len(response.Models)) + for _, model := range response.Models { + if model.Visibility != "list" || strings.TrimSpace(model.Slug) == "" { + continue + } + def := ModelDef{ + ID: strings.TrimSpace(model.Slug), + Name: strings.TrimSpace(model.DisplayName), + Backend: BackendCodexAgent, + DefaultEffort: model.DefaultReasoningLevel, + Priority: model.Priority, + } + if def.Name == "" { + def.Name = def.ID + } + if registry, ok := RegistryModelDef(BackendCodexAgent, def.ID); ok { + def.ReleaseDate = registry.ReleaseDate + def.CapabilitiesKnown = registry.CapabilitiesKnown + def.Reasoning = registry.Reasoning + def.Temperature = registry.Temperature + // models.dev is the source of truth for known model effort metadata. + // Codex debug output remains useful for the locally visible models, + // display names, and priority, but cannot override this catalog. + def.SupportedEfforts = registry.SupportedEfforts + def.DefaultEffort = registry.DefaultEffort + if def.Priority == 0 { + def.Priority = registry.Priority + } + } else { + for _, level := range model.SupportedReasoningLevels { + if level.Effort != api.EffortNone && level.Effort.Valid() { + def.SupportedEfforts = append(def.SupportedEfforts, level.Effort) + } + } + def.CapabilitiesKnown = true + def.Reasoning = len(def.SupportedEfforts) > 0 + } + models = append(models, def) + } + if len(models) == 0 { + return nil, fmt.Errorf("codex debug models returned no visible models") + } + SortModelsByReleaseDateDesc(models) + return models, nil +} diff --git a/pkg/ai/models_codex_test.go b/pkg/ai/models_codex_test.go new file mode 100644 index 0000000..5b3db8f --- /dev/null +++ b/pkg/ai/models_codex_test.go @@ -0,0 +1,37 @@ +package ai + +import ( + "reflect" + "testing" + + "github.com/flanksource/captain/pkg/api" +) + +func TestParseCodexDebugModels(t *testing.T) { + models, err := ParseCodexDebugModels([]byte(`{ + "models": [ + {"slug":"gpt-5.6-sol","display_name":"GPT-5.6-Sol","default_reasoning_level":"low","supported_reasoning_levels":[{"effort":"low"},{"effort":"high"},{"effort":"ultra"}],"visibility":"list","priority":1}, + {"slug":"codex-auto-review","visibility":"hide","priority":2} + ] +}`)) + if err != nil { + t.Fatalf("ParseCodexDebugModels: %v", err) + } + if len(models) != 1 { + t.Fatalf("models = %+v", models) + } + got := models[0] + if got.ID != "gpt-5.6-sol" || got.DefaultEffort != api.EffortNone || got.Priority != 1 || got.ReleaseDate != "2026-07-09" { + t.Fatalf("model = %+v", got) + } + wantEfforts := []api.Effort{api.EffortLow, api.EffortMedium, api.EffortHigh, api.EffortXHigh, api.EffortMax, api.EffortUltra} + if !reflect.DeepEqual(got.SupportedEfforts, wantEfforts) { + t.Fatalf("efforts = %v, want %v", got.SupportedEfforts, wantEfforts) + } +} + +func TestParseCodexDebugModelsRequiresVisibleRows(t *testing.T) { + if _, err := ParseCodexDebugModels([]byte(`{"models":[{"slug":"hidden","visibility":"hide"}]}`)); err == nil { + t.Fatal("expected empty visible catalog error") + } +} diff --git a/pkg/ai/models_remote.go b/pkg/ai/models_remote.go index d666c62..3315293 100644 --- a/pkg/ai/models_remote.go +++ b/pkg/ai/models_remote.go @@ -8,6 +8,8 @@ import ( "sort" "strings" "time" + + "github.com/flanksource/captain/pkg/api" ) // ModelDef is the in-memory shape used by the configure wizard and `captain @@ -15,10 +17,17 @@ import ( // and context data live in pkg/ai/pricing (sourced from OpenRouter) and are // looked up by id at render time. type ModelDef struct { - ID string - Name string - Backend Backend - ReleaseDate string `json:"-"` + ID string `json:"id"` + Name string `json:"label,omitempty"` + Backend Backend `json:"backend,omitempty"` + ReleaseDate string `json:"releaseDate,omitempty"` + CapabilitiesKnown bool `json:"capabilitiesKnown,omitempty"` + Reasoning bool `json:"reasoning,omitempty"` + Temperature bool `json:"temperature"` + InputMediaTypes []string `json:"inputMediaTypes,omitempty"` + SupportedEfforts []api.Effort `json:"supportedEfforts,omitempty"` + DefaultEffort api.Effort `json:"defaultEffort,omitempty"` + Priority int `json:"priority,omitempty"` } // remoteModelsTimeout caps each /v1/models call. The configure wizard is an @@ -43,6 +52,15 @@ type modelEntry struct { CreatedAt string `json:"created_at"` } +type ModelHTTPError struct { + Backend Backend + StatusCode int +} + +func (e ModelHTTPError) Error() string { + return fmt.Sprintf("%s models: HTTP %d", e.Backend, e.StatusCode) +} + // FetchOpenAIModels calls https://api.openai.com/v1/models and returns the // available model IDs as ModelDefs scoped to BackendOpenAI. apiKey is sent // as a Bearer token. An empty apiKey returns an error without making a @@ -78,17 +96,17 @@ func FetchAnthropicModels(ctx context.Context, apiKey string) ([]ModelDef, error // FetchGeminiModels calls Google's Generative Language ListModels endpoint // and returns the IDs scoped to BackendGemini. The endpoint authenticates via -// a `key=` query parameter (no header). The returned `name` field is shaped +// the x-goog-api-key header. The returned `name` field is shaped // "models/gemini-2.5-flash"; we strip the prefix so callers see the bare id. func FetchGeminiModels(ctx context.Context, apiKey string) ([]ModelDef, error) { if strings.TrimSpace(apiKey) == "" { return nil, fmt.Errorf("GEMINI_API_KEY is not set") } - url := "https://generativelanguage.googleapis.com/v1beta/models?key=" + apiKey - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://generativelanguage.googleapis.com/v1beta/models", nil) if err != nil { return nil, err } + req.Header.Set("x-goog-api-key", apiKey) return doModelsRequest(req, BackendGemini) } @@ -120,7 +138,7 @@ func doModelsRequest(req *http.Request, backend Backend) ([]ModelDef, error) { } defer func() { _ = resp.Body.Close() }() if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("%s models: HTTP %d", backend, resp.StatusCode) + return nil, ModelHTTPError{Backend: backend, StatusCode: resp.StatusCode} } var body modelsListResponse @@ -150,7 +168,16 @@ func doModelsRequest(req *http.Request, backend Backend) ([]ModelDef, error) { if releaseDate == "" { releaseDate = CatalogReleaseDate(backend, id) } - out = append(out, ModelDef{ID: id, Name: name, Backend: backend, ReleaseDate: releaseDate}) + def := ModelDef{ID: id, Name: name, Backend: backend, ReleaseDate: releaseDate} + if registry, ok := RegistryModelDef(backend, id); ok { + def.CapabilitiesKnown = registry.CapabilitiesKnown + def.Reasoning = registry.Reasoning + def.Temperature = registry.Temperature + def.SupportedEfforts = registry.SupportedEfforts + def.DefaultEffort = registry.DefaultEffort + def.Priority = registry.Priority + } + out = append(out, def) } return out, nil } @@ -169,7 +196,17 @@ func (m modelEntry) releaseDate() string { // models from the static catalog in pkg/cli (agentCatalogModels), so passing // one returns an error. func ListModels(ctx context.Context, backend Backend) ([]ModelDef, error) { - fetch, apiKey := remoteFetcherFor(backend) + resolved, err := ResolveAPIKey(backend) + if err != nil { + return nil, err + } + return ListModelsWithAPIKey(ctx, backend, resolved.Token) +} + +// ListModelsWithAPIKey validates a candidate credential directly against the +// provider model endpoint without reading or writing Captain's credential vault. +func ListModelsWithAPIKey(ctx context.Context, backend Backend, apiKey string) ([]ModelDef, error) { + fetch := remoteFetcherFor(backend) if fetch == nil { return nil, fmt.Errorf("backend %s has no live model listing", backend) } @@ -186,20 +223,20 @@ func ListModels(ctx context.Context, backend Backend) ([]ModelDef, error) { return models, nil } -// remoteFetcherFor returns the live-list function and API key for an API -// backend. Returns (nil, "") for any backend without a live listing endpoint +// remoteFetcherFor returns the live-list function for an API backend. It +// returns nil for any backend without a live listing endpoint // (every CLI/agent backend, which lists from the static catalog instead). -func remoteFetcherFor(backend Backend) (fetch func(context.Context, string) ([]ModelDef, error), apiKey string) { +func remoteFetcherFor(backend Backend) func(context.Context, string) ([]ModelDef, error) { switch backend { case BackendOpenAI: - return FetchOpenAIModels, GetAPIKeyFromEnv(backend) + return FetchOpenAIModels case BackendAnthropic: - return FetchAnthropicModels, GetAPIKeyFromEnv(backend) + return FetchAnthropicModels case BackendGemini: - return FetchGeminiModels, GetAPIKeyFromEnv(backend) + return FetchGeminiModels case BackendDeepSeek: - return FetchDeepSeekModels, GetAPIKeyFromEnv(backend) + return FetchDeepSeekModels default: - return nil, "" + return nil } } diff --git a/pkg/ai/models_remote_test.go b/pkg/ai/models_remote_test.go index 3150dcb..44d9716 100644 --- a/pkg/ai/models_remote_test.go +++ b/pkg/ai/models_remote_test.go @@ -6,9 +6,12 @@ import ( "io" "net/http" "net/http/httptest" + "path/filepath" "strings" "testing" "time" + + "github.com/flanksource/captain/pkg/credentials" ) // withTestServer redirects http.DefaultClient.Do to the given test server by @@ -27,8 +30,7 @@ type rewriteTransport struct { } func (r rewriteTransport) RoundTrip(req *http.Request) (*http.Response, error) { - // Strip scheme+host but preserve path AND query so the Gemini fetcher's - // `?key=…` survives the rewrite. + // Strip scheme+host while preserving request metadata for provider checks. target := r.base + req.URL.Path if req.URL.RawQuery != "" { target += "?" + req.URL.RawQuery @@ -183,6 +185,8 @@ func TestFetchDeepSeekModels_EmptyKey(t *testing.T) { } func TestListModels_ErrorsOnMissingKey(t *testing.T) { + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) t.Setenv("OPENAI_API_KEY", "") t.Setenv("ANTHROPIC_API_KEY", "") t.Setenv("GEMINI_API_KEY", "") @@ -255,14 +259,18 @@ func TestListModels_SortsAlphabetically(t *testing.T) { func TestFetchGeminiModels_StripsModelsPrefix(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Gemini sends the key as a query param, not a header. - if got := r.URL.Query().Get("key"); got != "g-test" { - t.Errorf("key= = %q", got) + if got := r.Header.Get("x-goog-api-key"); got != "g-test" { + t.Errorf("x-goog-api-key = %q", got) + } + if r.URL.RawQuery != "" { + t.Errorf("API key must not be present in URL query: %q", r.URL.RawQuery) } + // Both ids must be preferred registry models: the release-date fallback + // reads Catalog(), which only carries preferred entries. _ = json.NewEncoder(w).Encode(map[string]any{ "models": []map[string]any{ - {"name": "models/gemini-2.5-flash", "display_name": "Gemini 2.5 Flash"}, - {"name": "models/gemini-2.5-pro", "display_name": "Gemini 2.5 Pro"}, + {"name": "models/gemini-3.6-flash", "display_name": "Gemini 3.6 Flash"}, + {"name": "models/gemini-3.5-flash", "display_name": "Gemini 3.5 Flash"}, }, }) })) @@ -273,10 +281,10 @@ func TestFetchGeminiModels_StripsModelsPrefix(t *testing.T) { if err != nil { t.Fatalf("FetchGeminiModels: %v", err) } - if len(got) != 2 || got[0].ID != "gemini-2.5-flash" || got[0].Name != "Gemini 2.5 Flash" { + if len(got) != 2 || got[0].ID != "gemini-3.6-flash" || got[0].Name != "Gemini 3.6 Flash" { t.Errorf("unexpected: %+v", got) } - if got[1].ID != "gemini-2.5-pro" || got[1].ReleaseDate == "" { + if got[1].ID != "gemini-3.5-flash" || got[1].ReleaseDate == "" { t.Errorf("expected Gemini catalog release-date fallback, got %+v", got) } for _, m := range got { diff --git a/pkg/ai/pricing/catalog.go b/pkg/ai/pricing/catalog.go new file mode 100644 index 0000000..c96812f --- /dev/null +++ b/pkg/ai/pricing/catalog.go @@ -0,0 +1,47 @@ +package pricing + +// catalog is the generated model catalog. It is aliased because this package's +// own `registry` identifier is the OpenRouter-backed price map below. +import catalog "github.com/flanksource/captain/pkg/api/registry" + +// catalogInfo synthesizes a ModelInfo from the generated model catalog, which +// carries models.dev's published rate for every model captain can route to. It +// returns false for ids the catalog does not price. ContextWindow/MaxTokens are +// left zero — this fills prices only; context is supplied separately (catalog +// metadata or OpenRouter). +// +// This replaced a hand-written Claude family table that priced on the substrings +// "opus"/"sonnet"/"haiku": it billed every Opus at the retired 4.1 rate and had +// no entry at all for Fable, which therefore fell through to Sonnet's rate. +func catalogInfo(id string) (ModelInfo, bool) { + cost, ok := catalog.CostFor(id) + if !ok { + return ModelInfo{}, false + } + return ModelInfo{ + ModelID: id, + InputPrice: cost.Input, + OutputPrice: cost.Output, + CacheReadsPrice: cost.CacheRead, + CacheWritesPrice: cost.CacheWrite, + }, true +} + +// applyCatalogPrices overlays the catalog's rates onto registry rows, filling +// only the price fields OpenRouter left zero (OpenRouter stays primary, per the +// cost-merge decision). It does not add new ids: an id absent from the registry +// is priced lazily via catalogInfo on lookup (see GetModelInfo's lookup-on-miss). +func applyCatalogPrices() { + registryMu.RLock() + ids := make([]string, 0, len(registry)) + for id := range registry { + ids = append(ids, id) + } + registryMu.RUnlock() + + for _, id := range ids { + if info, ok := catalogInfo(id); ok { + MergeModel(id, info, MergeFillMissing) + } + } +} diff --git a/pkg/ai/pricing/static_test.go b/pkg/ai/pricing/catalog_test.go similarity index 52% rename from pkg/ai/pricing/static_test.go rename to pkg/ai/pricing/catalog_test.go index 25a1a88..c8382d0 100644 --- a/pkg/ai/pricing/static_test.go +++ b/pkg/ai/pricing/catalog_test.go @@ -4,7 +4,7 @@ import ( "errors" "testing" - "github.com/flanksource/captain/pkg/claude" + catalog "github.com/flanksource/captain/pkg/api/registry" ) // errTestShortCircuit makes EnsureLoaded return immediately so lookups in these @@ -73,60 +73,75 @@ func TestMergeModelOverlayReplaces(t *testing.T) { } } -func TestClaudeStaticInfo(t *testing.T) { - sonnet := claude.PricingTable[claude.ModelFamilySonnet4] - got, ok := claudeStaticInfo("claude-sonnet-4-6") - if !ok { - t.Fatal("expected sonnet id to classify") +func TestCatalogInfo(t *testing.T) { + for _, id := range []string{"claude-sonnet-4-6", "gemini-3.5-flash", "gpt-5.6"} { + want, ok := catalog.CostFor(id) + if !ok { + t.Fatalf("catalog must price %q", id) + } + got, ok := catalogInfo(id) + if !ok { + t.Fatalf("catalogInfo(%q) found no price", id) + } + if got.InputPrice != want.Input || got.OutputPrice != want.Output || + got.CacheReadsPrice != want.CacheRead || got.CacheWritesPrice != want.CacheWrite { + t.Errorf("catalogInfo(%q) = %+v, want catalog rate %+v", id, got, want) + } } - if got.InputPrice != sonnet.InputPerMTok || got.OutputPrice != sonnet.OutputPerMTok || - got.CacheReadsPrice != sonnet.CacheReadPerMTok || got.CacheWritesPrice != sonnet.CacheWritePerMTok { - t.Fatalf("static sonnet price mismatch: %+v", got) - } - if _, ok := claudeStaticInfo("gpt-4o"); ok { - t.Fatal("non-claude id must not classify") + if _, ok := catalogInfo("totally-unknown-model-zzz"); ok { + t.Fatal("an id the catalog does not price must not synthesize one") } } -func TestApplyStaticClaudeFillsMissingKeepsOpenRouter(t *testing.T) { - // OpenRouter gave a non-zero input price but no output price for this claude - // row; a non-claude row must be left entirely untouched. +func TestApplyCatalogPricesFillsMissingKeepsOpenRouter(t *testing.T) { + // OpenRouter gave a non-zero input price but no output price for this row; + // a row the catalog does not price must be left entirely untouched. withIsolatedRegistry(t, map[string]ModelInfo{ - "anthropic/claude-sonnet-4": {ModelID: "anthropic/claude-sonnet-4", InputPrice: 2.5}, - "openai/gpt-4o": {ModelID: "openai/gpt-4o", InputPrice: 5}, + "anthropic/claude-sonnet-4-6": {ModelID: "anthropic/claude-sonnet-4-6", InputPrice: 2.5}, + "openai/gpt-4o": {ModelID: "openai/gpt-4o", InputPrice: 5}, }) - applyStaticClaude() + applyCatalogPrices() - claudeRow, _ := GetModelInfo("anthropic/claude-sonnet-4") + sonnet, _ := catalog.CostFor("claude-sonnet-4-6") + claudeRow, _ := GetModelInfo("anthropic/claude-sonnet-4-6") if claudeRow.InputPrice != 2.5 { t.Fatalf("OpenRouter input price overwritten: %+v", claudeRow) } - if claudeRow.OutputPrice != claude.PricingTable[claude.ModelFamilySonnet4].OutputPerMTok { - t.Fatalf("static output price not filled: %+v", claudeRow) + if claudeRow.OutputPrice != sonnet.Output { + t.Fatalf("catalog output price not filled: %+v", claudeRow) } openaiRow, _ := GetModelInfo("openai/gpt-4o") if openaiRow.InputPrice != 5 || openaiRow.OutputPrice != 0 { - t.Fatalf("non-claude row mutated: %+v", openaiRow) + t.Fatalf("unpriced row mutated: %+v", openaiRow) } } -func TestGetModelInfoClassifyOnMiss(t *testing.T) { - // Registry intentionally empty: a Claude id absent from OpenRouter must - // still resolve via the static table. +// TestGetModelInfoLookupOnMiss pins that catalog models OpenRouter never lists +// still price — and that the fallback is per-model, not per-family: Opus must +// not come back at the retired $15/$75 that the old static table applied to +// every id containing "opus". +func TestGetModelInfoLookupOnMiss(t *testing.T) { withIsolatedRegistry(t, nil) - haiku := claude.PricingTable[claude.ModelFamilyHaiku4] - got, ok := GetModelInfo("claude-haiku-4-5") - if !ok { - t.Fatal("expected classify-on-miss to price a known claude family") + for _, id := range []string{"claude-haiku-4-5", "claude-opus-5", "gemini-3.5-flash"} { + want, _ := catalog.CostFor(id) + got, ok := GetModelInfo(id) + if !ok { + t.Fatalf("lookup-on-miss did not price catalog model %q", id) + } + if got.InputPrice != want.Input || got.OutputPrice != want.Output { + t.Errorf("GetModelInfo(%q) = %.2f/%.2f, want %.2f/%.2f", + id, got.InputPrice, got.OutputPrice, want.Input, want.Output) + } } - if got.InputPrice != haiku.InputPerMTok || got.OutputPrice != haiku.OutputPerMTok { - t.Fatalf("classify-on-miss price mismatch: %+v", got) + + if opus, _ := GetModelInfo("claude-opus-5"); opus.InputPrice == 15 || opus.OutputPrice == 75 { + t.Errorf("claude-opus-5 priced at the retired Opus 4.1 rate: %+v", opus) } if _, ok := GetModelInfo("totally-unknown-model-zzz"); ok { - t.Fatal("unknown non-claude id must remain a miss") + t.Fatal("unknown id must remain a miss") } } diff --git a/pkg/ai/pricing/openrouter.go b/pkg/ai/pricing/openrouter.go index 113f54d..7151ad0 100644 --- a/pkg/ai/pricing/openrouter.go +++ b/pkg/ai/pricing/openrouter.go @@ -75,7 +75,7 @@ func EnsureLoaded() { log.Debugf("Loaded OpenRouter pricing from cache (age: %s)", time.Since(cache.Timestamp)) pricingCache = cache MergeModels(cache.Models) - applyStaticClaude() + applyCatalogPrices() return } } @@ -89,7 +89,7 @@ func EnsureLoaded() { pricingCache = &PricingCache{Timestamp: time.Now(), Models: models} MergeModels(models) - applyStaticClaude() + applyCatalogPrices() } func fetchOpenRouterPricing() (map[string]*ModelInfo, error) { diff --git a/pkg/ai/pricing/registry.go b/pkg/ai/pricing/registry.go index 02d9054..ec10ad4 100644 --- a/pkg/ai/pricing/registry.go +++ b/pkg/ai/pricing/registry.go @@ -32,9 +32,10 @@ func GetModelInfo(model string) (ModelInfo, bool) { if ok { return info, true } - // Classify-on-miss: arbitrary Claude ids (e.g. catalog "claude-sonnet-4-6") - // that OpenRouter never lists are still priced from the static table. - return claudeStaticInfo(model) + // Lookup-on-miss: catalog ids that OpenRouter never lists (e.g. + // "claude-sonnet-4-6", every Gemini id) are still priced from the generated + // catalog's own rates. + return catalogInfo(model) } // Contains reports whether model is an exact id in the pricing registry. Unlike diff --git a/pkg/ai/pricing/static.go b/pkg/ai/pricing/static.go deleted file mode 100644 index 443eabe..0000000 --- a/pkg/ai/pricing/static.go +++ /dev/null @@ -1,42 +0,0 @@ -package pricing - -import "github.com/flanksource/captain/pkg/claude" - -// claudeStaticInfo synthesizes a ModelInfo from the hand-curated static Claude -// price table for any model id that classifies to a known Claude family -// (opus/sonnet/haiku). It returns false for non-Claude ids and unknown -// families. ContextWindow/MaxTokens are left zero — the static table carries -// prices only; context is supplied separately (catalog or OpenRouter). -func claudeStaticInfo(id string) (ModelInfo, bool) { - p, ok := claude.PricingTable[claude.ClassifyModel(id)] - if !ok { - return ModelInfo{}, false - } - return ModelInfo{ - ModelID: id, - InputPrice: p.InputPerMTok, - OutputPrice: p.OutputPerMTok, - CacheReadsPrice: p.CacheReadPerMTok, - CacheWritesPrice: p.CacheWritePerMTok, - }, true -} - -// applyStaticClaude overlays the static Claude price table onto registry rows -// that classify to a known Claude family, filling only the price fields -// OpenRouter left zero (OpenRouter stays primary, per the cost-merge decision). -// It does not add new ids: an id absent from the registry is priced lazily via -// claudeStaticInfo on lookup (see GetModelInfo's classify-on-miss). -func applyStaticClaude() { - registryMu.RLock() - ids := make([]string, 0, len(registry)) - for id := range registry { - ids = append(ids, id) - } - registryMu.RUnlock() - - for _, id := range ids { - if info, ok := claudeStaticInfo(id); ok { - MergeModel(id, info, MergeFillMissing) - } - } -} diff --git a/pkg/ai/pricing_coverage_test.go b/pkg/ai/pricing_coverage_test.go new file mode 100644 index 0000000..72d7599 --- /dev/null +++ b/pkg/ai/pricing_coverage_test.go @@ -0,0 +1,100 @@ +package ai + +import ( + "testing" + + "github.com/flanksource/captain/pkg/ai/pricing" + "github.com/flanksource/captain/pkg/api/registry" +) + +// knownPricingGaps are preferred catalog models the pricing source does not list +// under any key, so captain reports $0 for them. These are upstream data gaps, +// not prefix bugs: the id resolves to no entry whether or not it is namespaced. +// Listed explicitly so the hole is visible and shrinks when upstream catches up, +// rather than being hidden by a permissive assertion. +// +// Empty since the generated catalog carries its own list prices: a model captain +// knows about now prices even when OpenRouter never lists it — the api-only +// "gpt-5.6" base id (OpenRouter has only the {sol,terra,luna} variants) was the +// last hole. Keep the mechanism for the next upstream gap. +var knownPricingGaps = map[string]string{} + +// TestPricingIDsCoverEveryCatalogModel guards the failure mode that motivated +// unifying the pricing prefixes: a wrong namespace does not error, it just +// misses, and the run reports $0. +// +// Gemini is the trap — its catalog namespace is "googleai" but OpenRouter keys +// it under "google" — and captain had three hand-written copies of that mapping +// (PricingIDs, orPrefix, pricingModelID) that could drift apart independently. +func TestPricingIDsCoverEveryCatalogModel(t *testing.T) { + pricing.EnsureLoaded() + + for _, p := range registry.Providers() { + backend, err := p.BackendFor(registry.ModeAPI) + if err != nil { + t.Fatalf("%s has no API backend: %v", p.Name, err) + } + for _, m := range p.Models() { + if !m.Preferred { + continue + } + t.Run(p.Name+"/"+m.ID, func(t *testing.T) { + info, ok := lookupPricing(backend, m.ID) + if reason, gap := knownPricingGaps[m.ID]; gap { + if ok { + t.Fatalf("%s now has pricing — remove it from knownPricingGaps (%s)", m.ID, reason) + } + t.Skipf("known upstream pricing gap: %s", reason) + } + if !ok { + t.Fatalf("no pricing for %s via %v; a missing prefix reports $0 rather than failing", + m.ID, PricingIDs(backend, m.ID)) + } + if info.InputPrice <= 0 && info.OutputPrice <= 0 { + t.Errorf("%s priced at zero (input=%v output=%v)", m.ID, info.InputPrice, info.OutputPrice) + } + }) + } + } +} + +// TestPricingPrefixIsNotCatalogPrefix pins the two namespaces apart. If someone +// "simplifies" PricingPrefix to reuse CatalogPrefix, every Gemini price silently +// resolves to nothing — this fails first. +func TestPricingPrefixIsNotCatalogPrefix(t *testing.T) { + if registry.Google.CatalogPrefix != "googleai" { + t.Errorf("Google.CatalogPrefix = %q, want googleai (genkit/menu namespace)", registry.Google.CatalogPrefix) + } + if registry.Google.PricingPrefix != "google" { + t.Errorf("Google.PricingPrefix = %q, want google (OpenRouter key)", registry.Google.PricingPrefix) + } +} + +// TestPricingAgreesAcrossCatalogAndBilling: the catalog and the cost path must +// price a model identically. They used to run different lookups — the catalog +// tried the bare id first (which the classify-on-miss Claude family table always +// answered, at the wrong rate) while billing tried the prefixed key first — so +// the price shown could differ from the price charged. +func TestPricingAgreesAcrossCatalogAndBilling(t *testing.T) { + pricing.EnsureLoaded() + + for _, model := range []string{"claude-sonnet-5", "claude-opus-4-8"} { + t.Run(model, func(t *testing.T) { + catalog, ok := lookupPricing(BackendAnthropic, model) + if !ok { + t.Fatalf("catalog has no price for %s", model) + } + var billing pricing.ModelInfo + for _, id := range PricingIDs(BackendAnthropic, model) { + if info, found := pricing.GetModelInfo(id); found { + billing = info + break + } + } + if catalog.InputPrice != billing.InputPrice || catalog.OutputPrice != billing.OutputPrice { + t.Errorf("catalog prices %s at in=%v/out=%v but billing uses in=%v/out=%v", + model, catalog.InputPrice, catalog.OutputPrice, billing.InputPrice, billing.OutputPrice) + } + }) + } +} diff --git a/pkg/ai/prompt/document.go b/pkg/ai/prompt/document.go index eab3e4e..143d666 100644 --- a/pkg/ai/prompt/document.go +++ b/pkg/ai/prompt/document.go @@ -1,6 +1,7 @@ package prompt import ( + "bytes" "fmt" "strings" @@ -23,9 +24,11 @@ type Document struct { // no frontmatter. It is the lossless source of truth for reserialization. Frontmatter map[string]any // Spec is the frontmatter decoded into the typed spec (the dotprompt-only - // keys config/input/output/name/description are stripped). Zero when the - // source is body-only or declares no spec-native keys. + // keys config/input/output/name/description/runtimes are stripped). Zero + // when the source is body-only or declares no spec-native keys. Spec api.Spec + // Runtimes are the prompt's default parallel execution targets. + Runtimes []api.Model // Body is the unrendered Handlebars template body. Body string } @@ -46,12 +49,61 @@ func Parse(source string) (*Document, error) { return nil, fmt.Errorf("parse prompt frontmatter: %w", err) } doc.Frontmatter = raw + runtimes, err := decodePromptRuntimes(raw) + if err != nil { + return nil, fmt.Errorf("decode prompt runtimes: %w", err) + } + doc.Runtimes = runtimes if err := decodeSpecFrontmatter(raw, &doc.Spec); err != nil { return nil, fmt.Errorf("decode prompt frontmatter into spec: %w", err) } return doc, nil } +func decodePromptRuntimes(raw map[string]any) ([]api.Model, error) { + value, ok := raw["runtimes"] + if !ok { + return nil, nil + } + values, ok := value.([]any) + if !ok { + return nil, fmt.Errorf("runtimes must be a list") + } + if len(values) == 0 { + return nil, fmt.Errorf("runtimes must contain at least two entries") + } + runtimes := make([]api.Model, 0, len(values)) + for i, value := range values { + encoded, err := yaml.Marshal(value) + if err != nil { + return nil, fmt.Errorf("runtime %d: encode: %w", i+1, err) + } + if _, compact := value.(string); compact { + encoded, err = yaml.Marshal([]any{value}) + if err != nil { + return nil, fmt.Errorf("runtime %d: encode compact value: %w", i+1, err) + } + var list api.ModelList + if err := yaml.Unmarshal(encoded, &list); err != nil { + return nil, fmt.Errorf("runtime %d: %w", i+1, err) + } + if len(list) != 1 { + return nil, fmt.Errorf("runtime %d: expected one compact model", i+1) + } + runtimes = append(runtimes, list[0]) + continue + } + var runtime api.Model + dec := yaml.NewDecoder(bytes.NewReader(encoded)) + dec.KnownFields(true) + if err := dec.Decode(&runtime); err != nil { + return nil, fmt.Errorf("runtime %d: %w", i+1, err) + } + runtimes = append(runtimes, runtime) + } + return runtimes, nil +} + // String reserializes the document to .prompt text: "---\n\n---\n", // or the body alone when there is no frontmatter. The output satisfies the // dotprompt frontmatter grammar so Parse can read it back. diff --git a/pkg/ai/prompt/prompt.go b/pkg/ai/prompt/prompt.go index e3806f3..d945b8a 100644 --- a/pkg/ai/prompt/prompt.go +++ b/pkg/ai/prompt/prompt.go @@ -178,7 +178,7 @@ func decodeSpecFrontmatter(raw map[string]any, req *ai.Request) error { specRaw := map[string]any{} for key, value := range raw { switch key { - case "config", "input", "output", "name", "description": + case "config", "input", "output", "name", "description", "runtimes": continue default: specRaw[key] = value diff --git a/pkg/ai/prompt/prompt_test.go b/pkg/ai/prompt/prompt_test.go index 447a5d8..f01dafb 100644 --- a/pkg/ai/prompt/prompt_test.go +++ b/pkg/ai/prompt/prompt_test.go @@ -21,12 +21,48 @@ func TestRender_FrontmatterAndMessages(t *testing.T) { tmpl, err := LoadFS(library, "testdata/commit.prompt") require.NoError(t, err) - req, cfg, err := tmpl.Render(map[string]any{"diff": "+added a line"}, nil) + req, cfg, err := tmpl.Render(map[string]any{ + "patch": "+func Login() bool { return a < b && c > d }", + "maxBodyLines": 3, + }, nil) require.NoError(t, err) - assert.Contains(t, req.Prompt.System, "Conventional Commit") - assert.Contains(t, req.Prompt.User, "+added a line") - assert.NotContains(t, req.Prompt.User, "Conventional Commit", "system text must not leak into the user prompt") + assert.Contains(t, req.Prompt.System, "commit message generator") + assert.Contains(t, req.Prompt.User, "a < b && c > d") + assert.NotContains(t, req.Prompt.User, "<", "patch content must not be HTML-escaped") + assert.Contains(t, req.Prompt.User, "body: at most 3 line(s)") + assert.NotContains(t, req.Prompt.User, "commit message generator", "system text must not leak into the user prompt") + assert.JSONEq(t, `{ + "type": "object", + "additionalProperties": false, + "required": ["type", "subject"], + "properties": { + "type": { + "type": "string", + "description": "Conventional commit type: feat|fix|perf|refactor|test|docs|build|ci|chore|revert" + }, + "scope": { + "type": "string", + "description": "Optional scope, e.g. db, api, fe, kubernetes" + }, + "subject": { + "type": "string", + "description": "Imperative subject line, max 100 chars, no trailing period" + }, + "body": { + "type": "string", + "description": "Optional body explaining why and impact" + } + } + }`, string(req.Prompt.SchemaJSON)) + + withoutCap, _, err := tmpl.Render(map[string]any{ + "patch": "+trivial change", + "maxBodyLines": 0, + }, nil) + require.NoError(t, err) + assert.Contains(t, withoutCap.Prompt.User, "body: omit unless the change is non-trivial") + assert.NotContains(t, withoutCap.Prompt.User, "body: at most") assert.Equal(t, "claude-sonnet-4-6", cfg.Model.Name) assert.Equal(t, ai.BackendAnthropic, cfg.Model.Backend, "model name should infer the anthropic backend") @@ -111,7 +147,10 @@ func TestRender_FrontmatterOutputSchema(t *testing.T) { func TestLibrary_Render(t *testing.T) { lib := NewLibrary(library) - req, cfg, err := lib.Render("testdata/commit.prompt", map[string]any{"diff": "x"}, nil) + req, cfg, err := lib.Render("testdata/commit.prompt", map[string]any{ + "patch": "x", + "maxBodyLines": 0, + }, nil) require.NoError(t, err) assert.Equal(t, "claude-sonnet-4-6", cfg.Model.Name) assert.Contains(t, req.Prompt.User, "x") @@ -130,6 +169,7 @@ func TestRender_BackendFixtureExamples(t *testing.T) { "testdata/fixtures/claude-cmux-sonnet.prompt": {backend: api.BackendClaudeCmux, model: "claude-cmux-sonnet"}, "testdata/fixtures/claude-cli-opus.prompt": {backend: api.BackendClaudeCLI, model: "claude-agent-opus"}, "testdata/fixtures/claude-cli-sonnet.prompt": {backend: api.BackendClaudeCLI, model: "claude-agent-sonnet"}, + "testdata/fixtures/codex-agent.prompt": {backend: api.BackendCodexAgent, model: "gpt-5-codex"}, "testdata/fixtures/codex-cmux.prompt": {backend: api.BackendCodexCmux, model: "gpt-5-codex"}, "testdata/fixtures/deepseek.prompt": {backend: api.BackendDeepSeek, model: "deepseek-reasoner"}, "testdata/fixtures/codex-cli.prompt": {backend: api.BackendCodexCLI, model: "gpt-5-codex"}, diff --git a/pkg/ai/prompt/testdata/commit.prompt b/pkg/ai/prompt/testdata/commit.prompt index 21796cc..6b4223f 100644 --- a/pkg/ai/prompt/testdata/commit.prompt +++ b/pkg/ai/prompt/testdata/commit.prompt @@ -5,10 +5,49 @@ config: temperature: 0.2 input: schema: - diff: string + type: object + additionalProperties: false + required: [patch, maxBodyLines] + properties: + patch: + type: string + description: "Git patch to summarize" + maxBodyLines: + type: integer + description: "Maximum commit-message body lines; zero omits the cap" +output: + schema: + type: object + additionalProperties: false + required: [type, subject] + properties: + type: + type: string + description: "Conventional commit type: feat|fix|perf|refactor|test|docs|build|ci|chore|revert" + scope: + type: string + description: "Optional scope, e.g. db, api, fe, kubernetes" + subject: + type: string + description: "Imperative subject line, max 100 chars, no trailing period" + body: + type: string + description: "Optional body explaining why and impact" --- {{role "system"}} -You write Conventional Commit messages. +You are a commit message generator. Analyze the diff below and produce a Conventional Commit message. {{role "user"}} -Diff: -{{diff}} +DIFF INPUT: + +{{patch}} + +REQUIREMENTS: +- type: one of feat|fix|perf|refactor|test|docs|build|ci|chore|revert +- scope (optional): the area of the codebase affected, e.g. db|api|fe|crd|chart|docker|kubernetes|terraform|aws +- subject: imperative mood, no trailing period, <=100 characters +{{#if maxBodyLines}} +- body: at most {{maxBodyLines}} line(s); omit entirely for trivial changes. Explain why and impact; include "BREAKING CHANGE: ..." if behavior changes or APIs are removed; reference issues like "Refs #123" or "Fixes #123" if applicable +{{else}} +- body: omit unless the change is non-trivial. When present, explain why and impact; include "BREAKING CHANGE: ..." if behavior changes or APIs are removed; reference issues like "Refs #123" or "Fixes #123" if applicable +{{/if}} +- Explain intention and effect; do not restate the code. diff --git a/pkg/ai/prompt/testdata/fixtures/codex-agent.prompt b/pkg/ai/prompt/testdata/fixtures/codex-agent.prompt new file mode 100644 index 0000000..6677efc --- /dev/null +++ b/pkg/ai/prompt/testdata/fixtures/codex-agent.prompt @@ -0,0 +1,19 @@ +--- +model: gpt-5-codex +backend: codex-agent +effort: medium +setup: + cwd: . +permissions: + presets: + - edit + mcp: + disabled: true +memory: + bare: true +--- +{{role "system"}} +You are a Codex agent running through the app-server SDK. Make scoped edits and verify the result. +{{role "user"}} +Code task: +{{task}} diff --git a/pkg/ai/provider.go b/pkg/ai/provider.go index 97c687e..dc980f9 100644 --- a/pkg/ai/provider.go +++ b/pkg/ai/provider.go @@ -20,6 +20,7 @@ const ( BackendCodexCLI = api.BackendCodexCLI BackendGeminiCLI = api.BackendGeminiCLI BackendClaudeAgent = api.BackendClaudeAgent + BackendCodexAgent = api.BackendCodexAgent BackendClaudeCmux = api.BackendClaudeCmux BackendCodexCmux = api.BackendCodexCmux ) diff --git a/pkg/ai/provider/claude_cli.go b/pkg/ai/provider/claude_cli.go new file mode 100644 index 0000000..a18f4d6 --- /dev/null +++ b/pkg/ai/provider/claude_cli.go @@ -0,0 +1,308 @@ +package provider + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" +) + +var claudeCLICommand = "claude" + +type ClaudeCLI struct { + model string +} + +func NewClaudeCLI(model string) *ClaudeCLI { + if strings.TrimSpace(model) == "" { + model = "opus" + } + model = ai.NormalizeModelForBackend(ai.BackendClaudeCLI, model) + return &ClaudeCLI{model: model} +} + +func (c *ClaudeCLI) GetModel() string { return c.model } +func (c *ClaudeCLI) GetBackend() ai.Backend { return ai.BackendClaudeCLI } + +func (c *ClaudeCLI) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { + start := time.Now() + events, err := c.ExecuteStream(ctx, req) + if err != nil { + return nil, err + } + resp, err := CoalesceStreamForBackend(ctx, ai.BackendClaudeCLI, c.model, events, start) + if err != nil { + return nil, err + } + if req.Prompt.Schema != nil { + raw, _ := resp.StructuredData.(json.RawMessage) + if len(raw) == 0 { + raw = json.RawMessage(resp.Text) + } + if err := ai.BindStructuredOutput(req.Prompt.Schema, raw); err != nil { + return nil, err + } + resp.StructuredData = req.Prompt.Schema + resp.Text = "" + } + return resp, nil +} + +func (c *ClaudeCLI) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai.Event, error) { + args, cleanup, err := buildClaudeCLIArgs(c.model, req) + if err != nil { + return nil, err + } + env := []string(nil) + if req.Setup != nil { + env = commandEnv(req.Setup.Env) + } + cmd, stdout, stderrBuf, err := startCLIStream(ctx, claudeCLICommand, args, []byte(composePrompt(req)), req.Cwd(), env) + if err != nil { + cleanup() + return nil, err + } + out := make(chan ai.Event, 16) + go func() { + defer close(out) + defer cleanup() + defer func() { _ = stdout.Close() }() + iterator := claude.NewStreamJSONIterator(stdout) + for iterator.Next() { + for _, ev := range claudeEntryEvents(iterator.Entry(), c.model) { + emit(ctx, out, ev) + } + } + if err := iterator.Err(); err != nil { + emit(ctx, out, ai.Event{Kind: ai.EventError, Error: fmt.Sprintf("claude stream-json: %v", err), Model: c.model}) + } + if err := finishCLIStream(ctx, cmd, stderrBuf); err != nil { + emit(ctx, out, ai.Event{Kind: ai.EventError, Error: err.Error(), Model: c.model}) + } + }() + return out, nil +} + +func buildClaudeCLIArgs(model string, req ai.Request) ([]string, func(), error) { + args := []string{"-p", "--verbose", "--output-format", "stream-json"} + cleanup := func() {} + if m := claudeCLIModel(model); m != "" { + args = append(args, "--model", m) + } + if req.Prompt.System != "" { + args = append(args, "--system-prompt", req.Prompt.System) + } + if req.Prompt.AppendSystem != "" { + args = append(args, "--append-system-prompt", req.Prompt.AppendSystem) + } + if req.SessionID != "" { + args = append(args, "--resume", req.SessionID) + } + if req.Effort != "" { + args = append(args, "--effort", string(req.Effort)) + } + if req.Budget.Cost > 0 { + args = append(args, "--max-budget-usd", fmt.Sprintf("%g", req.Budget.Cost)) + } + if mode := cliClaudePermissionMode(req.Permissions.Mode); mode != "" { + args = append(args, "--permission-mode", mode) + } + if len(req.Permissions.Tools.Allow) > 0 { + args = append(args, "--allowedTools", strings.Join(req.Permissions.Tools.Allow, ",")) + } + if len(req.Permissions.Tools.Deny) > 0 { + args = append(args, "--disallowedTools", strings.Join(req.Permissions.Tools.Deny, ",")) + } + for _, dir := range req.Memory.Skills { + if strings.TrimSpace(dir) != "" { + args = append(args, "--plugin-dir", dir) + } + } + if req.Memory.SkipSkills { + args = append(args, "--disable-slash-commands") + } + if req.Memory.Bare || req.Permissions.HasPreset(api.PresetBare) { + args = append(args, "--bare") + } + if req.Permissions.MCP.Disabled { + args = append(args, "--mcp-config", `{"mcpServers":{}}`, "--strict-mcp-config") + } + if binary, ok := captainBinary(); ok && api.MonitorHooksEnabled(req) { + settingsPath, remove, err := writeClaudeMonitorSettings(binary) + if err != nil { + return nil, cleanup, err + } + cleanup = remove + args = append(args, "--settings", settingsPath) + } + schema, err := ai.SchemaJSONForBackend(ai.BackendClaudeCLI, req.Prompt) + if err != nil { + return nil, cleanup, fmt.Errorf("claude-cli: cannot derive structured-output schema: %w", err) + } + if len(schema) > 0 { + args = append(args, "--json-schema", string(schema)) + } + return args, cleanup, nil +} + +func claudeCLIModel(model string) string { + model = strings.TrimSpace(model) + if model == "claude" { + return "" + } + return ai.NormalizeModelForBackend(ai.BackendClaudeCLI, model) +} + +func cliClaudePermissionMode(mode api.PermissionMode) string { + switch mode { + case "", api.PermissionDefault: + return "" + case api.PermissionAuto: + return "auto" + case api.PermissionAcceptEdits: + return "acceptEdits" + case api.PermissionBypass: + return "bypassPermissions" + case api.PermissionDontAsk: + return "dontAsk" + case api.PermissionPlan: + return "plan" + default: + return string(mode) + } +} + +func claudeEntryEvents(entry claude.HistoryEntry, fallbackModel string) []ai.Event { + model := firstNonEmpty(entry.Message.Model, fallbackModel) + var out []ai.Event + for _, block := range entry.Message.Content { + switch block.Type { + case claude.ContentTypeText: + if block.Text != "" && entry.Message.Role == claude.MessageRoleAssistant { + out = append(out, ai.Event{Kind: ai.EventText, Text: block.Text, SessionID: entry.SessionID, Model: model}) + } + case claude.ContentTypeThinking: + if block.Thinking != "" { + out = append(out, ai.Event{Kind: ai.EventThinking, Text: block.Thinking, SessionID: entry.SessionID, Model: model}) + } + case claude.ContentTypeToolUse: + input := map[string]any{} + if len(block.Input) > 0 { + _ = json.Unmarshal(block.Input, &input) + } + ev := ai.Event{Kind: ai.EventToolUse, Tool: block.Name, Input: input, ToolCallID: block.ID, SessionID: entry.SessionID, Model: model, Raw: entry} + switch block.Name { + case "SessionInit": + ev.Kind = ai.EventSystem + ev.Tool = "SessionInit" + if session, _ := input["session_id"].(string); session != "" { + ev.SessionID = session + } + case "Result": + ev = claudeResultEvent(entry, block, input, model) + case "ApiError": + ev.Kind = ai.EventError + ev.Error = firstString(input, "error", "message", "result") + ev.Raw = entry + } + out = append(out, ev) + case claude.ContentTypeToolResult: + ev := ai.Event{Kind: ai.EventToolResult, Text: string(block.Content), ToolCallID: block.ToolUseID, Success: !block.IsError, SessionID: entry.SessionID, Model: model, Raw: entry} + out = append(out, ev) + } + } + return out +} + +func claudeResultEvent(entry claude.HistoryEntry, block claude.ContentBlock, input map[string]any, model string) ai.Event { + ev := ai.Event{ + Kind: ai.EventResult, + Tool: "Result", + Success: !truthy(input["is_error"]), + SessionID: entry.SessionID, + Model: model, + Input: input, + Raw: entry, + } + if result, _ := input["result"].(string); result != "" { + ev.Text = result + } + if ev.SessionID == "" { + if session, _ := input["session_id"].(string); session != "" { + ev.SessionID = session + } + } + if errText := firstString(input, "error", "result"); !ev.Success && errText != "" { + ev.Error = errText + } + if cost, ok := number(input["total_cost_usd"]); ok { + ev.CostUSD = cost + } + if entry.Message.Usage != nil { + ev.Usage = &ai.Usage{ + InputTokens: entry.Message.Usage.InputTokens, + OutputTokens: entry.Message.Usage.OutputTokens, + CacheReadTokens: entry.Message.Usage.CacheReadInputTokens, + CacheWriteTokens: entry.Message.Usage.CacheCreationInputTokens, + } + } + if len(block.Input) > 0 { + ev.StructuredData = structuredJSONFromResult(input) + } + return ev +} + +func structuredJSONFromResult(input map[string]any) json.RawMessage { + result, _ := input["result"].(string) + if strings.TrimSpace(result) == "" { + return nil + } + var js json.RawMessage + if json.Unmarshal([]byte(result), &js) == nil { + return js + } + return nil +} + +func firstString(input map[string]any, keys ...string) string { + for _, key := range keys { + if v, ok := input[key].(string); ok && v != "" { + return v + } + } + return "" +} + +func truthy(v any) bool { + switch x := v.(type) { + case bool: + return x + case string: + return x == "true" + default: + return false + } +} + +func number(v any) (float64, bool) { + switch x := v.(type) { + case float64: + return x, true + case float32: + return float64(x), true + case int: + return float64(x), true + case int64: + return float64(x), true + case json.Number: + f, err := x.Float64() + return f, err == nil + default: + return 0, false + } +} diff --git a/pkg/ai/provider/claude_cli_test.go b/pkg/ai/provider/claude_cli_test.go new file mode 100644 index 0000000..6150650 --- /dev/null +++ b/pkg/ai/provider/claude_cli_test.go @@ -0,0 +1,151 @@ +package provider + +import ( + "encoding/json" + "reflect" + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" +) + +func TestBuildClaudeCLIArgs(t *testing.T) { + req := ai.Request{ + Prompt: api.Prompt{ + User: "hello", + System: "system prompt", + AppendSystem: "append system", + SchemaJSON: json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string","minLength":2}}}`), + }, + Model: api.Model{Effort: api.EffortHigh}, + Budget: api.Budget{Cost: 1.25}, + SessionID: "sess-1", + Memory: api.Memory{ + Skills: []string{"/skills/a", " "}, + SkipSkills: true, + Bare: true, + }, + Permissions: api.Permissions{ + Mode: api.PermissionAcceptEdits, + Tools: api.Tools{ + Allow: []string{"Read", "Grep"}, + Deny: []string{"Bash"}, + }, + MCP: api.MCP{Disabled: true}, + }, + } + + args, cleanup, err := buildClaudeCLIArgs("claude-sonnet-5", req) + if err != nil { + t.Fatalf("buildClaudeCLIArgs: %v", err) + } + defer cleanup() + wantPrefix := []string{"-p", "--verbose", "--output-format", "stream-json"} + if !reflect.DeepEqual(args[:len(wantPrefix)], wantPrefix) { + t.Fatalf("args prefix = %v, want %v", args[:len(wantPrefix)], wantPrefix) + } + requireFlagValue(t, args, "--model", "claude-sonnet-5") + requireFlagValue(t, args, "--system-prompt", "system prompt") + requireFlagValue(t, args, "--append-system-prompt", "append system") + requireFlagValue(t, args, "--resume", "sess-1") + requireFlagValue(t, args, "--effort", "high") + requireFlagValue(t, args, "--max-budget-usd", "1.25") + requireFlagValue(t, args, "--permission-mode", "acceptEdits") + requireFlagValue(t, args, "--allowedTools", "Read,Grep") + requireFlagValue(t, args, "--disallowedTools", "Bash") + requireFlagValue(t, args, "--plugin-dir", "/skills/a") + requireFlagValue(t, args, "--mcp-config", `{"mcpServers":{}}`) + requireHasArg(t, args, "--strict-mcp-config") + requireHasArg(t, args, "--disable-slash-commands") + requireHasArg(t, args, "--bare") + schema := flagValue(t, args, "--json-schema") + if !json.Valid([]byte(schema)) { + t.Fatalf("--json-schema = %q, want valid JSON", schema) + } + var decoded map[string]any + if err := json.Unmarshal([]byte(schema), &decoded); err != nil { + t.Fatalf("unmarshal schema: %v", err) + } + answer := decoded["properties"].(map[string]any)["answer"].(map[string]any) + if _, ok := answer["minLength"]; ok { + t.Fatalf("--json-schema should strip minLength key, got %s", schema) + } + if answer["description"] == "" { + t.Fatalf("--json-schema should describe removed constraints, got %s", schema) + } +} + +func TestClaudeEntryEventsMapsTextThinkingAndResult(t *testing.T) { + entry := claude.HistoryEntry{ + SessionID: "sess-1", + Message: claude.Message{ + Model: "sonnet", + Role: claude.MessageRoleAssistant, + Usage: &claude.Usage{InputTokens: 11, OutputTokens: 7, CacheReadInputTokens: 3, CacheCreationInputTokens: 2}, + Content: []claude.ContentBlock{ + {Type: claude.ContentTypeThinking, Thinking: "thinking"}, + {Type: claude.ContentTypeText, Text: "hello"}, + { + Type: claude.ContentTypeToolUse, + ID: "result-1", + Name: "Result", + Input: json.RawMessage(`{"result":"{\"answer\":\"42\"}","total_cost_usd":0.5}`), + }, + }, + }, + } + + events := claudeEntryEvents(entry, "fallback") + if len(events) != 3 { + t.Fatalf("events = %+v, want 3", events) + } + if events[0].Kind != ai.EventThinking || events[0].Text != "thinking" { + t.Fatalf("thinking event = %+v", events[0]) + } + if events[1].Kind != ai.EventText || events[1].Text != "hello" { + t.Fatalf("text event = %+v", events[1]) + } + result := events[2] + if result.Kind != ai.EventResult || !result.Success || result.Text != `{"answer":"42"}` { + t.Fatalf("result event = %+v", result) + } + if string(result.StructuredData) != `{"answer":"42"}` { + t.Fatalf("StructuredData = %s, want answer JSON", result.StructuredData) + } + if result.Usage == nil || result.Usage.InputTokens != 11 || result.Usage.OutputTokens != 7 { + t.Fatalf("usage = %+v", result.Usage) + } +} + +func requireHasArg(t *testing.T, args []string, want string) { + t.Helper() + for _, arg := range args { + if arg == want { + return + } + } + t.Fatalf("args %v do not contain %q", args, want) +} + +func requireFlagValue(t *testing.T, args []string, flag, want string) { + t.Helper() + got := flagValue(t, args, flag) + if got != want { + t.Fatalf("%s = %q, want %q (args: %v)", flag, got, want, args) + } +} + +func flagValue(t *testing.T, args []string, flag string) string { + t.Helper() + for i, arg := range args { + if arg == flag { + if i+1 >= len(args) { + t.Fatalf("%s has no value in args %v", flag, args) + } + return args[i+1] + } + } + t.Fatalf("args %v do not contain flag %s", args, flag) + return "" +} diff --git a/pkg/ai/provider/claudeagent/agent.ts b/pkg/ai/provider/claudeagent/agent.ts index 8b0216f..98b7e6d 100644 --- a/pkg/ai/provider/claudeagent/agent.ts +++ b/pkg/ai/provider/claudeagent/agent.ts @@ -8,7 +8,7 @@ // maxTurns, maxBudgetUsd, permissionMode, resume, approvalMode, // outputSchema} // -> reply {ok:true} -// prompt {text} -> reply {accepted:true} +// prompt {text, attachments?} -> reply {accepted:true} // interrupt -> reply {} // shutdown -> reply {} then exit // server -> client notifications: @@ -62,6 +62,9 @@ interface InitializeParams { // structured-output target. Present => the SDK is asked for validated JSON // (options.outputFormat) and every turn's result carries structured_output. outputSchema?: Record; + // monitorUrl is the captain serve base URL session-monitoring lifecycle + // hooks POST to. Empty/absent disables monitoring hook injection. + monitorUrl?: string; } type JsonRpcId = number | string | null; @@ -143,10 +146,43 @@ class TurnQueue implements AsyncIterable { private waiters: ((r: IteratorResult) => void)[] = []; private ended = false; - push(text: string) { + push(params: PromptParams) { + const content: Exclude = []; + if (params.text) { + content.push({ type: "text", text: params.text }); + } + for (const attachment of params.attachments ?? []) { + if (attachment.mediaType === "application/pdf") { + content.push({ + type: "document", + source: { + type: "base64", + media_type: "application/pdf", + data: attachment.data, + }, + title: attachment.filename || undefined, + }); + } else if (isClaudeImageMediaType(attachment.mediaType)) { + content.push({ + type: "image", + source: { + type: "base64", + media_type: attachment.mediaType, + data: attachment.data, + }, + }); + } else { + throw new Error( + `unsupported attachment media type: ${attachment.mediaType}`, + ); + } + } const msg: SDKUserMessage = { type: "user", - message: { role: "user", content: text }, + message: { + role: "user", + content, + }, parent_tool_use_id: null, session_id: "", }; @@ -188,6 +224,31 @@ class TurnQueue implements AsyncIterable { let turns: TurnQueue | null = null; let activeQuery: Query | null = null; +interface PromptAttachment { + mediaType: string; + data: string; + filename?: string; +} + +type ClaudeImageMediaType = + | "image/png" + | "image/jpeg" + | "image/gif" + | "image/webp"; + +function isClaudeImageMediaType( + mediaType: string, +): mediaType is ClaudeImageMediaType { + return ["image/png", "image/jpeg", "image/gif", "image/webp"].includes( + mediaType, + ); +} + +interface PromptParams { + text?: string; + attachments?: PromptAttachment[]; +} + function buildOptions(params: InitializeParams): Options { // brokered: the host vets each tool over the can_use_tool round-trip, so the // SDK must consult canUseTool rather than auto-approving. bypassPermissions / @@ -228,6 +289,44 @@ function buildOptions(params: InitializeParams): Options { }, }; + // Session-monitoring lifecycle hooks: fire-and-forget POSTs to captain + // serve so the session appears in the database in real time. A monitoring + // failure (serve down, slow) must never block or slow the agent turn. + if (params.monitorUrl) { + const monitorUrl = params.monitorUrl.replace(/\/+$/, ""); + const monitorEvents = [ + "SessionStart", + "UserPromptSubmit", + "Stop", + "SubagentStop", + "SessionEnd", + ] as const; + const hooks = options.hooks as unknown as Record; + for (const event of monitorEvents) { + hooks[event] = [ + { + hooks: [ + async (input: Record) => { + fetch(`${monitorUrl}/api/captain/hooks/claude`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + event, + sessionId: input.session_id, + transcriptPath: input.transcript_path, + cwd: input.cwd, + detail: input.source ?? input.reason ?? undefined, + }), + signal: AbortSignal.timeout(1000), + }).catch(() => {}); + return {}; + }, + ], + }, + ]; + } + } + // appendSystemPrompt is not a top-level Options field; it must ride on the // claude_code preset. A custom systemPrompt string replaces the default. if (params.systemPrompt && params.appendSystemPrompt) { @@ -307,13 +406,17 @@ function handleInitialize(id: JsonRpcId, params: InitializeParams) { } } -function handlePrompt(id: JsonRpcId, params: { text?: string }) { +function handlePrompt(id: JsonRpcId, params: PromptParams) { if (!turns) { replyError(id, -32002, "not initialized"); return; } - turns.push(params.text || ""); - reply(id, { accepted: true }); + try { + turns.push(params); + reply(id, { accepted: true }); + } catch (err) { + replyError(id, -32602, `invalid prompt: ${(err as Error)?.message || err}`); + } } async function handleInterrupt(id: JsonRpcId) { @@ -454,7 +557,7 @@ rl.on("line", (line) => { handleInitialize(id, (req.params as InitializeParams) || {}); break; case "prompt": - handlePrompt(id, (req.params as { text?: string }) || {}); + handlePrompt(id, (req.params as PromptParams) || {}); break; case "interrupt": handleInterrupt(id); diff --git a/pkg/ai/provider/claudeagent/attachments_ginkgo_test.go b/pkg/ai/provider/claudeagent/attachments_ginkgo_test.go new file mode 100644 index 0000000..12c369f --- /dev/null +++ b/pkg/ai/provider/claudeagent/attachments_ginkgo_test.go @@ -0,0 +1,43 @@ +package claudeagent + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" +) + +var _ = Describe("Claude Agent prompt parameters", func() { + It("derives the required SDK version from the embedded package manifest", func() { + version, err := requiredSDKVersion() + Expect(err).NotTo(HaveOccurred()) + Expect(version).To(Equal("0.3.210")) + }) + + It("materializes the structured attachment bridge source", func() { + directory, err := prepareAgentDir() + Expect(err).NotTo(HaveOccurred()) + content, err := os.ReadFile(filepath.Join(directory, "agent.ts")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(content)).To(ContainSubstring("attachments?: PromptAttachment[]")) + }) + + It("encodes prepared image and PDF data as ordered structured inputs", func() { + image := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"}. + WithPreparedContent(api.AttachmentContent{Bytes: []byte("image")}) + pdf := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "application/pdf"}. + WithPreparedContent(api.AttachmentContent{Bytes: []byte("pdf")}) + + params, err := buildPromptParams(ai.Request{Prompt: api.Prompt{User: "inspect", Attachments: []api.AttachmentRef{image, pdf}}}) + Expect(err).NotTo(HaveOccurred()) + Expect(params.Text).To(Equal("inspect")) + Expect(params.Attachments).To(Equal([]promptAttachment{ + {MediaType: "image/png", Data: "aW1hZ2U="}, + {MediaType: "application/pdf", Data: "cGRm"}, + })) + }) +}) diff --git a/pkg/ai/provider/claudeagent/attachments_suite_test.go b/pkg/ai/provider/claudeagent/attachments_suite_test.go new file mode 100644 index 0000000..7712f0f --- /dev/null +++ b/pkg/ai/provider/claudeagent/attachments_suite_test.go @@ -0,0 +1,13 @@ +package claudeagent + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAttachments(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Claude Agent Attachments Suite") +} diff --git a/pkg/ai/provider/claudeagent/mapping_test.go b/pkg/ai/provider/claudeagent/mapping_test.go index 9b52296..d30e363 100644 --- a/pkg/ai/provider/claudeagent/mapping_test.go +++ b/pkg/ai/provider/claudeagent/mapping_test.go @@ -10,7 +10,7 @@ import ( "github.com/stretchr/testify/require" ) -const testModel = "claude-agent-sonnet" +const testModel = "claude-sonnet-5" func mustMap(t *testing.T, method, params string) ai.Event { t.Helper() @@ -215,13 +215,16 @@ func TestDecodeUsage(t *testing.T) { func TestAliasModel(t *testing.T) { cases := map[string]string{ - "claude-agent-sonnet": "sonnet", - "claude-agent-opus": "opus", - "claude-code-haiku": "haiku", - "claude-sonnet-4-5": "claude-sonnet-4-5", - "claude-agent-": "sonnet", - "": "sonnet", - "claude-agent-opus-4-1": "opus-4-1", + "claude-agent-sonnet": "claude-sonnet-5", + "claude-agent-opus": "claude-opus-5", + "claude-code-haiku": "claude-haiku-4-5", + "claude-sonnet-4-5": "claude-sonnet-4-6", + // A bare family-less token resolves through the provider's emptyFamily, + // which is opus (see registry/providers.go). The "" case is short-circuited + // by aliasModel itself and so still answers sonnet. + "claude-agent-": "claude-opus-5", + "": "claude-sonnet-5", + "claude-agent-opus-4-1": "claude-opus-4-8", } for in, want := range cases { assert.Equalf(t, want, aliasModel(in), "aliasModel(%q)", in) diff --git a/pkg/ai/provider/claudeagent/package.json b/pkg/ai/provider/claudeagent/package.json index 4add5c4..0695c3f 100644 --- a/pkg/ai/provider/claudeagent/package.json +++ b/pkg/ai/provider/claudeagent/package.json @@ -3,7 +3,7 @@ "private": true, "type": "module", "dependencies": { - "@anthropic-ai/claude-agent-sdk": "latest", + "@anthropic-ai/claude-agent-sdk": "0.3.210", "tsx": "latest", "typescript": "latest" } diff --git a/pkg/ai/provider/claudeagent/provider.go b/pkg/ai/provider/claudeagent/provider.go index d1e70ce..32c49fc 100644 --- a/pkg/ai/provider/claudeagent/provider.go +++ b/pkg/ai/provider/claudeagent/provider.go @@ -47,7 +47,7 @@ const ( const methodCanUseTool = "can_use_tool" const ( - defaultModel = "claude-agent-sonnet" + defaultModel = "claude-sonnet-5" // initTimeout bounds the initialize handshake (after provisioning). The npm // install / tsx cold start happen synchronously before this window. initTimeout = 2 * time.Minute @@ -126,6 +126,7 @@ func New(cfg ai.Config) (*Provider, error) { if model == "" { model = defaultModel } + model = ai.NormalizeModelForBackend(ai.BackendClaudeAgent, model) ctx, cancel := context.WithCancel(context.Background()) return &Provider{ model: model, @@ -160,8 +161,18 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e success = true sawResult bool lastErr string + outcome *ai.TerminalOutcome + outcomeErr error ) for ev := range events { + if outcomeErr == nil { + parsed, parseErr := ai.TerminalOutcomeFromEvent(ev) + if parseErr != nil { + outcomeErr = parseErr + } else if parsed != nil { + outcome = parsed + } + } switch ev.Kind { case ai.EventText: text.WriteString(ev.Text) @@ -188,6 +199,9 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e lastErr = ev.Error } } + if outcomeErr != nil { + return nil, fmt.Errorf("claude-agent: invalid terminal outcome: %w", outcomeErr) + } if !sawResult && lastErr != "" { return nil, fmt.Errorf("%w: %s", ai.ErrCLIExecutionFailed, lastErr) @@ -197,15 +211,19 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e } resp := &ai.Response{ - Text: text.String(), - Model: p.model, - Backend: ai.BackendClaudeAgent, - Usage: usage, - Duration: time.Since(start), + Text: text.String(), + Model: p.model, + Backend: ai.BackendClaudeAgent, + Usage: usage, + Duration: time.Since(start), + TerminalOutcome: outcome, } if sessionID != "" { resp.Raw = map[string]any{"session_id": sessionID} } + if outcome != nil { + return resp, nil + } if req.Prompt.Schema != nil { if err := ai.BindStructuredOutput(req.Prompt.Schema, structured); err != nil { return nil, err @@ -265,7 +283,7 @@ func (p *Provider) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai // Prompt.SchemaJSON), or nil for a text-mode request. A non-struct target fails // loudly rather than silently dropping the schema. func requestSchemaJSON(req ai.Request) (json.RawMessage, error) { - schema, err := ai.SchemaJSONFor(req.Prompt) + schema, err := ai.SchemaJSONForBackend(ai.BackendClaudeAgent, req.Prompt) if err != nil { return nil, fmt.Errorf("claude-agent: cannot derive structured-output schema: %w", err) } @@ -393,6 +411,14 @@ func (p *Provider) initializeParams(req ai.Request) initializeParams { resume = p.cfg.SessionID } + // Prefer the per-request budget (like claude-cli, which reads req.Budget.Cost) + // and fall back to the config default, so both Claude backends resolve the + // ceiling from the same source (finding A4). + maxBudget := req.Budget.Cost + if maxBudget == 0 { + maxBudget = p.cfg.Budget.Cost + } + return initializeParams{ Cwd: req.Cwd(), Model: aliasModel(p.model), @@ -400,12 +426,23 @@ func (p *Provider) initializeParams(req ai.Request) initializeParams { AppendSystemPrompt: req.Prompt.AppendSystem, AllowedTools: allowed, MaxTurns: req.Budget.MaxTurns, - MaxBudgetUsd: p.cfg.Budget.Cost, + MaxBudgetUsd: maxBudget, PermissionMode: mode, Resume: resume, ApprovalMode: approvalMode, OutputSchema: p.sessionSchema, + MonitorURL: monitorHooksURL(req), + } +} + +// monitorHooksURL resolves the captain serve URL the SDK's session-monitoring +// hooks deliver lifecycle events to; empty disables injection (bare runs, +// CAPTAIN_MONITOR_HOOKS=off). +func monitorHooksURL(req ai.Request) string { + if !api.MonitorHooksEnabled(req) { + return "" } + return api.ServeBaseURL() } type initializeParams struct { @@ -420,6 +457,7 @@ type initializeParams struct { Resume string `json:"resume,omitempty"` ApprovalMode string `json:"approvalMode,omitempty"` OutputSchema json.RawMessage `json:"outputSchema,omitempty"` + MonitorURL string `json:"monitorUrl,omitempty"` } func (p *Provider) setInitResult(err error) { diff --git a/pkg/ai/provider/claudeagent/provider_test.go b/pkg/ai/provider/claudeagent/provider_test.go index 7de86ff..8c54c31 100644 --- a/pkg/ai/provider/claudeagent/provider_test.go +++ b/pkg/ai/provider/claudeagent/provider_test.go @@ -68,6 +68,7 @@ func runFakeServer() { // initHadSchema records whether the host sent an outputSchema on initialize, // so the "structured" turn can prove the Go→TS schema wiring end to end. initHadSchema := false + promptCount := 0 scanner := bufio.NewScanner(os.Stdin) scanner.Buffer(make([]byte, 64*1024), 4*1024*1024) @@ -98,15 +99,33 @@ func runFakeServer() { "session_id": "fake-sess", "model": "claude-sonnet-4-5", "tools": []string{"Read", "Bash"}, }}) case "prompt": + promptCount++ enc(map[string]any{"jsonrpc": "2.0", "id": id(frame.ID), "result": map[string]any{"accepted": true}}) enc(map[string]any{"jsonrpc": "2.0", "method": "message/text", "params": map[string]any{"text": "hi from fake"}}) switch mode { + case "steer": + if promptCount == 2 { + completed("first prompt complete") + completed("steered prompt complete") + } case "approval": // Ask the host to vet a Bash tool use; the turn completes when the // host replies (handled above). enc(map[string]any{"jsonrpc": "2.0", "id": "perm-1", "method": "can_use_tool", "params": map[string]any{ "tool": "Bash", "input": map[string]any{"command": "ls"}, "tool_use_id": "tu1", }}) + case "plan-approval": + // A plan-mode turn ending in ExitPlanMode: the tool_use streams first + // (the SDK yields the assistant message before executing the tool), + // then the permission check; the turn completes when the host replies. + enc(map[string]any{"jsonrpc": "2.0", "method": "message/tool_use", "params": map[string]any{ + "tool": "ExitPlanMode", "id": "tu-plan", + "input": map[string]any{"plan": "1. change the seam", "planFilePath": "/repo/.claude/plans/x.md"}, + }}) + enc(map[string]any{"jsonrpc": "2.0", "id": "perm-plan", "method": "can_use_tool", "params": map[string]any{ + "tool": "ExitPlanMode", "tool_use_id": "tu-plan", + "input": map[string]any{"plan": "1. change the seam", "planFilePath": "/repo/.claude/plans/x.md"}, + }}) case "hang": // Emit nothing more; wait for the interrupt control request. case "structured": @@ -161,12 +180,12 @@ func withFakeAgentProcessEnv(t *testing.T, env map[string]string) { func TestProvider_StreamLifecycle(t *testing.T) { withFakeAgentProcess(t) - p, err := New(ai.Config{Model: api.Model{Name: "claude-agent-sonnet"}}) + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}}) require.NoError(t, err) t.Cleanup(func() { _ = p.Close() }) assert.Equal(t, ai.BackendClaudeAgent, p.GetBackend()) - assert.Equal(t, "claude-agent-sonnet", p.GetModel()) + assert.Equal(t, "claude-sonnet-5", p.GetModel()) ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() @@ -197,7 +216,7 @@ func TestProvider_StreamLifecycle(t *testing.T) { func TestProvider_ExecuteCoalesce(t *testing.T) { withFakeAgentProcess(t) - p, err := New(ai.Config{Model: api.Model{Name: "claude-agent-sonnet"}}) + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}}) require.NoError(t, err) t.Cleanup(func() { _ = p.Close() }) @@ -225,7 +244,7 @@ type companyInfo struct { func TestProvider_StructuredOutput(t *testing.T) { withFakeAgentProcessEnv(t, map[string]string{fakeServerEnv: "1", fakeModeEnv: "structured"}) - p, err := New(ai.Config{Model: api.Model{Name: "claude-agent-sonnet"}}) + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}}) require.NoError(t, err) t.Cleanup(func() { _ = p.Close() }) @@ -263,7 +282,7 @@ func TestRequestSchemaJSON(t *testing.T) { require.NoError(t, err) require.NotEmpty(t, raw) - p := &Provider{model: "claude-agent-sonnet", sessionSchema: raw} + p := &Provider{model: "claude-sonnet-5", sessionSchema: raw} ip := p.initializeParams(ai.Request{Prompt: api.Prompt{User: "x", Schema: &plan{}}}) require.NotEmpty(t, ip.OutputSchema, "outputSchema should be set from the session schema") @@ -276,6 +295,20 @@ func TestRequestSchemaJSON(t *testing.T) { assert.Contains(t, props, "steps") }) + t.Run("raw schema is sanitized for anthropic structured output", func(t *testing.T) { + raw, err := requestSchemaJSON(ai.Request{Prompt: api.Prompt{ + SchemaJSON: json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string","maxLength":12}}}`), + }}) + require.NoError(t, err) + var decoded map[string]any + require.NoError(t, json.Unmarshal(raw, &decoded)) + answer := decoded["properties"].(map[string]any)["answer"].(map[string]any) + if _, ok := answer["maxLength"]; ok { + t.Fatalf("request schema should strip maxLength key, got %s", raw) + } + assert.NotEmpty(t, answer["description"]) + }) + t.Run("non-struct target fails loudly", func(t *testing.T) { _, err := requestSchemaJSON(ai.Request{Prompt: api.Prompt{Schema: "not a struct"}}) require.Error(t, err) @@ -285,7 +318,7 @@ func TestRequestSchemaJSON(t *testing.T) { func TestProvider_MultiTurnSerialized(t *testing.T) { withFakeAgentProcess(t) - p, err := New(ai.Config{Model: api.Model{Name: "claude-agent-sonnet"}}) + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}}) require.NoError(t, err) t.Cleanup(func() { _ = p.Close() }) @@ -323,7 +356,7 @@ func TestProvider_CanUseTool(t *testing.T) { } // CanUseTool is a runtime concern, set on the provider's Config (not the request). - p, err := New(ai.Config{Model: api.Model{Name: "claude-agent-sonnet"}, CanUseTool: canUseTool}) + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}, CanUseTool: canUseTool}) require.NoError(t, err) t.Cleanup(func() { _ = p.Close() }) @@ -371,6 +404,62 @@ func TestProvider_CanUseTool(t *testing.T) { assert.Contains(t, resultText, "ls -la") } +// TestProvider_PlanModeExitPlanModeAutoDenied verifies the shared plan-mode +// terminal policy: in a plan-only run the ExitPlanMode can_use_tool is answered +// by the provider itself (deny + guidance) — the host broker is never invoked, +// no EventPermission is surfaced, and the already-streamed tool_use still +// yields the plan terminal outcome. +func TestProvider_PlanModeExitPlanModeAutoDenied(t *testing.T) { + withFakeAgentProcessEnv(t, map[string]string{fakeServerEnv: "1", fakeModeEnv: "plan-approval"}) + + canUseTool := func(_ context.Context, r ai.PermissionRequest) (ai.PermissionDecision, error) { + t.Errorf("CanUseTool must not be invoked for ExitPlanMode in plan mode, got %s", r.Tool) + return ai.PermissionDecision{Allow: true}, nil + } + + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}, CanUseTool: canUseTool}) + require.NoError(t, err) + t.Cleanup(func() { _ = p.Close() }) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + events, err := p.ExecuteStream(ctx, ai.Request{ + Prompt: api.Prompt{User: "plan it"}, + Permissions: api.Permissions{Mode: api.PermissionPlan}, + }) + require.NoError(t, err) + + var outcome *ai.TerminalOutcome + var result *ai.Event + for ev := range events { + switch ev.Kind { + case ai.EventPermission: + t.Errorf("plan-terminal ExitPlanMode must not surface an EventPermission, got %s", ev.Tool) + case ai.EventToolUse: + parsed, perr := ai.TerminalOutcomeFromEvent(ev) + require.NoError(t, perr) + if parsed != nil { + outcome = parsed + } + case ai.EventResult: + cp := ev + result = &cp + } + } + + // The streamed tool_use still derives the plan terminal outcome. + require.NotNil(t, outcome, "expected the ExitPlanMode tool_use to stream") + assert.Equal(t, ai.TerminalOutcomePlan, outcome.Kind) + assert.Equal(t, "/repo/.claude/plans/x.md", outcome.Plan.Path) + + // The provider answered the can_use_tool itself with the deny + guidance. + require.NotNil(t, result, "expected a terminal result event") + resultText, _ := result.Input["result"].(string) + assert.Contains(t, resultText, `"allow":false`) + assert.Contains(t, resultText, "Plan captured for human review") +} + // TestProvider_InterruptNoKill verifies a cancelled turn is ended by the // interrupt control request — the fake records receiving it — rather than by // killing the process. @@ -382,7 +471,7 @@ func TestProvider_InterruptNoKill(t *testing.T) { fakeMarkerEnv: marker, }) - p, err := New(ai.Config{Model: api.Model{Name: "claude-agent-sonnet"}}) + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}}) require.NoError(t, err) t.Cleanup(func() { _ = p.Close() }) @@ -408,3 +497,31 @@ func TestProvider_InterruptNoKill(t *testing.T) { // written the marker — proof the graceful interrupt reached it. require.FileExists(t, marker) } + +func TestProvider_SteerKeepsTurnOpenForEveryAcceptedPrompt(t *testing.T) { + withFakeAgentProcessEnv(t, map[string]string{ + fakeServerEnv: "1", + fakeModeEnv: "steer", + }) + + p, err := New(ai.Config{Model: api.Model{Name: "claude-sonnet-5"}}) + require.NoError(t, err) + t.Cleanup(func() { _ = p.Close() }) + + events, err := p.ExecuteStream(context.Background(), ai.Request{Prompt: api.Prompt{User: "first"}}) + require.NoError(t, err) + for event := range events { + if event.Kind == ai.EventText { + break + } + } + require.NoError(t, p.Steer(context.Background(), ai.Request{Prompt: api.Prompt{User: "btw"}})) + + results := 0 + for event := range events { + if event.Kind == ai.EventResult { + results++ + } + } + assert.Equal(t, 2, results) +} diff --git a/pkg/ai/provider/claudeagent/runner.go b/pkg/ai/provider/claudeagent/runner.go index a199f9f..0578d25 100644 --- a/pkg/ai/provider/claudeagent/runner.go +++ b/pkg/ai/provider/claudeagent/runner.go @@ -4,11 +4,14 @@ import ( "crypto/sha256" _ "embed" "encoding/hex" + "encoding/json" "fmt" "os" "os/exec" "path/filepath" "strings" + + "github.com/flanksource/captain/pkg/ai" ) //go:embed agent.ts @@ -61,15 +64,23 @@ func contentHash(data []byte) string { return hex.EncodeToString(h[:]) } -// ensureDependencies installs node_modules the first time (the SDK package -// missing is the signal). It fails loudly when npm is unavailable rather than -// silently degrading — the provider cannot run without the SDK. +// ensureDependencies installs the pinned SDK when it is missing or its version +// differs. It fails loudly when npm is unavailable because the provider cannot +// run without the exact bridge contract declared in package.json. func ensureDependencies(agentDir string) error { - sdkDir := filepath.Join(agentDir, "node_modules", "@anthropic-ai", "claude-agent-sdk") - if _, err := os.Stat(sdkDir); err == nil { - return nil + requiredVersion, err := requiredSDKVersion() + if err != nil { + return err + } + sdkPackage := filepath.Join(agentDir, "node_modules", "@anthropic-ai", "claude-agent-sdk", "package.json") + if data, err := os.ReadFile(sdkPackage); err == nil { + var installed struct { + Version string `json:"version"` + } + if json.Unmarshal(data, &installed) == nil && installed.Version == requiredVersion { + return nil + } } - npmPath, err := exec.LookPath("npm") if err != nil { return fmt.Errorf("npm not found in PATH (required to install the Claude Agent SDK): %w", err) @@ -85,6 +96,20 @@ func ensureDependencies(agentDir string) error { return nil } +func requiredSDKVersion() (string, error) { + var manifest struct { + Dependencies map[string]string `json:"dependencies"` + } + if err := json.Unmarshal([]byte(agentPackageJSON), &manifest); err != nil { + return "", fmt.Errorf("parse embedded Claude Agent package manifest: %w", err) + } + version := strings.TrimSpace(manifest.Dependencies["@anthropic-ai/claude-agent-sdk"]) + if version == "" { + return "", fmt.Errorf("embedded Claude Agent package manifest does not pin @anthropic-ai/claude-agent-sdk") + } + return version, nil +} + // findTsx resolves the tsx runner, preferring the version installed into the // agent dir's node_modules over a global one. func findTsx(agentDir string) (string, error) { @@ -115,15 +140,13 @@ func nestingEnvOverrides(environ []string) map[string]string { return overrides } -// aliasModel strips the captain backend prefix from a model id and passes the -// remainder (e.g. "sonnet", "opus", "claude-sonnet-4-5") straight to the SDK, -// which resolves the short aliases via the Claude Code CLI. Kept local so the -// claudeagent package never imports pkg/ai/provider (which imports it back). +// aliasModel is retained as the local model renderer for the agent.ts bridge. +// It accepts legacy backend-prefixed aliases as input compatibility, but returns +// the exact Claude model ID the SDK should receive. func aliasModel(model string) string { - m := strings.TrimPrefix(model, "claude-agent-") - m = strings.TrimPrefix(m, "claude-code-") - if m == "" { - return "sonnet" + m := strings.TrimSpace(model) + if m == "" || m == "claude" { + return "claude-sonnet-5" } - return m + return ai.NormalizeModelForBackend(ai.BackendClaudeAgent, m) } diff --git a/pkg/ai/provider/claudeagent/turn.go b/pkg/ai/provider/claudeagent/turn.go index 89402c2..8e4723d 100644 --- a/pkg/ai/provider/claudeagent/turn.go +++ b/pkg/ai/provider/claudeagent/turn.go @@ -2,12 +2,16 @@ package claudeagent import ( "context" + "encoding/base64" "encoding/json" "fmt" + "os" + "sync" "time" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/ai/provider/jsonrpc" + "github.com/flanksource/captain/pkg/api" ) // turnState carries the channels a single in-flight turn uses to receive mapped @@ -25,16 +29,54 @@ type turnState struct { ctx context.Context canUseTool ai.PermissionFunc + // planMode marks a plan-only turn: ExitPlanMode is then the terminal signal + // and its can_use_tool is answered by the shared policy, never the broker. + planMode bool + + promptMu sync.Mutex + pending int + ended bool } type promptParams struct { - Text string `json:"text"` + Text string `json:"text"` + Attachments []promptAttachment `json:"attachments,omitempty"` +} + +type promptAttachment struct { + MediaType string `json:"mediaType"` + Data string `json:"data"` + Filename string `json:"filename,omitempty"` } func composePrompt(req ai.Request) string { return req.Prompt.User } +func buildPromptParams(req ai.Request) (promptParams, error) { + params := promptParams{Text: composePrompt(req), Attachments: make([]promptAttachment, 0, len(req.Prompt.Attachments))} + for i, attachment := range req.Prompt.Attachments { + content, ok := attachment.PreparedContent() + if !ok { + return promptParams{}, fmt.Errorf("attachment %d (%s) is not prepared", i+1, attachment.ID) + } + data := content.Bytes + if data == nil && content.Path != "" { + var err error + data, err = os.ReadFile(content.Path) + if err != nil { + return promptParams{}, fmt.Errorf("read prepared attachment %s: %w", attachment.ID, err) + } + } + params.Attachments = append(params.Attachments, promptAttachment{ + MediaType: attachment.MediaType, + Data: base64.StdEncoding.EncodeToString(data), + Filename: attachment.Filename, + }) + } + return params, nil +} + // runTurn owns one turn's event channel: it sends the prompt, forwards mapped // notifications, and emits the terminal result (or a loud error on cancellation // / mid-turn process exit). @@ -50,6 +92,8 @@ func (p *Provider) runTurn(ctx context.Context, req ai.Request, events chan ai.E quit: make(chan struct{}), ctx: ctx, canUseTool: p.cfg.CanUseTool, + planMode: req.Permissions.Mode == api.PermissionPlan, + pending: 1, } p.setActive(ts) defer func() { @@ -57,7 +101,16 @@ func (p *Provider) runTurn(ctx context.Context, req ai.Request, events chan ai.E p.clearActive() }() - if _, err := p.rpc.Call(ctx, methodPrompt, promptParams{Text: composePrompt(req)}); err != nil { + if err := ai.ValidateAttachmentCompatibility([]api.Model{{Name: p.model, Backend: api.BackendClaudeAgent}}, req.Prompt.Attachments); err != nil { + emit(ctx, events, ai.Event{Kind: ai.EventError, Error: err.Error(), Model: p.model}) + return + } + params, err := buildPromptParams(req) + if err != nil { + emit(ctx, events, ai.Event{Kind: ai.EventError, Error: err.Error(), Model: p.model}) + return + } + if _, err := p.rpc.Call(ctx, methodPrompt, params); err != nil { emit(ctx, events, ai.Event{Kind: ai.EventError, Error: fmt.Sprintf("claude-agent prompt failed: %v", err), Model: p.model}) return } @@ -109,11 +162,7 @@ func (p *Provider) onNotification(method string, params json.RawMessage) { } } if method == notifyTurnDone || method == notifyTurnError { - select { - case <-ts.term: - default: - close(ts.term) - } + ts.completePrompt() } } @@ -154,7 +203,7 @@ func (p *Provider) handleCanUseTool(params json.RawMessage) (any, *jsonrpc.RPCEr p.activeMu.Lock() ts := p.active p.activeMu.Unlock() - if ts == nil || ts.canUseTool == nil { + if ts == nil { return canUseToolResult{Allow: true, UpdatedInput: in.Input}, nil } @@ -162,6 +211,19 @@ func (p *Provider) handleCanUseTool(params json.RawMessage) (any, *jsonrpc.RPCEr sessionID := p.sessionID p.sessMu.Unlock() + // The plan-mode terminal signal is answered here, never brokered: nothing is + // awaiting a human, so no EventPermission is surfaced either. The tool_use + // already streamed, so the turn still ends in a plan terminal outcome. + if decision, handled := ai.PlanTerminalPermission(ts.planMode, ai.PermissionRequest{ + Tool: in.Tool, Input: in.Input, ToolUseID: in.ToolUseID, SessionID: sessionID, + }); handled { + return canUseToolResult{Allow: decision.Allow, Message: decision.Message}, nil + } + + if ts.canUseTool == nil { + return canUseToolResult{Allow: true, UpdatedInput: in.Input}, nil + } + p.deliver(ts, ai.Event{ Kind: ai.EventPermission, Tool: in.Tool, @@ -197,13 +259,39 @@ func (p *Provider) deliver(ts *turnState, ev ai.Event) { } } -func (p *Provider) interrupt() { +func (p *Provider) Steer(ctx context.Context, req api.Spec) error { + p.activeMu.Lock() + ts := p.active + p.activeMu.Unlock() + if ts == nil || !ts.addPrompt() { + return fmt.Errorf("claude-agent: no active turn to steer") + } + params, err := buildPromptParams(req) + if err != nil { + ts.completePrompt() + return err + } + if _, err := p.rpc.Call(ctx, methodPrompt, params); err != nil { + ts.completePrompt() + return fmt.Errorf("claude-agent steer failed: %w", err) + } + return nil +} + +func (p *Provider) Interrupt(ctx context.Context) error { if p.rpc == nil { - return + return fmt.Errorf("claude-agent: provider not started") + } + if _, err := p.rpc.Call(ctx, methodInterrupt, nil); err != nil { + return fmt.Errorf("claude-agent interrupt failed: %w", err) } + return nil +} + +func (p *Provider) interrupt() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - _, _ = p.rpc.Call(ctx, methodInterrupt, nil) + _ = p.Interrupt(ctx) } func (p *Provider) setActive(ts *turnState) { @@ -224,6 +312,30 @@ func (p *Provider) rememberSession(id string) { p.sessMu.Unlock() } +func (ts *turnState) addPrompt() bool { + ts.promptMu.Lock() + defer ts.promptMu.Unlock() + if ts.ended { + return false + } + ts.pending++ + return true +} + +func (ts *turnState) completePrompt() { + ts.promptMu.Lock() + defer ts.promptMu.Unlock() + if ts.ended { + return + } + ts.pending-- + if ts.pending > 0 { + return + } + ts.ended = true + close(ts.term) +} + // emit sends ev on events, honouring ctx cancellation. Returns false if ctx was // cancelled before the send completed. func emit(ctx context.Context, events chan ai.Event, ev ai.Event) bool { diff --git a/pkg/ai/provider/cli.go b/pkg/ai/provider/cli.go index b38d360..80da09f 100644 --- a/pkg/ai/provider/cli.go +++ b/pkg/ai/provider/cli.go @@ -6,6 +6,8 @@ import ( "context" "errors" "fmt" + "io" + "os" "os/exec" "strings" @@ -73,53 +75,84 @@ func HandleExitError(exitCode int, stderr string) error { } } -func runCLI(ctx context.Context, command string, stdinData []byte, cwd string, env ...[]string) (stdout []byte, stderr string, err error) { - cmd := exec.CommandContext(ctx, command) +func startCLIStream(ctx context.Context, command string, args []string, stdinData []byte, cwd string, env []string) (*exec.Cmd, io.ReadCloser, *bytes.Buffer, error) { + cmd := exec.CommandContext(ctx, command, args...) if cwd != "" { cmd.Dir = cwd } - if len(env) > 0 && len(env[0]) > 0 { - cmd.Env = env[0] + if len(env) > 0 { + cmd.Env = env + } + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, nil, nil, fmt.Errorf("failed to create stdout pipe: %w", err) } - - // Buffer stdout/stderr instead of using StdoutPipe/StderrPipe: cmd.Wait closes - // those pipes when the process exits, so reading them in a goroutine that races - // Wait fails with "file already closed" (deterministic on Linux). With buffers, - // Wait blocks until the os/exec output copiers finish, making the reads safe. - var stdoutBuf, stderrBuf bytes.Buffer - cmd.Stdout = &stdoutBuf - cmd.Stderr = &stderrBuf - stdin, err := cmd.StdinPipe() if err != nil { - return nil, "", fmt.Errorf("failed to create stdin pipe: %w", err) + return nil, nil, nil, fmt.Errorf("failed to create stdin pipe: %w", err) } - + stderrBuf := &bytes.Buffer{} + cmd.Stderr = stderrBuf if err := cmd.Start(); err != nil { if IsCommandNotFound(err) { - return nil, "", fmt.Errorf("%w: %v", ai.ErrCLINotFound, err) + return nil, nil, nil, fmt.Errorf("%w: %v", ai.ErrCLINotFound, err) } - return nil, "", fmt.Errorf("failed to start %s: %w", command, err) + return nil, nil, nil, fmt.Errorf("failed to start %s: %w", command, err) } - - // Feed stdin in the background so a large prompt cannot deadlock against a - // process that only drains stdin after producing its output. go func() { - _, _ = stdin.Write(stdinData) - _, _ = stdin.Write([]byte("\n")) + if len(stdinData) > 0 { + _, _ = stdin.Write(stdinData) + if stdinData[len(stdinData)-1] != '\n' { + _, _ = stdin.Write([]byte("\n")) + } + } _ = stdin.Close() }() + return cmd, stdout, stderrBuf, nil +} - // CommandContext kills the process when ctx is cancelled, which unblocks Wait. +func finishCLIStream(ctx context.Context, cmd *exec.Cmd, stderrBuf *bytes.Buffer) error { waitErr := cmd.Wait() if ctx.Err() != nil { - return nil, "", fmt.Errorf("%w: context cancelled", ai.ErrTimeout) + return fmt.Errorf("%w: context cancelled", ai.ErrTimeout) + } + if waitErr == nil { + return nil } + return HandleExitError(GetExitCode(waitErr), ParseStderr(stderrBuf.String())) +} - stderrData := stderrBuf.String() - if waitErr != nil { - return nil, stderrData, HandleExitError(GetExitCode(waitErr), ParseStderr(stderrData)) +func commandEnv(setupEnv []string) []string { + if len(setupEnv) == 0 { + return nil + } + merged := os.Environ() + positions := map[string]int{} + for i, item := range merged { + if key, _, ok := strings.Cut(item, "="); ok { + positions[key] = i + } + } + for _, item := range setupEnv { + key, _, ok := strings.Cut(item, "=") + if !ok || key == "" { + continue + } + if idx, ok := positions[key]; ok { + merged[idx] = item + continue + } + positions[key] = len(merged) + merged = append(merged, item) } + return merged +} - return stdoutBuf.Bytes(), stderrData, nil +func emit(ctx context.Context, events chan<- ai.Event, ev ai.Event) bool { + select { + case events <- ev: + return true + case <-ctx.Done(): + return false + } } diff --git a/pkg/ai/provider/cli_test.go b/pkg/ai/provider/cli_test.go index e5e635c..a91124d 100644 --- a/pkg/ai/provider/cli_test.go +++ b/pkg/ai/provider/cli_test.go @@ -33,33 +33,6 @@ func TestParseStderr(t *testing.T) { } } -func TestRunCLIUsesContextDir(t *testing.T) { - cwd := filepath.Join(t.TempDir(), "workspace") - if err := os.MkdirAll(cwd, 0o755); err != nil { - t.Fatalf("mkdir workspace: %v", err) - } - bin := filepath.Join(t.TempDir(), "fake-cli") - if err := os.WriteFile(bin, []byte("#!/bin/sh\npwd\ncat >/dev/null\n"), 0o755); err != nil { - t.Fatalf("write fake cli: %v", err) - } - - stdout, _, err := runCLI(context.Background(), bin, []byte("input"), cwd) - if err != nil { - t.Fatalf("runCLI: %v", err) - } - got, err := filepath.EvalSymlinks(filepath.Clean(strings.TrimSpace(string(stdout)))) - if err != nil { - t.Fatalf("eval pwd: %v", err) - } - want, err := filepath.EvalSymlinks(cwd) - if err != nil { - t.Fatalf("eval cwd: %v", err) - } - if got != want { - t.Errorf("pwd = %q, want %q", got, want) - } -} - func TestGeminiCLIUsesContextDir(t *testing.T) { cwd := filepath.Join(t.TempDir(), "workspace") if err := os.MkdirAll(cwd, 0o755); err != nil { @@ -67,7 +40,12 @@ func TestGeminiCLIUsesContextDir(t *testing.T) { } binDir := t.TempDir() gemini := filepath.Join(binDir, "gemini") - script := "#!/bin/sh\nprintf '{\"text\":\"%s\",\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}\\n' \"$(pwd)\"\ncat >/dev/null\n" + // A fake gemini speaking --output-format stream-json: it reports its own cwd + // as the assistant reply, then closes the run with a terminal result. + script := "#!/bin/sh\ncat >/dev/null\n" + + "printf '{\"type\":\"init\",\"session_id\":\"sess-1\",\"model\":\"gemini-3.5-flash\"}\\n'\n" + + "printf '{\"type\":\"message\",\"role\":\"assistant\",\"content\":\"%s\",\"delta\":true}\\n' \"$(pwd)\"\n" + + "printf '{\"type\":\"result\",\"status\":\"success\",\"stats\":{\"input_tokens\":3,\"output_tokens\":2,\"cached\":2}}\\n'\n" if err := os.WriteFile(gemini, []byte(script), 0o755); err != nil { t.Fatalf("write fake gemini: %v", err) } @@ -80,7 +58,7 @@ func TestGeminiCLIUsesContextDir(t *testing.T) { if err != nil { t.Fatalf("Execute: %v", err) } - got, err := filepath.EvalSymlinks(filepath.Clean(resp.Text)) + got, err := filepath.EvalSymlinks(filepath.Clean(strings.TrimSpace(resp.Text))) if err != nil { t.Fatalf("eval response text: %v", err) } @@ -91,7 +69,7 @@ func TestGeminiCLIUsesContextDir(t *testing.T) { if got != want { t.Errorf("response text = %q, want %q", got, want) } - if resp.Usage.InputTokens != 1 || resp.Usage.OutputTokens != 2 { - t.Errorf("usage = %+v, want input=1 output=2", resp.Usage) + if resp.Usage.InputTokens != 1 || resp.Usage.CacheReadTokens != 2 || resp.Usage.OutputTokens != 2 { + t.Errorf("usage = %+v, want input=1 cacheRead=2 output=2", resp.Usage) } } diff --git a/pkg/ai/provider/cmux/executor.go b/pkg/ai/provider/cmux/executor.go index f66f128..27f77a8 100644 --- a/pkg/ai/provider/cmux/executor.go +++ b/pkg/ai/provider/cmux/executor.go @@ -112,13 +112,33 @@ type run struct { // approvals tracks the in-flight approval handlers spawned by the stall // watchdog so the driver can wait for them (and the EventPermission they emit) // before the event channel is closed. - approvals sync.WaitGroup + approvals sync.WaitGroup + planMode bool + planExitSeen bool + dismissPlanOnce sync.Once + dismissPlanErr error lastSurface WorkspaceRef lastSessionID string lastWorkDir string } +func (r *run) dismissPlanSurface(ctx context.Context, ref WorkspaceRef, sessionID string) error { + r.dismissPlanOnce.Do(func() { + if err := r.client.SendKeySurface(ctx, ref.String(), ref.SurfaceID, "Escape"); err != nil { + r.dismissPlanErr = fmt.Errorf("dismiss claude plan surface for session %s: %w", sessionID, err) + } + }) + return r.dismissPlanErr +} + +func (r *run) dismissCompletedPlan(ctx context.Context, ref WorkspaceRef, sessionID string) error { + if !r.planMode || !r.planExitSeen { + return nil + } + return r.dismissPlanSurface(ctx, ref, sessionID) +} + // readScreen returns the normalized surface contents, or "" if the read failed. func (r *run) readScreen(ctx context.Context, ref WorkspaceRef) string { screen, err := r.client.ReadScreen(ctx, ReadScreenOpts{ @@ -237,7 +257,7 @@ type AgentCommandOpts struct { // --disallowedTools. codex ignores them. AllowedTools []string DisallowedTools []string - // Effort maps onto claude --effort (from Spec.Model.Effort). codex ignores it. + // Effort maps onto claude --effort or Codex's model_reasoning_effort config. Effort api.Effort // Memory drives claude --bare / --disable-slash-commands / --setting-sources // (from Spec.Memory). codex ignores it. @@ -274,6 +294,9 @@ func AgentCommand(opts AgentCommandOpts) string { if opts.Model != "" { tokens = append(tokens, "-m", opts.Model) } + if opts.Effort != api.EffortNone { + tokens = append(tokens, "-c", fmt.Sprintf("model_reasoning_effort=%q", opts.Effort)) + } if extra, ok := opts.Extra.(*api.CodexCmuxOptions); ok && extra != nil { tokens = append(tokens, flagArgs(*extra)...) } @@ -423,14 +446,20 @@ func shellSingleQuote(v string) string { return "'" + strings.ReplaceAll(v, "'", `'\''`) + "'" } -// modelFlag normalizes the configured model into the value claude/codex --model -// expects: the bare agent names ("claude"/"codex") and an empty string carry no -// concrete model, so they yield "" (no flag). +// modelFlag normalizes the configured model into the exact value claude/codex +// --model expects. The bare agent names ("claude"/"codex") and an empty string +// carry no concrete model, so they yield "" (no flag). func modelFlag(agent, model string) string { lower := strings.ToLower(strings.TrimSpace(model)) if lower == "" || lower == agent { return "" } + if agent == "claude" { + return ai.NormalizeModelForBackend(ai.BackendClaudeCmux, model) + } + if agent == "codex" { + return ai.NormalizeModelForBackend(ai.BackendCodexCmux, model) + } return model } diff --git a/pkg/ai/provider/cmux/executor_test.go b/pkg/ai/provider/cmux/executor_test.go index 2507ef6..af3903b 100644 --- a/pkg/ai/provider/cmux/executor_test.go +++ b/pkg/ai/provider/cmux/executor_test.go @@ -35,6 +35,7 @@ func TestAgentCommand(t *testing.T) { {"claude ignores codex extras", AgentCommandOpts{Agent: "claude", Extra: &api.CodexCmuxOptions{Search: true}}, "claude"}, {"codex bare", AgentCommandOpts{Agent: "codex"}, "codex"}, {"codex with model", AgentCommandOpts{Agent: "codex", Model: "gpt-5"}, "codex -m gpt-5"}, + {"codex effort config", AgentCommandOpts{Agent: "codex", Model: "gpt-5.6-sol", Effort: api.EffortUltra}, `codex -m gpt-5.6-sol -c 'model_reasoning_effort="ultra"'`}, {"codex ignores tools and permission", AgentCommandOpts{Agent: "codex", AllowedTools: []string{"Read"}, PermissionMode: api.PermissionPlan}, "codex"}, {"codex extra posture and flags", AgentCommandOpts{Agent: "codex", Model: "gpt-5", Extra: &api.CodexCmuxOptions{Sandbox: api.CodexSandboxWorkspaceWrite, AskForApproval: api.CodexApprovalOnRequest, Search: true}}, "codex -m gpt-5 --sandbox workspace-write --ask-for-approval on-request --search"}, {"codex ignores claude extras", AgentCommandOpts{Agent: "codex", Extra: &api.ClaudeCmuxOptions{AddDir: []string{"x"}}}, "codex"}, @@ -123,7 +124,9 @@ func TestModelFlag(t *testing.T) { }{ {"claude", "claude", ""}, {"claude", "", ""}, - {"claude", "opus", "opus"}, + {"claude", "opus", "claude-opus-5"}, + {"claude", "claude-agent-opus", "claude-opus-5"}, + {"claude", "claude-code-sonnet", "claude-sonnet-5"}, {"codex", "codex", ""}, {"codex", "gpt-5", "gpt-5"}, } diff --git a/pkg/ai/provider/cmux/provider.go b/pkg/ai/provider/cmux/provider.go index 88aecd6..0b14ec9 100644 --- a/pkg/ai/provider/cmux/provider.go +++ b/pkg/ai/provider/cmux/provider.go @@ -48,8 +48,9 @@ func New(cfg ai.Config) (*Provider, error) { if err != nil { return nil, err } + model := ai.NormalizeModelForBackend(cfg.Model.Backend, cfg.Model.Name) return &Provider{ - model: cfg.Model.Name, + model: model, agent: agent, cfg: cfg, }, nil @@ -77,10 +78,9 @@ func (p *Provider) GetBackend() ai.Backend { // Execute drains its own ExecuteStream into a buffered ai.Response. cmux cannot // constrain output natively, so a structured-output request is served by -// appending a schema instruction to the prompt (see ExecuteStream) and -// extracting the JSON object from the reply here. +// appending a schema instruction to the prompt (see ExecuteStream); the stream +// carries the validated JSON on its terminal EventResult. func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { - schemaRequested := req.Prompt.HasSchema() start := time.Now() events, err := p.ExecuteStream(ctx, req) if err != nil { @@ -88,14 +88,25 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e } var ( - text strings.Builder - usage ai.Usage - sessionID string - success = true - sawResult bool - lastErr string + text strings.Builder + usage ai.Usage + sessionID string + success = true + sawResult bool + lastErr string + structured json.RawMessage + outcome *ai.TerminalOutcome + outcomeErr error ) for ev := range events { + if outcomeErr == nil { + parsed, parseErr := ai.TerminalOutcomeFromEvent(ev) + if parseErr != nil { + outcomeErr = parseErr + } else if parsed != nil { + outcome = parsed + } + } switch ev.Kind { case ai.EventText: text.WriteString(ev.Text) @@ -106,6 +117,9 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e case ai.EventResult: sawResult = true success = ev.Success + if len(ev.StructuredData) > 0 { + structured = ev.StructuredData + } if ev.Usage != nil { usage = *ev.Usage } @@ -116,6 +130,9 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e lastErr = ev.Error } } + if outcomeErr != nil { + return nil, fmt.Errorf("cmux: invalid terminal outcome: %w", outcomeErr) + } if !sawResult && lastErr != "" { return nil, fmt.Errorf("%w: %s", ai.ErrCLIExecutionFailed, lastErr) @@ -129,19 +146,24 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e } resp := &ai.Response{ - Text: text.String(), - Model: p.model, - Backend: p.GetBackend(), - Usage: usage, - Duration: time.Since(start), + Text: text.String(), + Model: p.model, + Backend: p.GetBackend(), + Usage: usage, + Duration: time.Since(start), + TerminalOutcome: outcome, } if sessionID != "" { resp.Raw = map[string]any{"session_id": sessionID} } - if schemaRequested { - if obj, ok := ai.ExtractJSONObject(resp.Text); ok { - resp.StructuredData = json.RawMessage(obj) + if len(structured) > 0 && req.Prompt.Schema != nil { + if err := ai.BindStructuredOutput(req.Prompt.Schema, structured); err != nil { + return nil, err } + resp.StructuredData = req.Prompt.Schema + resp.Text = "" + } else if len(structured) > 0 { + resp.StructuredData = structured } return resp, nil } @@ -149,7 +171,7 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e // ExecuteStream drives one cmux run in a goroutine and streams ai.Events on a // buffered channel, closing it when done (always after exactly one EventResult). func (p *Provider) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai.Event, error) { - req, err := withSchemaPrompt(req) + req, schema, err := ai.WithSchemaPrompt(req) if err != nil { return nil, err } @@ -157,45 +179,56 @@ func (p *Provider) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai return nil, fmt.Errorf("cmux provider: prompt is required") } events := make(chan ai.Event, 32) - go p.drive(ctx, req, events) + go p.drive(ctx, req, schema, events) return events, nil } -// withSchemaPrompt makes a structured-output request runnable on cmux (which -// cannot enforce a schema natively) by appending a JSON-only schema instruction -// to the user prompt and clearing the native schema fields so the run is a plain -// text turn. The reply's JSON is recovered in Execute. -func withSchemaPrompt(req ai.Request) (ai.Request, error) { - schema, err := ai.SchemaJSONFor(req.Prompt) - if err != nil { - return req, err - } - if len(schema) > 0 { - req.Prompt.User = strings.TrimRight(req.Prompt.User, "\n") + "\n\n" + ai.SchemaInstruction(string(schema)) - } - req.Prompt.Schema = nil - req.Prompt.SchemaJSON = nil - return req, nil -} - // drive runs the session and translates its outcome into the single terminal // EventResult (preceded by an EventError on failure) before closing the channel. -func (p *Provider) drive(ctx context.Context, req ai.Request, events chan ai.Event) { +func (p *Provider) drive(ctx context.Context, req ai.Request, schema json.RawMessage, events chan ai.Event) { defer close(events) + var text strings.Builder + var outcome *ai.TerminalOutcome + var outcomeErr error r := &run{ client: NewClient(""), - emit: func(ev ai.Event) { emit(ctx, events, ev) }, canUseTool: p.cfg.CanUseTool, } + r.emit = func(ev ai.Event) { + switch ev.Kind { + case ai.EventText: + text.WriteString(ev.Text) + case ai.EventToolUse: + if ev.Tool == "ExitPlanMode" { + r.planExitSeen = true + } + if outcomeErr == nil { + parsed, err := ai.TerminalOutcomeFromEvent(ev) + if err != nil { + outcomeErr = err + } else if parsed != nil { + outcome = parsed + } + } + } + emit(ctx, events, ev) + } usage, cost, err := p.execute(ctx, req, r) + if err == nil && outcomeErr != nil { + err = fmt.Errorf("cmux: invalid terminal outcome: %w", outcomeErr) + } + var structured json.RawMessage + if err == nil { + structured, err = ai.ValidatedStructuredData(schema, text.String(), outcome) + } if err != nil { emit(ctx, events, ai.Event{Kind: ai.EventError, Error: err.Error(), Model: p.model}) emit(ctx, events, ai.Event{Kind: ai.EventResult, Success: false, Error: err.Error(), Model: p.model}) return } - emit(ctx, events, ai.Event{Kind: ai.EventResult, Success: true, Usage: usage, CostUSD: cost, Model: p.model}) + emit(ctx, events, ai.Event{Kind: ai.EventResult, Success: true, Usage: usage, CostUSD: cost, Model: p.model, StructuredData: structured}) } // cmuxExtraArgs decodes req.CLIArgs into the backend's typed "extra cmux args" @@ -246,6 +279,7 @@ func decodeCLIArgs(args map[string]any, out any) error { func (p *Provider) execute(ctx context.Context, req ai.Request, r *run) (*ai.Usage, float64, error) { start := time.Now() agent := p.agent + r.planMode = req.Permissions.Mode == api.PermissionPlan model := modelFlag(agent, p.model) workDir := groupWorkDir(req.Cwd()) @@ -374,6 +408,23 @@ func (p *Provider) execute(ctx context.Context, req ai.Request, r *run) (*ai.Usa _, completed, serr := r.awaitWithStallWatchdog(ctx, ref, sessionID, workDir, timeout, resume, acc) acc.Finish() snap := acc.snapshot() + pausedForQuestion := snap.State == sessionStateAsk + if err := r.dismissCompletedPlan(ctx, ref, sessionID); err != nil { + return nil, 0, err + } + if pausedForQuestion { + // Claude's terminal UI remains inside the interactive question picker even + // though the structured tool event contains everything the host needs to + // render the form. Dismiss that surface once; the host resumes the same + // session later with SendFeedback and a structured answers payload. + if r.planMode { + if err := r.dismissPlanSurface(ctx, ref, sessionID); err != nil { + return nil, 0, err + } + } else if err := r.client.SendKeySurface(ctx, ref.String(), ref.SurfaceID, "Escape"); err != nil { + return nil, 0, fmt.Errorf("dismiss claude question for session %s: %w", sessionID, err) + } + } switch { case errors.Is(serr, errSessionLogNotFound): @@ -385,7 +436,11 @@ func (p *Provider) execute(ctx context.Context, req ai.Request, r *run) (*ai.Usa case !completed: return nil, 0, fmt.Errorf("claude session %s did not complete within %s", sessionID, timeout) default: - log.Infof("cmux: claude session %s completed", sessionID) + if pausedForQuestion { + log.Infof("cmux: claude session %s paused for user input", sessionID) + } else { + log.Infof("cmux: claude session %s completed", sessionID) + } r.lastSurface = ref r.lastSessionID = sessionID r.lastWorkDir = workDir diff --git a/pkg/ai/provider/cmux/provider_test.go b/pkg/ai/provider/cmux/provider_test.go index 4f777f4..d94e65e 100644 --- a/pkg/ai/provider/cmux/provider_test.go +++ b/pkg/ai/provider/cmux/provider_test.go @@ -3,7 +3,6 @@ package cmux import ( "context" "errors" - "strings" "testing" "time" @@ -12,48 +11,17 @@ import ( "github.com/flanksource/commons-db/shell" ) -func TestWithSchemaPrompt(t *testing.T) { - // A raw JSON schema is appended to the prompt and the native fields cleared, - // so the cmux run is a plain text turn that still asks for JSON. - req := ai.Request{Prompt: api.Prompt{ - User: "review the diff", - SchemaJSON: []byte(`{"type":"object","required":["pass"]}`), - }} - got, err := withSchemaPrompt(req) - if err != nil { - t.Fatalf("withSchemaPrompt: %v", err) - } - if got.Prompt.Schema != nil || got.Prompt.SchemaJSON != nil { - t.Errorf("native schema fields must be cleared, got Schema=%v SchemaJSON=%s", got.Prompt.Schema, got.Prompt.SchemaJSON) - } - if !strings.Contains(got.Prompt.User, "review the diff") { - t.Errorf("original prompt lost: %q", got.Prompt.User) - } - if !strings.Contains(got.Prompt.User, `"required":["pass"]`) { - t.Errorf("schema not appended to prompt: %q", got.Prompt.User) - } - - // A text-mode request is returned unchanged. - plain := ai.Request{Prompt: api.Prompt{User: "hi"}} - got, err = withSchemaPrompt(plain) - if err != nil { - t.Fatalf("withSchemaPrompt(text): %v", err) - } - if got.Prompt.User != "hi" { - t.Errorf("text prompt altered: %q", got.Prompt.User) - } -} - func TestNewDerivesAgentAndBackend(t *testing.T) { cases := []struct { name string backend api.Backend model string + wantModel string wantAgent string wantBackend api.Backend }{ - {"claude cmux", api.BackendClaudeCmux, "opus", "claude", api.BackendClaudeCmux}, - {"codex cmux", api.BackendCodexCmux, "gpt-5", "codex", api.BackendCodexCmux}, + {"claude cmux", api.BackendClaudeCmux, "opus", "claude-opus-5", "claude", api.BackendClaudeCmux}, + {"codex cmux", api.BackendCodexCmux, "gpt-5", "gpt-5", "codex", api.BackendCodexCmux}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { @@ -64,8 +32,8 @@ func TestNewDerivesAgentAndBackend(t *testing.T) { if p.agent != tc.wantAgent { t.Fatalf("agent = %q, want %q", p.agent, tc.wantAgent) } - if p.GetModel() != tc.model { - t.Fatalf("GetModel() = %q, want %q", p.GetModel(), tc.model) + if p.GetModel() != tc.wantModel { + t.Fatalf("GetModel() = %q, want %q", p.GetModel(), tc.wantModel) } if p.GetBackend() != tc.wantBackend { t.Fatalf("GetBackend() = %q, want %q", p.GetBackend(), tc.wantBackend) diff --git a/pkg/ai/provider/cmux/sessionlog.go b/pkg/ai/provider/cmux/sessionlog.go index 77535ce..d42567a 100644 --- a/pkg/ai/provider/cmux/sessionlog.go +++ b/pkg/ai/provider/cmux/sessionlog.go @@ -15,8 +15,14 @@ import ( ) const ( - defaultSessionLogPollInterval = 500 * time.Millisecond - defaultSessionLogAppearTimeout = 30 * time.Second + defaultSessionLogPollInterval = 500 * time.Millisecond + // defaultSessionLogAppearTimeout bounds how long we wait for claude to create + // its session log after launch. A cold first invocation (auth check, MCP + // server startup, model-registry fetch) routinely takes longer than 30s to + // write the first log line, so a short window mis-reports a slow-but-healthy + // start as a failed launch. 2 minutes covers cold starts while still failing + // loudly if the session never begins. + defaultSessionLogAppearTimeout = 2 * time.Minute ) // errSessionLogNotFound is returned by the tailer when Claude never created the @@ -53,8 +59,8 @@ func fileSize(path string) int64 { } // sessionTailer follows a Claude session log, streaming each parsed event to a -// callback until the assistant turn ends, the log goes quiet, or the context is -// cancelled. +// callback until the assistant turn ends, pauses for a structured user question, +// the log goes quiet, or the context is cancelled. type sessionTailer struct { path string pollInterval time.Duration @@ -69,6 +75,9 @@ type sessionTailer struct { // into events. It feeds the session-stats accumulator so token/cost totals // stay live during a run. The slice is only valid for the call's duration. onLine func([]byte) + // terminal overrides the normal end-turn predicate for modes whose native + // terminal signal is a tool outcome rather than every assistant end_turn. + terminal func(history.SessionEvent) bool } func (st sessionTailer) poll() time.Duration { @@ -109,7 +118,7 @@ func (st sessionTailer) tail(ctx context.Context, onEvent func(history.SessionEv } if progressed { lastActivity = time.Now() - } else if sawActivity && st.quiescePeriod > 0 && time.Since(lastActivity) >= st.quiescePeriod { + } else if sawActivity && st.terminal == nil && st.quiescePeriod > 0 && time.Since(lastActivity) >= st.quiescePeriod { return true, nil } sawActivity = sawActivity || progressed @@ -149,7 +158,7 @@ func (st sessionTailer) drain(f *os.File, pending *[]byte, buf []byte, onEvent f // EventTurnEnd is a normal completion; EventError is a terminal // API/network failure. Both end the tail — the caller distinguishes // success from failure from the events it saw. - if ev.Kind == history.EventTurnEnd || ev.Kind == history.EventError { + if st.isTerminal(ev) { return progressed, true, nil } } @@ -164,6 +173,24 @@ func (st sessionTailer) drain(f *os.File, pending *[]byte, buf []byte, onEvent f } } +func (st sessionTailer) isTerminal(ev history.SessionEvent) bool { + if st.terminal != nil { + return st.terminal(ev) + } + return ev.Kind == history.EventTurnEnd || ev.Kind == history.EventError || isUserQuestionEvent(ev) +} + +func planModeTerminalEvent(ev history.SessionEvent) bool { + if ev.Kind == history.EventError || isUserQuestionEvent(ev) { + return true + } + return ev.Kind == history.EventToolUse && ev.ToolUse.Tool == "ExitPlanMode" +} + +func isUserQuestionEvent(ev history.SessionEvent) bool { + return ev.Kind == history.EventToolUse && ev.ToolUse.Tool == "AskUserQuestion" +} + // waitForFile polls until the session log exists or the appear timeout elapses. func (st sessionTailer) waitForFile(ctx context.Context) (*os.File, error) { appear := st.appearTimeout @@ -207,6 +234,9 @@ func (r *run) awaitSessionCompletion(ctx context.Context, sessionID, workDir str quiescePeriod: r.sessionLogQuiescePeriod(), seekToEnd: resume, } + if r.planMode { + tailer.terminal = planModeTerminalEvent + } if acc != nil { tailer.onLine = acc.AddLine } diff --git a/pkg/ai/provider/cmux/sessionlog_test.go b/pkg/ai/provider/cmux/sessionlog_test.go index 3d6b350..5e9974d 100644 --- a/pkg/ai/provider/cmux/sessionlog_test.go +++ b/pkg/ai/provider/cmux/sessionlog_test.go @@ -95,6 +95,83 @@ func TestSessionTailerTerminatesOnAPIError(t *testing.T) { } } +func TestSessionTailerPausesOnAskUserQuestion(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + writeSessionLog(t, path, + `{"type":"assistant","sessionId":"s","message":{"stop_reason":"tool_use","content":[{"type":"tool_use","id":"ask-1","name":"AskUserQuestion","input":{"questions":[{"question":"Which database?"}]}}]}}`, + ) + + tailer := sessionTailer{path: path, pollInterval: time.Millisecond, appearTimeout: 200 * time.Millisecond} + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + var events []history.SessionEvent + completed, err := tailer.tail(ctx, func(ev history.SessionEvent) { events = append(events, ev) }) + if err != nil || !completed { + t.Fatalf("tail() = (%v, %v), want paused terminal outcome", completed, err) + } + if len(events) != 1 || !isUserQuestionEvent(events[0]) || events[0].ToolUse.ToolUseID != "ask-1" { + t.Fatalf("events = %+v, want AskUserQuestion ask-1", events) + } +} + +func TestSessionTailerPlanModeIgnoresEndTurnUntilExitPlanMode(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + writeSessionLog(t, path, + `{"type":"assistant","sessionId":"s","message":{"stop_reason":"end_turn","content":[{"type":"text","text":"intermediate"}]}}`, + `{"type":"assistant","sessionId":"s","message":{"stop_reason":"tool_use","content":[{"type":"tool_use","id":"plan-1","name":"ExitPlanMode","input":{"plan":"1. Inspect\n2. Implement","planFilePath":"/repo/.claude/plans/example.md"}}]}}`, + ) + + tailer := sessionTailer{ + path: path, + pollInterval: time.Millisecond, + appearTimeout: 200 * time.Millisecond, + terminal: planModeTerminalEvent, + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + var events []history.SessionEvent + completed, err := tailer.tail(ctx, func(ev history.SessionEvent) { events = append(events, ev) }) + + if err != nil || !completed { + t.Fatalf("tail() = (%v, %v), want terminal plan outcome", completed, err) + } + if len(events) != 3 { + t.Fatalf("events = %+v, want intermediate text + end turn + ExitPlanMode", events) + } + last := events[len(events)-1] + if last.Kind != history.EventToolUse || last.ToolUse.Tool != "ExitPlanMode" { + t.Fatalf("last event = %+v, want ExitPlanMode", last) + } +} + +func TestSessionTailerPlanModeUsesTimeoutInsteadOfQuiescence(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + writeSessionLog(t, path, + `{"type":"assistant","sessionId":"s","message":{"stop_reason":"end_turn","content":[{"type":"text","text":"waiting for background agents"}]}}`, + ) + + tailer := sessionTailer{ + path: path, + pollInterval: time.Millisecond, + appearTimeout: 200 * time.Millisecond, + quiescePeriod: 5 * time.Millisecond, + terminal: planModeTerminalEvent, + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + + completed, err := tailer.tail(ctx, func(history.SessionEvent) {}) + + if completed { + t.Fatal("plan-mode tail completed on quiescence before ExitPlanMode") + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("tail() error = %v, want context deadline", err) + } +} + func TestSessionTailerCompletesOnQuiescence(t *testing.T) { path := filepath.Join(t.TempDir(), "s.jsonl") // No end_turn — only a mid-turn entry; completion must come from quiescence. diff --git a/pkg/ai/provider/cmux/stall.go b/pkg/ai/provider/cmux/stall.go index 177afff..5581c44 100644 --- a/pkg/ai/provider/cmux/stall.go +++ b/pkg/ai/provider/cmux/stall.go @@ -225,7 +225,20 @@ func (w *stallWatchdog) markAwaitingHuman() { // interactive terminal user to answer (the stall clock is held by awaitingHuman). func (w *stallWatchdog) maybeRequestApproval(ctx context.Context, screen string) { req, ok := detectApprovalRequest(w.sessionID, screen) - if !ok || w.r.canUseTool == nil { + if !ok { + return + } + // The shared plan-mode policy: a plan-only run must return ExitPlanMode to + // its caller — accepting this dialog would start implementation inside the + // same agent turn. On this transport the deny is applied by dismissing the + // plan surface. + if _, handled := ai.PlanTerminalPermission(w.r.planMode, req); handled { + if err := w.r.dismissPlanSurface(ctx, w.ref, w.sessionID); err != nil { + log.Warnf("cmux: failed to dismiss plan-only approval: %v", err) + } + return + } + if w.r.canUseTool == nil { return } if !w.approving.CompareAndSwap(false, true) { diff --git a/pkg/ai/provider/cmux/stall_test.go b/pkg/ai/provider/cmux/stall_test.go index 1954404..59ff772 100644 --- a/pkg/ai/provider/cmux/stall_test.go +++ b/pkg/ai/provider/cmux/stall_test.go @@ -312,6 +312,54 @@ func TestStallWatchdogPlanApprovalBrokeredOnce(t *testing.T) { } } +func TestStallWatchdogPlanRunNeverApprovesExitPlanMode(t *testing.T) { + runner := &stallRunner{} + runner.setScreen(planApprovalScreen) + r := newTestRun(runConfig{}, runner.run) + r.planMode = true + var broker atomic.Int32 + r.canUseTool = func(_ context.Context, _ ai.PermissionRequest) (ai.PermissionDecision, error) { + broker.Add(1) + return ai.PermissionDecision{Allow: true}, nil + } + wd := newWatchdog(r, "plan-safe", filepath.Join(t.TempDir(), "s.jsonl"), nil) + + wd.maybeRequestApproval(context.Background(), planApprovalScreen) + wd.maybeRequestApproval(context.Background(), planApprovalScreen) + r.approvals.Wait() + + if got := broker.Load(); got != 0 { + t.Fatalf("plan approval broker called %d times, want 0", got) + } + if got := runner.enterCount(); got != 0 { + t.Fatalf("Enter sent %d times, want 0 for a plan-only run", got) + } + if got := runner.escapeCount(); got != 1 { + t.Fatalf("Escape sent %d times, want exactly 1 for a plan-only run", got) + } +} + +func TestCompletedPlanExitDismissesSurfaceWhenWatchdogMissesDialog(t *testing.T) { + runner := &stallRunner{} + r := newTestRun(runConfig{}, runner.run) + r.planMode = true + r.planExitSeen = true + + if err := r.dismissCompletedPlan(context.Background(), testSurface, "plan-complete"); err != nil { + t.Fatalf("dismissCompletedPlan() error = %v", err) + } + if err := r.dismissCompletedPlan(context.Background(), testSurface, "plan-complete"); err != nil { + t.Fatalf("second dismissCompletedPlan() error = %v", err) + } + + if got := runner.enterCount(); got != 0 { + t.Fatalf("Enter sent %d times, want 0 for a completed plan-only run", got) + } + if got := runner.escapeCount(); got != 1 { + t.Fatalf("Escape sent %d times, want exactly 1 for a completed plan-only run", got) + } +} + func TestStallWatchdogLogGrowthResetsTimer(t *testing.T) { logPath := filepath.Join(t.TempDir(), "s.jsonl") writeSessionLog(t, logPath, "seed") diff --git a/pkg/ai/provider/coalesce.go b/pkg/ai/provider/coalesce.go index e29bfee..aebf6e0 100644 --- a/pkg/ai/provider/coalesce.go +++ b/pkg/ai/provider/coalesce.go @@ -14,19 +14,40 @@ import ( // the live event stream call this on a tee'd channel; for one-shot use, prefer // Provider.Execute. func CoalesceStream(ctx context.Context, model string, events <-chan ai.Event, start time.Time) (*ai.Response, error) { + return CoalesceStreamForBackend(ctx, ai.BackendClaudeCLI, model, events, start) +} + +func CoalesceStreamForBackend(ctx context.Context, backend ai.Backend, model string, events <-chan ai.Event, start time.Time) (*ai.Response, error) { var ( text strings.Builder usage ai.Usage lastResult *ai.Event errEvents []ai.Event sessionID string + outcome *ai.TerminalOutcome + outcomeErr error ) for { select { case ev, ok := <-events: if !ok { - return finaliseCoalescedResponse(model, text.String(), usage, lastResult, errEvents, sessionID, start) + if outcomeErr != nil { + return nil, fmt.Errorf("invalid terminal outcome: %w", outcomeErr) + } + resp, err := finaliseCoalescedResponse(backend, model, text.String(), usage, lastResult, errEvents, sessionID, start) + if resp != nil { + resp.TerminalOutcome = outcome + } + return resp, err + } + if outcomeErr == nil { + parsed, err := ai.TerminalOutcomeFromEvent(ev) + if err != nil { + outcomeErr = err + } else if parsed != nil { + outcome = parsed + } } switch ev.Kind { case ai.EventText: @@ -34,6 +55,9 @@ func CoalesceStream(ctx context.Context, model string, events <-chan ai.Event, s case ai.EventResult: cp := ev lastResult = &cp + if ev.Text != "" && text.Len() == 0 { + text.WriteString(ev.Text) + } if ev.Usage != nil { usage = *ev.Usage } @@ -50,7 +74,7 @@ func CoalesceStream(ctx context.Context, model string, events <-chan ai.Event, s } } -func finaliseCoalescedResponse(model, text string, usage ai.Usage, lastResult *ai.Event, errEvents []ai.Event, sessionID string, start time.Time) (*ai.Response, error) { +func finaliseCoalescedResponse(backend ai.Backend, model, text string, usage ai.Usage, lastResult *ai.Event, errEvents []ai.Event, sessionID string, start time.Time) (*ai.Response, error) { if lastResult != nil && !lastResult.Success { msg := lastResult.Error if msg == "" { @@ -65,12 +89,16 @@ func finaliseCoalescedResponse(model, text string, usage ai.Usage, lastResult *a resp := &ai.Response{ Text: text, Model: model, - Backend: ai.BackendClaudeCLI, + Backend: backend, Usage: usage, Duration: time.Since(start), } if lastResult != nil { resp.Raw = lastResult.Input + // Carry the provider-reported cost (e.g. claude-cli total_cost_usd) so the + // buffered Execute path does not lose it — buffered callers otherwise fall + // back to a list-price recompute (finding D4). + resp.CostUSD = lastResult.CostUSD // Carry any validated structured output (raw JSON) the terminal result // event holds; the caller's Execute binds it into its schema target. if len(lastResult.StructuredData) > 0 { diff --git a/pkg/ai/provider/coalesce_test.go b/pkg/ai/provider/coalesce_test.go index fc8a755..8bded16 100644 --- a/pkg/ai/provider/coalesce_test.go +++ b/pkg/ai/provider/coalesce_test.go @@ -66,6 +66,57 @@ func TestCoalesceStream_CarriesStructuredData(t *testing.T) { } } +func TestCoalesceStream_CarriesTerminalOutcome(t *testing.T) { + events := []ai.Event{ + {Kind: ai.EventToolUse, Tool: "ExitPlanMode", Input: map[string]any{ + "plan": "1. Inspect\n2. Implement", + "planFilePath": "/repo/.claude/plans/example.md", + }}, + {Kind: ai.EventResult, Success: true}, + } + + resp, err := CoalesceStream(context.Background(), "claude", feedEvents(events), time.Now()) + if err != nil { + t.Fatalf("CoalesceStream err: %v", err) + } + if resp.TerminalOutcome == nil || resp.TerminalOutcome.Plan == nil { + t.Fatalf("TerminalOutcome = %+v, want plan", resp.TerminalOutcome) + } + if resp.TerminalOutcome.Plan.Content != "1. Inspect\n2. Implement" { + t.Errorf("plan content = %q", resp.TerminalOutcome.Plan.Content) + } +} + +func TestCoalesceStream_FailsOnMalformedTerminalOutcome(t *testing.T) { + events := []ai.Event{ + {Kind: ai.EventToolUse, Tool: "ExitPlanMode", Input: map[string]any{"planFilePath": "/repo/plan.md"}}, + {Kind: ai.EventResult, Success: true}, + } + + resp, err := CoalesceStream(context.Background(), "claude", feedEvents(events), time.Now()) + if err == nil { + t.Fatalf("CoalesceStream error = nil, response = %+v", resp) + } + if !strings.Contains(err.Error(), "plan is required") { + t.Fatalf("CoalesceStream error = %v, want missing plan", err) + } +} + +func TestCoalesceStream_UsesResultTextWhenNoTextEvents(t *testing.T) { + const model = "claude-sonnet-5" + events := []ai.Event{ + {Kind: ai.EventResult, Tool: "Result", Success: true, Text: "final answer", Model: model}, + } + + resp, err := CoalesceStream(context.Background(), model, feedEvents(events), time.Now()) + if err != nil { + t.Fatalf("CoalesceStream err: %v", err) + } + if resp.Text != "final answer" { + t.Errorf("Text = %q, want final answer", resp.Text) + } +} + func TestCoalesceStream_ResultErrorReturnsError(t *testing.T) { const wantMsg = "upstream 500" events := []ai.Event{ diff --git a/pkg/ai/provider/codex_appserver.go b/pkg/ai/provider/codex_appserver.go index 891063f..f474c82 100644 --- a/pkg/ai/provider/codex_appserver.go +++ b/pkg/ai/provider/codex_appserver.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" osexec "os/exec" + "strings" "sync" "time" @@ -29,11 +30,12 @@ type CodexAppServer struct { turnMu sync.Mutex // serializes turns; held by ExecuteStream, freed by its driver - mu sync.Mutex // guards the rpc/supervisor handles + the active turn - sup *exec.SupervisedProcess - rpc *jsonrpc.Client - rpcDone chan struct{} // closed by the rpc Run goroutine when the child exits - active *turnState + mu sync.Mutex // guards the rpc/supervisor handles + the active turn + sup *exec.SupervisedProcess + rpc *jsonrpc.Client + rpcDone chan struct{} // closed by the rpc Run goroutine when the child exits + active *turnState + threadID string } const ( @@ -47,11 +49,12 @@ func NewCodexAppServer(model string) (*CodexAppServer, error) { if model == "" { model = CodexCLIDefaultModel } + model = ai.NormalizeModelForBackend(ai.BackendCodexAgent, model) return &CodexAppServer{model: model}, nil } func (c *CodexAppServer) GetModel() string { return c.model } -func (c *CodexAppServer) GetBackend() ai.Backend { return ai.BackendCodexCLI } +func (c *CodexAppServer) GetBackend() ai.Backend { return ai.BackendCodexAgent } // Execute drains the streaming output into a buffered ai.Response. When the // request carries a structured-output schema, the final agent message's JSON is @@ -63,7 +66,7 @@ func (c *CodexAppServer) Execute(ctx context.Context, req ai.Request) (*ai.Respo if err != nil { return nil, err } - resp, err := CoalesceStream(ctx, c.model, events, start) + resp, err := CoalesceStreamForBackend(ctx, ai.BackendCodexAgent, c.model, events, start) if err != nil { return nil, err } @@ -89,7 +92,7 @@ func (c *CodexAppServer) Execute(ctx context.Context, req ai.Request) (*ai.Respo // a text-mode request. A non-struct target fails loudly rather than silently // dropping the schema. func codexOutputSchema(req ai.Request) (json.RawMessage, error) { - schema, err := ai.SchemaJSONFor(req.Prompt) + schema, err := ai.SchemaJSONForBackend(ai.BackendCodexAgent, req.Prompt) if err != nil { return nil, fmt.Errorf("codex app-server: cannot derive structured-output schema: %w", err) } @@ -118,7 +121,9 @@ func (c *CodexAppServer) ExecuteStream(ctx context.Context, req ai.Request) (<-c usage: &ai.Usage{}, model: c.model, streamed: map[string]bool{}, + toolOutput: map[string]string{}, terminal: make(chan struct{}), + started: make(chan struct{}), outputSchema: schema, } c.setActive(ts) @@ -142,6 +147,7 @@ func (c *CodexAppServer) driveTurn(ctx context.Context, req ai.Request, ts *turn c.failTurn(ts, err) return } + ts.setIDs(threadID, turnID) select { case <-ts.terminal: @@ -186,10 +192,12 @@ func (c *CodexAppServer) ensureStarted(ctx context.Context) error { } ready := make(chan error, 1) + var process *exec.Process sup := exec.NewExec("codex", "app-server").WithStdioPipe().Supervise(exec.SuperviseOptions{ // No restart: a crash surfaces as EventError, never a silent retry. RestartPolicy: exec.RestartNo, OnStarted: func(p *exec.Process) { + process = p rpc := jsonrpc.New(p.Stdin(), p.StdoutReader(), true, jsonrpc.Handlers{ OnNotification: c.handleNotification, OnRequest: c.handleApproval, @@ -212,6 +220,9 @@ func (c *CodexAppServer) ensureStarted(ctx context.Context) error { select { case err := <-ready: if err != nil { + if process != nil { + err = appServerProcessError(err, process.GetStderr()) + } c.teardown(true) return err } @@ -225,6 +236,13 @@ func (c *CodexAppServer) ensureStarted(ctx context.Context) error { } } +func appServerProcessError(err error, stderr string) error { + if detail := strings.TrimSpace(stderr); detail != "" { + return fmt.Errorf("%w: %s", err, detail) + } + return err +} + // handshake performs the required initialize → initialized exchange (any request // before it errors "Not initialized" server-side). func (c *CodexAppServer) handshake(ctx context.Context, rpc *jsonrpc.Client) error { @@ -243,7 +261,7 @@ func (c *CodexAppServer) handshake(ctx context.Context, rpc *jsonrpc.Client) err func (c *CodexAppServer) teardown(stop bool) { c.mu.Lock() sup := c.sup - c.sup, c.rpc, c.rpcDone = nil, nil, nil + c.sup, c.rpc, c.rpcDone, c.threadID = nil, nil, nil, "" c.mu.Unlock() if stop && sup != nil { sup.Stop() @@ -255,18 +273,37 @@ func (c *CodexAppServer) startThread(ctx context.Context, req ai.Request) (strin if rpc == nil { return "", fmt.Errorf("codex app-server: not started") } + c.mu.Lock() + threadID := c.threadID + c.mu.Unlock() + if threadID != "" { + return threadID, nil + } if req.SessionID != "" { raw, err := rpc.Call(ctx, "thread/resume", buildResumeParams(req)) if err != nil { return "", err } - return firstNonEmpty(parseAppServerNotif(raw).threadID(), req.SessionID), nil + threadID = firstNonEmpty(parseAppServerNotif(raw).threadID(), req.SessionID) + c.rememberThread(threadID) + return threadID, nil } raw, err := rpc.Call(ctx, "thread/start", buildThreadStartParams(c.model, req)) if err != nil { return "", err } - return parseAppServerNotif(raw).threadID(), nil + threadID = parseAppServerNotif(raw).threadID() + if threadID == "" { + return "", fmt.Errorf("codex app-server: thread/start returned no thread id") + } + c.rememberThread(threadID) + return threadID, nil +} + +func (c *CodexAppServer) rememberThread(threadID string) { + c.mu.Lock() + c.threadID = threadID + c.mu.Unlock() } func (c *CodexAppServer) startTurn(ctx context.Context, req ai.Request, threadID string, outputSchema json.RawMessage) (string, error) { @@ -274,7 +311,11 @@ func (c *CodexAppServer) startTurn(ctx context.Context, req ai.Request, threadID if rpc == nil { return "", fmt.Errorf("codex app-server: not started") } - raw, err := rpc.Call(ctx, "turn/start", buildTurnStartParams(c.model, req, threadID, outputSchema)) + params, err := buildTurnStartParams(c.model, req, threadID, outputSchema) + if err != nil { + return "", err + } + raw, err := rpc.Call(ctx, "turn/start", params) if err != nil { return "", err } @@ -296,6 +337,30 @@ func (c *CodexAppServer) interrupt(threadID, turnID string) { _, _ = rpc.Call(ctx, "turn/interrupt", map[string]string{"threadId": threadID, "turnId": turnID}) } +func (c *CodexAppServer) Interrupt(ctx context.Context) error { + ts := c.currentTurn() + if ts == nil { + return fmt.Errorf("codex app-server: no active turn to interrupt") + } + threadID, turnID, err := ts.waitIDs(ctx) + if err != nil { + return err + } + rpc := c.client() + if rpc == nil { + return fmt.Errorf("codex app-server: not started") + } + if _, err := rpc.Call(ctx, "turn/interrupt", map[string]string{"threadId": threadID, "turnId": turnID}); err != nil { + return fmt.Errorf("codex app-server interrupt failed: %w", err) + } + return nil +} + +func (c *CodexAppServer) Close() error { + c.teardown(true) + return nil +} + // handleNotification routes one notification to the active turn. It runs on the // rpc Run goroutine (notifications dispatch sequentially), so the per-turn // dedup/usage state needs no extra locking. @@ -304,23 +369,40 @@ func (c *CodexAppServer) handleNotification(method string, params json.RawMessag if ts == nil { return } + ctx := appServerEventContext{Model: ts.model, Usage: ts.usage} switch method { case "item/agentMessage/delta": if id := parseAppServerNotif(params).ItemID; id != "" { ts.streamed[id] = true } + if len(ts.outputSchema) > 0 { + return + } + case "item/commandExecution/outputDelta": + n := parseAppServerNotif(params) + if n.ItemID != "" && n.Delta != "" { + ts.toolOutput[n.ItemID] += n.Delta + } + return case "item/completed": // The completed agent message carries the full text (structured runs use it // as the validated JSON result). Capture it, then drop the duplicate so // CoalesceStream / renderers don't double-count already-streamed deltas. if it := parseAppServerNotif(params).Item; it != nil && it.Type == "agentMessage" { ts.lastAgentMessage = it.Text + if len(ts.outputSchema) > 0 { + return + } } if appServerStreamedAgentMessage(params, ts.streamed) { return } + if it := parseAppServerNotif(params).Item; it != nil { + ctx.ToolOutput = ts.toolOutput[it.ID] + delete(ts.toolOutput, it.ID) + } } - if ev, ok := mapAppServerNotification(method, params, ts.model, ts.usage); ok { + if ev, ok := mapAppServerNotification(method, params, ctx); ok { if ev.Kind == ai.EventResult && len(ts.outputSchema) > 0 { ev.StructuredData = json.RawMessage(ts.lastAgentMessage) } @@ -353,50 +435,6 @@ func (c *CodexAppServer) handleApproval(method string, _ json.RawMessage) (any, // turnState is the routing target for one turn's server notifications. send and // finish are guarded so a late notification can never send on a closed channel; // terminal doubles as the "stop sending" signal that unblocks a blocked send. -type turnState struct { - ch chan ai.Event - usage *ai.Usage - model string - streamed map[string]bool // agent-message item IDs already streamed via deltas - - // outputSchema is non-nil when the turn requested structured output; it drives - // the turn/start outputSchema and gates capturing lastAgentMessage as the - // structured result. lastAgentMessage is the full text of the most recent - // completed agentMessage (codex returns structured JSON as that text). - outputSchema json.RawMessage - lastAgentMessage string - - terminal chan struct{} // closed when the turn is over (by codex or by finish) - termOnce sync.Once - sendMu sync.Mutex - closed bool -} - -func (ts *turnState) signalTerminal() { ts.termOnce.Do(func() { close(ts.terminal) }) } - -func (ts *turnState) send(ev ai.Event) { - ts.sendMu.Lock() - defer ts.sendMu.Unlock() - if ts.closed { - return - } - select { - case ts.ch <- ev: - case <-ts.terminal: - } -} - -func (ts *turnState) finish() { - ts.signalTerminal() // unblock a blocked send before locking - ts.sendMu.Lock() - defer ts.sendMu.Unlock() - if ts.closed { - return - } - ts.closed = true - close(ts.ch) -} - // compile-time assertion: the provider satisfies the streaming interface. // The pure mapping, parse structs, and request-param builders live in the // sibling codex_appserver_protocol.go (same package). diff --git a/pkg/ai/provider/codex_appserver_lifecycle_ginkgo_test.go b/pkg/ai/provider/codex_appserver_lifecycle_ginkgo_test.go new file mode 100644 index 0000000..c2b2da2 --- /dev/null +++ b/pkg/ai/provider/codex_appserver_lifecycle_ginkgo_test.go @@ -0,0 +1,137 @@ +package provider + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + . "github.com/onsi/gomega/gstruct" +) + +func TestCodexAppServerLifecycle(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Codex App Server Lifecycle Suite") +} + +var _ = Describe("Codex app-server tool lifecycle", func() { + It("emits one command use and one correlated result", func() { + client, turn := activeGinkgoTurn() + client.handleNotification("item/started", json.RawMessage(`{ + "threadId":"thread-1", + "item":{"id":"cmd-1","type":"commandExecution","command":"pwd","cwd":"/work","status":"inProgress"} + }`)) + client.handleNotification("item/commandExecution/outputDelta", json.RawMessage(`{ + "threadId":"thread-1","itemId":"cmd-1","delta":"/work\n" + }`)) + client.handleNotification("item/completed", json.RawMessage(`{ + "threadId":"thread-1", + "item":{"id":"cmd-1","type":"commandExecution","command":"pwd","cwd":"/work","status":"completed","exitCode":0,"aggregatedOutput":""} + }`)) + + events := drainEvents(turn) + Expect(events).To(HaveLen(2)) + Expect(events[0]).To(MatchFields(IgnoreExtras, Fields{ + "Kind": Equal(ai.EventToolUse), + "Tool": Equal("Bash"), + "Input": Equal(map[string]any{"command": "pwd"}), + "ToolCallID": Equal("cmd-1"), + "SessionID": Equal("thread-1"), + })) + Expect(events[1]).To(MatchFields(IgnoreExtras, Fields{ + "Kind": Equal(ai.EventToolResult), + "Text": Equal("/work\n"), + "ToolCallID": Equal("cmd-1"), + "SessionID": Equal("thread-1"), + "Success": BeTrue(), + })) + + tool, ok := events[0].Raw.(claude.ToolUse) + Expect(ok).To(BeTrue()) + Expect(tool.ToolUseID).To(Equal("cmd-1")) + Expect(tool.SessionID).To(Equal("thread-1")) + }) + + It("marks nonzero command completion as failed", func() { + client, turn := activeGinkgoTurn() + client.handleNotification("item/started", json.RawMessage(`{ + "threadId":"thread-1", + "item":{"id":"cmd-2","type":"commandExecution","command":"false","status":"inProgress"} + }`)) + client.handleNotification("item/completed", json.RawMessage(`{ + "threadId":"thread-1", + "item":{"id":"cmd-2","type":"commandExecution","command":"false","status":"failed","exitCode":1,"aggregatedOutput":"failed\n"} + }`)) + + events := drainEvents(turn) + Expect(events).To(HaveLen(2)) + Expect(events[1].Kind).To(Equal(ai.EventToolResult)) + Expect(events[1].Success).To(BeFalse()) + Expect(events[1].Text).To(Equal("failed\n")) + }) +}) + +var _ = Describe("Codex app-server turn control", func() { + It("waits for thread and turn identifiers before interrupting", func() { + turn := &turnState{terminal: make(chan struct{}), started: make(chan struct{})} + go func() { + time.Sleep(10 * time.Millisecond) + turn.setIDs("thread-1", "turn-1") + }() + + threadID, turnID, err := turn.waitIDs(context.Background()) + + Expect(err).NotTo(HaveOccurred()) + Expect(threadID).To(Equal("thread-1")) + Expect(turnID).To(Equal("turn-1")) + }) +}) + +var _ = Describe("Codex app-server attachments", func() { + It("sends ordered localImage inputs after optional text", func() { + attachment := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"}. + WithPreparedContent(api.AttachmentContent{Path: "/work/diagram.png"}) + params, err := buildTurnStartParams("gpt-5", ai.Request{Prompt: api.Prompt{ + User: "inspect", + Attachments: []api.AttachmentRef{attachment}, + }}, "thread-1", nil) + Expect(err).NotTo(HaveOccurred()) + Expect(params["input"]).To(Equal([]map[string]any{ + {"type": "text", "text": "inspect"}, + {"type": "localImage", "path": "/work/diagram.png"}, + })) + }) +}) + +var _ = Describe("Codex CLI attachments", func() { + It("adds one native --image argument per prepared attachment", func() { + first := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"}. + WithPreparedContent(api.AttachmentContent{Path: "/work/first.png"}) + second := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/jpeg"}. + WithPreparedContent(api.AttachmentContent{Path: "/work/second.jpg"}) + args, cleanup, err := buildCodexCLIArgs("codex", ai.Request{Prompt: api.Prompt{Attachments: []api.AttachmentRef{first, second}}}) + DeferCleanup(cleanup) + Expect(err).NotTo(HaveOccurred()) + Expect(args[:6]).To(Equal([]string{"exec", "--json", "--image", "/work/first.png", "--image", "/work/second.jpg"})) + }) +}) + +func activeGinkgoTurn() (*CodexAppServer, *turnState) { + client, err := NewCodexAppServer("gpt-5") + Expect(err).NotTo(HaveOccurred()) + turn := &turnState{ + ch: make(chan ai.Event, 16), + usage: &ai.Usage{}, + model: "gpt-5", + streamed: map[string]bool{}, + toolOutput: map[string]string{}, + terminal: make(chan struct{}), + } + client.setActive(turn) + return client, turn +} diff --git a/pkg/ai/provider/codex_appserver_protocol.go b/pkg/ai/provider/codex_appserver_protocol.go index 451c549..b706d24 100644 --- a/pkg/ai/provider/codex_appserver_protocol.go +++ b/pkg/ai/provider/codex_appserver_protocol.go @@ -2,8 +2,10 @@ package provider import ( "encoding/json" + "fmt" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/history" "github.com/flanksource/captain/pkg/api" ) @@ -38,11 +40,22 @@ type appServerThread struct { } type appServerItemBody struct { - ID string `json:"id"` - Type string `json:"type"` - Text string `json:"text"` - Command string `json:"command"` - Tool string `json:"tool"` + ID string `json:"id"` + Type string `json:"type"` + Text string `json:"text"` + Command string `json:"command"` + Tool string `json:"tool"` + Server string `json:"server"` + CWD string `json:"cwd"` + Status string `json:"status"` + Arguments json.RawMessage `json:"arguments"` + AggregatedOutput *string `json:"aggregatedOutput"` + ExitCode *int `json:"exitCode"` + Success *bool `json:"success"` + Result json.RawMessage `json:"result"` + ContentItems json.RawMessage `json:"contentItems"` + Changes json.RawMessage `json:"changes"` + Error *appServerErrorBody `json:"error"` } type appServerErrorBody struct { @@ -93,9 +106,12 @@ func (n appServerNotif) foldUsage(usage *ai.Usage) { if usage == nil || n.TokenUsage == nil { return } + // Codex uses OpenAI accounting: inputTokens folds in the cached prefix and + // outputTokens folds in reasoning. Net both so pricing/totals do not + // double-count against cachedInputTokens/reasoningOutputTokens. t := n.TokenUsage.Total - usage.InputTokens = t.InputTokens - usage.OutputTokens = t.OutputTokens + usage.InputTokens = ai.NetInputTokens(t.InputTokens, t.CachedInputTokens) + usage.OutputTokens = ai.NetOutputTokens(t.OutputTokens, t.ReasoningOutputTokens) usage.CacheReadTokens = t.CachedInputTokens usage.ReasoningTokens = t.ReasoningOutputTokens } @@ -105,86 +121,149 @@ func (n appServerNotif) foldUsage(usage *ai.Usage) { // mapAppServerNotification maps one notification into an ai.Event. Pure: it // folds thread/tokenUsage/updated into usage (ok=false) and reads the folded // usage back out on turn/completed. -func mapAppServerNotification(method string, params json.RawMessage, model string, usage *ai.Usage) (ai.Event, bool) { +type appServerEventContext struct { + Model string + Usage *ai.Usage + ToolOutput string +} + +func mapAppServerNotification(method string, params json.RawMessage, ctx appServerEventContext) (ai.Event, bool) { n := parseAppServerNotif(params) switch method { case "thread/started": sid := n.threadID() - out := ai.Event{Kind: ai.EventSystem, Tool: "SessionInit", SessionID: sid, Model: model} - out.Raw = codexSessionToolUse(sid, model) + out := ai.Event{Kind: ai.EventSystem, Tool: "SessionInit", SessionID: sid, Model: ctx.Model} + out.Raw = codexSessionToolUse(sid, ctx.Model) return out, true case "item/agentMessage/delta": if n.Delta == "" { return ai.Event{}, false } - return ai.Event{Kind: ai.EventText, Text: n.Delta, Model: model}, true + return ai.Event{Kind: ai.EventText, Text: n.Delta, SessionID: n.threadID(), Model: ctx.Model}, true case "item/reasoning/textDelta", "item/reasoning/summaryTextDelta": if n.Delta == "" { return ai.Event{}, false } - return ai.Event{Kind: ai.EventThinking, Text: n.Delta, Model: model}, true + return ai.Event{Kind: ai.EventThinking, Text: n.Delta, SessionID: n.threadID(), Model: ctx.Model}, true case "item/commandExecution/outputDelta": - if n.Delta == "" { - return ai.Event{}, false - } - input := map[string]any{"delta": n.Delta} - out := ai.Event{Kind: ai.EventToolUse, Tool: "command", Input: input, Model: model} - out.Raw = codexToolUse("command", input, n.ItemID, model) - return out, true + return ai.Event{}, false case "item/started", "item/completed": - return mapAppServerItem(n.Item, model) + return mapAppServerItem(method, n.Item, n.threadID(), ctx) case "thread/tokenUsage/updated": - n.foldUsage(usage) + n.foldUsage(ctx.Usage) return ai.Event{}, false case "turn/completed": - out := ai.Event{Kind: ai.EventResult, Tool: "Result", Model: model, Success: true} - if usage != nil && usage.TotalTokens() > 0 { - u := *usage + out := ai.Event{Kind: ai.EventResult, Tool: "Result", SessionID: n.threadID(), Model: ctx.Model, Success: true} + if ctx.Usage != nil && ctx.Usage.TotalTokens() > 0 { + u := *ctx.Usage out.Usage = &u } out.Raw = codexResultToolUse(out, n.ThreadID) return out, true case "turn/failed", "error": - return ai.Event{Kind: ai.EventError, Error: extractCodexErrorText(n.errorText()), Model: model}, true + return ai.Event{Kind: ai.EventError, Error: extractCodexErrorText(n.errorText()), SessionID: n.threadID(), Model: ctx.Model}, true } return ai.Event{}, false } // mapAppServerItem dispatches item/started and item/completed on the item type: -// agent messages become text, command/tool/file items become tool-use events, -// and reasoning/user/hook items are dropped (reasoning text arrives via deltas). -func mapAppServerItem(it *appServerItemBody, model string) (ai.Event, bool) { +// agent messages become text, command/tool/file items become correlated use and +// result events, and reasoning/user/hook items are dropped. +func mapAppServerItem(method string, it *appServerItemBody, sessionID string, ctx appServerEventContext) (ai.Event, bool) { if it == nil { return ai.Event{}, false } switch it.Type { case "agentMessage", "plan": - if it.Text == "" { + if method != "item/completed" || it.Text == "" { return ai.Event{}, false } - return ai.Event{Kind: ai.EventText, Text: it.Text, Model: model}, true + return ai.Event{Kind: ai.EventText, Text: it.Text, SessionID: sessionID, Model: ctx.Model}, true case "reasoning", "userMessage", "hookPrompt", "": return ai.Event{}, false - default: - name := firstNonEmpty(it.Tool, it.Command, it.Type) - input := map[string]any{} - if it.Command != "" { - input["command"] = it.Command - } - if it.Text != "" { - input["text"] = it.Text + } + use := history.NormalizeCodexToolCall(appServerToolCall(it, sessionID, ctx.Model)) + if method == "item/started" { + out := ai.Event{ + Kind: ai.EventToolUse, Tool: use.Tool, Input: use.Input, + ToolCallID: use.ToolUseID, SessionID: sessionID, Model: ctx.Model, } - out := ai.Event{Kind: ai.EventToolUse, Tool: name, Input: input, Model: model} - out.Raw = codexToolUse(name, input, it.ID, model) + out.Raw = codexToolUse(use, ctx.Model) return out, true } + if method != "item/completed" { + return ai.Event{}, false + } + text := appServerToolResultText(it, ctx.ToolOutput) + success := appServerToolSucceeded(it) + use.Response = text + raw := codexToolUse(use, ctx.Model) + raw.IsError = !success + return ai.Event{ + Kind: ai.EventToolResult, Text: text, ToolCallID: use.ToolUseID, + Success: success, SessionID: sessionID, Model: ctx.Model, Raw: raw, + }, true +} + +func appServerToolCall(it *appServerItemBody, sessionID, model string) history.CodexToolCall { + name := it.Tool + input := map[string]any{} + switch it.Type { + case "commandExecution": + name = "" + case "fileChange": + name = "CodexPatchApply" + if len(it.Changes) > 0 { + var changes any + _ = json.Unmarshal(it.Changes, &changes) + input["changes"] = changes + } + default: + name = firstNonEmpty(name, it.Type) + } + return history.CodexToolCall{ + Name: name, Namespace: it.Server, Arguments: it.Arguments, + Command: it.Command, Input: input, CWD: it.CWD, + SessionID: sessionID, ID: it.ID, Model: model, + } +} + +func appServerToolResultText(it *appServerItemBody, buffered string) string { + if it.AggregatedOutput != nil && *it.AggregatedOutput != "" { + return *it.AggregatedOutput + } + if buffered != "" { + return buffered + } + if it.AggregatedOutput != nil { + return *it.AggregatedOutput + } + if it.Error != nil { + return firstNonEmpty(it.Error.Message, it.Error.AdditionalDetails) + } + for _, raw := range []json.RawMessage{it.Result, it.ContentItems, it.Changes} { + if len(raw) > 0 && string(raw) != "null" { + return string(raw) + } + } + return it.Text +} + +func appServerToolSucceeded(it *appServerItemBody) bool { + if it.Success != nil { + return *it.Success + } + if it.ExitCode != nil && *it.ExitCode != 0 { + return false + } + return it.Status == "completed" } // appServerErrorIsFatal reports whether an error/turn-failure notification ends @@ -256,10 +335,25 @@ func codexSafety(req ai.Request) (sandbox, approval string) { // non-empty, is sent as the turn-scoped `outputSchema` that constrains the final // assistant message to validated JSON (structured output); the raw JSON Schema // bytes are embedded inline verbatim. -func buildTurnStartParams(model string, req ai.Request, threadID string, outputSchema json.RawMessage) map[string]any { + +func buildTurnStartParams(model string, req ai.Request, threadID string, outputSchema json.RawMessage) (map[string]any, error) { + if err := ai.ValidateAttachmentCompatibility([]api.Model{{Name: model, Backend: api.BackendCodexAgent}}, req.Prompt.Attachments); err != nil { + return nil, err + } + input := make([]map[string]any, 0, len(req.Prompt.Attachments)) + if text := composePrompt(req); text != "" { + input = append(input, map[string]any{"type": "text", "text": text}) + } + for i, attachment := range req.Prompt.Attachments { + content, ok := attachment.PreparedContent() + if !ok || content.Path == "" { + return nil, fmt.Errorf("attachment %d (%s) has no prepared local path", i+1, attachment.ID) + } + input = append(input, map[string]any{"type": "localImage", "path": content.Path}) + } p := map[string]any{ "threadId": threadID, - "input": []map[string]any{{"type": "text", "text": composePrompt(req)}}, + "input": input, } if model != "" { p["model"] = model @@ -270,7 +364,7 @@ func buildTurnStartParams(model string, req ai.Request, threadID string, outputS if len(outputSchema) > 0 { p["outputSchema"] = outputSchema } - return p + return p, nil } func buildResumeParams(req ai.Request) map[string]any { diff --git a/pkg/ai/provider/codex_appserver_test.go b/pkg/ai/provider/codex_appserver_test.go index 220b5b1..63056f8 100644 --- a/pkg/ai/provider/codex_appserver_test.go +++ b/pkg/ai/provider/codex_appserver_test.go @@ -2,6 +2,7 @@ package provider import ( "encoding/json" + "errors" "testing" "github.com/flanksource/captain/pkg/ai" @@ -16,11 +17,16 @@ func TestNewCodexAppServer_Defaults(t *testing.T) { c, err := NewCodexAppServer("") require.NoError(t, err) assert.Equal(t, CodexCLIDefaultModel, c.GetModel()) - assert.Equal(t, ai.BackendCodexCLI, c.GetBackend()) + assert.Equal(t, ai.BackendCodexAgent, c.GetBackend()) - c2, err := NewCodexAppServer("gpt-5-codex") + c2, err := NewCodexAppServer("gpt-5.4") require.NoError(t, err) - assert.Equal(t, "gpt-5-codex", c2.GetModel()) + assert.Equal(t, "gpt-5.4", c2.GetModel()) +} + +func TestAppServerProcessErrorIncludesStderr(t *testing.T) { + err := appServerProcessError(errors.New("jsonrpc: client closed"), " state runtime unavailable \n") + assert.EqualError(t, err, "jsonrpc: client closed: state runtime unavailable") } func TestMapAppServerNotification_Kinds(t *testing.T) { @@ -76,16 +82,10 @@ func TestMapAppServerNotification_Kinds(t *testing.T) { check: func(t *testing.T, ev ai.Event) { assert.Equal(t, "because", ev.Text) }, }, { - name: "command output delta becomes tool use", + name: "command output delta is buffered by the turn", method: "item/commandExecution/outputDelta", params: `{"delta":"line1\n","itemId":"cmd-1"}`, - want: ai.EventToolUse, - check: func(t *testing.T, ev ai.Event) { - assert.Equal(t, "line1\n", ev.Input["delta"]) - tu, ok := ev.Raw.(claude.ToolUse) - require.True(t, ok) - assert.Equal(t, "cmd-1", tu.SessionID) - }, + drop: true, }, { name: "agent message item completed becomes text", @@ -95,29 +95,32 @@ func TestMapAppServerNotification_Kinds(t *testing.T) { check: func(t *testing.T, ev ai.Event) { assert.Equal(t, "final answer", ev.Text) }, }, { - name: "command execution item completed becomes tool use", + name: "command execution item completed becomes tool result", method: "item/completed", - params: `{"item":{"id":"c1","type":"commandExecution","command":"ls -la"}}`, - want: ai.EventToolUse, + params: `{"threadId":"thread-1","item":{"id":"c1","type":"commandExecution","command":"ls -la","status":"completed","exitCode":0}}`, + want: ai.EventToolResult, check: func(t *testing.T, ev ai.Event) { - assert.Equal(t, "ls -la", ev.Tool) - assert.Equal(t, "ls -la", ev.Input["command"]) + assert.Equal(t, "c1", ev.ToolCallID) + assert.True(t, ev.Success) tu, ok := ev.Raw.(claude.ToolUse) require.True(t, ok) - assert.Equal(t, "codex", tu.Source) + assert.Equal(t, "c1", tu.ToolUseID) }, }, { name: "command execution item started becomes tool use", method: "item/started", - params: `{"item":{"id":"c1","type":"commandExecution","command":"pwd"}}`, + params: `{"threadId":"thread-1","item":{"id":"c1","type":"commandExecution","command":"pwd","status":"inProgress"}}`, want: ai.EventToolUse, - check: func(t *testing.T, ev ai.Event) { assert.Equal(t, "pwd", ev.Tool) }, + check: func(t *testing.T, ev ai.Event) { + assert.Equal(t, "Bash", ev.Tool) + assert.Equal(t, "pwd", ev.Input["command"]) + }, }, { name: "mcp tool call item uses tool name", - method: "item/completed", - params: `{"item":{"id":"m1","type":"mcpToolCall","tool":"search"}}`, + method: "item/started", + params: `{"threadId":"thread-1","item":{"id":"m1","type":"mcpToolCall","tool":"search","arguments":{"query":"captain"},"status":"inProgress"}}`, want: ai.EventToolUse, check: func(t *testing.T, ev ai.Event) { assert.Equal(t, "search", ev.Tool) }, }, @@ -160,7 +163,9 @@ func TestMapAppServerNotification_Kinds(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - ev, ok := mapAppServerNotification(tc.method, json.RawMessage(tc.params), "gpt-5", &ai.Usage{}) + ev, ok := mapAppServerNotification(tc.method, json.RawMessage(tc.params), appServerEventContext{ + Model: "gpt-5", Usage: &ai.Usage{}, + }) if tc.drop { assert.False(t, ok, "expected notification to be dropped, got %+v", ev) return @@ -199,6 +204,7 @@ func activeTurn(t *testing.T, schema json.RawMessage) (*CodexAppServer, *turnSta usage: &ai.Usage{}, model: "gpt-5", streamed: map[string]bool{}, + toolOutput: map[string]string{}, terminal: make(chan struct{}), outputSchema: schema, } @@ -217,7 +223,7 @@ func resultEvent(t *testing.T, evs []ai.Event) ai.Event { return ai.Event{} } -// A structured turn attaches the final agentMessage JSON to the terminal result. +// A structured turn attaches the final agentMessage JSON only to the terminal result. func TestHandleNotification_StructuredOutput(t *testing.T) { schema, err := ai.SchemaJSONFor(api.Prompt{Schema: &struct { Answer string `json:"answer"` @@ -225,18 +231,67 @@ func TestHandleNotification_StructuredOutput(t *testing.T) { require.NoError(t, err) c, ts := activeTurn(t, schema) - // The final agent message's text IS the validated JSON (codex has no separate - // structured field). + c.handleNotification("item/agentMessage/delta", + json.RawMessage(`{"itemId":"a1","delta":"{\"answer\":"}`)) + c.handleNotification("item/agentMessage/delta", + json.RawMessage(`{"itemId":"a1","delta":"\"42\"}"}`)) c.handleNotification("item/completed", json.RawMessage(`{"item":{"id":"a1","type":"agentMessage","text":"{\"answer\":\"42\"}"}}`)) c.handleNotification("turn/completed", json.RawMessage(`{"threadId":"t","turn":{"id":"u","status":"completed"}}`)) - result := resultEvent(t, drainEvents(ts)) + events := drainEvents(ts) + var resultCount int + for _, ev := range events { + assert.NotEqual(t, ai.EventText, ev.Kind, "structured agent messages must not emit text progress") + if ev.Kind == ai.EventResult { + resultCount++ + } + } + assert.Equal(t, 1, resultCount, "structured turn should emit exactly one terminal result") + result := resultEvent(t, events) + assert.True(t, result.Success) require.NotEmpty(t, result.StructuredData, "structured turn should carry the final JSON on the result") assert.JSONEq(t, `{"answer":"42"}`, string(result.StructuredData)) } +func TestHandleNotification_StructuredOutputWithoutDeltas(t *testing.T) { + schema := json.RawMessage(`{"type":"object"}`) + c, ts := activeTurn(t, schema) + c.handleNotification("item/completed", + json.RawMessage(`{"item":{"id":"a1","type":"agentMessage","text":"{\"answer\":\"42\"}"}}`)) + c.handleNotification("turn/completed", + json.RawMessage(`{"threadId":"t","turn":{"id":"u","status":"completed"}}`)) + + events := drainEvents(ts) + for _, ev := range events { + assert.NotEqual(t, ai.EventText, ev.Kind, "completed structured message must not leak as text progress") + } + assert.JSONEq(t, `{"answer":"42"}`, string(resultEvent(t, events).StructuredData)) +} + +func TestHandleNotification_TextModeStreamsDeltasAndDeduplicatesCompletedMessage(t *testing.T) { + c, ts := activeTurn(t, nil) + c.handleNotification("item/agentMessage/delta", + json.RawMessage(`{"itemId":"a1","delta":"plain "}`)) + c.handleNotification("item/agentMessage/delta", + json.RawMessage(`{"itemId":"a1","delta":"answer"}`)) + c.handleNotification("item/completed", + json.RawMessage(`{"item":{"id":"a1","type":"agentMessage","text":"plain answer"}}`)) + c.handleNotification("turn/completed", + json.RawMessage(`{"threadId":"t","turn":{"id":"u","status":"completed"}}`)) + + events := drainEvents(ts) + var text []string + for _, ev := range events { + if ev.Kind == ai.EventText { + text = append(text, ev.Text) + } + } + assert.Equal(t, []string{"plain ", "answer"}, text) + assert.Empty(t, resultEvent(t, events).StructuredData) +} + // A text-mode turn (no schema) leaves the result's StructuredData empty. func TestHandleNotification_NoStructuredWithoutSchema(t *testing.T) { c, ts := activeTurn(t, nil) @@ -253,41 +308,49 @@ func TestMapAppServerNotification_ErrorUnwrapping(t *testing.T) { nested := `The 'gpt-5.5-codex' model is not supported when using Codex with a ChatGPT account.` params := `{"threadId":"t","turnId":"u","willRetry":false,"error":{"message":"{\"type\":\"error\",\"status\":400,\"error\":{\"type\":\"invalid_request_error\",\"message\":\"` + nested + `\"}}"}}` - ev, ok := mapAppServerNotification("error", json.RawMessage(params), "gpt-5", &ai.Usage{}) + ev, ok := mapAppServerNotification("error", json.RawMessage(params), appServerEventContext{Model: "gpt-5", Usage: &ai.Usage{}}) require.True(t, ok) assert.Equal(t, ai.EventError, ev.Kind) assert.Equal(t, nested, ev.Error, "stringified upstream error payload should be unwrapped one level") } func TestMapAppServerNotification_TurnFailed(t *testing.T) { - ev, ok := mapAppServerNotification("turn/failed", json.RawMessage(`{"error":{"message":"boom"}}`), "m", &ai.Usage{}) + ev, ok := mapAppServerNotification("turn/failed", json.RawMessage(`{"error":{"message":"boom"}}`), appServerEventContext{Model: "m", Usage: &ai.Usage{}}) require.True(t, ok) assert.Equal(t, ai.EventError, ev.Kind) assert.Equal(t, "boom", ev.Error) } // Token usage folded from thread/tokenUsage/updated must surface on the -// subsequent turn/completed result event. +// subsequent turn/completed result event, normalized to captain's disjoint +// buckets: codex reports inputTokens inclusive of cachedInputTokens and +// outputTokens inclusive of reasoningOutputTokens, so foldUsage nets both out to +// avoid the cache/reasoning double-count (findings B1/B2). func TestMapAppServerNotification_UsageFolding(t *testing.T) { usage := &ai.Usage{} + ctx := appServerEventContext{Model: "m", Usage: usage} _, ok := mapAppServerNotification( "thread/tokenUsage/updated", json.RawMessage(`{"tokenUsage":{"total":{"inputTokens":120,"outputTokens":40,"cachedInputTokens":12,"reasoningOutputTokens":7}}}`), - "m", usage) + ctx) assert.False(t, ok, "token usage update emits no event") - assert.Equal(t, 120, usage.InputTokens) - assert.Equal(t, 40, usage.OutputTokens) + assert.Equal(t, 108, usage.InputTokens, "input net of cache (120-12)") + assert.Equal(t, 33, usage.OutputTokens, "output net of reasoning (40-7)") assert.Equal(t, 12, usage.CacheReadTokens) assert.Equal(t, 7, usage.ReasoningTokens) - ev, ok := mapAppServerNotification("turn/completed", json.RawMessage(`{"threadId":"t","turn":{"id":"u"}}`), "m", usage) + ev, ok := mapAppServerNotification("turn/completed", json.RawMessage(`{"threadId":"t","turn":{"id":"u"}}`), ctx) require.True(t, ok) require.NotNil(t, ev.Usage, "turn/completed should carry the folded usage") - assert.Equal(t, 120, ev.Usage.InputTokens) - assert.Equal(t, 40, ev.Usage.OutputTokens) + assert.Equal(t, 108, ev.Usage.InputTokens) + assert.Equal(t, 33, ev.Usage.OutputTokens) assert.Equal(t, 12, ev.Usage.CacheReadTokens) assert.Equal(t, 7, ev.Usage.ReasoningTokens) + // Disjoint buckets: netted input+cache and output+reasoning recover the raw + // codex totals, so pricing cannot bill the overlap twice. + assert.Equal(t, 120, ev.Usage.InputTokens+ev.Usage.CacheReadTokens) + assert.Equal(t, 40, ev.Usage.OutputTokens+ev.Usage.ReasoningTokens) } func TestAppServerErrorIsFatal(t *testing.T) { @@ -349,13 +412,14 @@ func req(p api.Prompt) ai.Request { } func TestBuildTurnStartParams(t *testing.T) { - p := buildTurnStartParams("gpt-5", ai.Request{ + p, err := buildTurnStartParams("gpt-5", ai.Request{ Prompt: api.Prompt{User: "hi"}, - Model: api.Model{Effort: api.EffortHigh}, + Model: api.Model{Effort: api.EffortUltra}, }, "thread-1", nil) + require.NoError(t, err) assert.Equal(t, "thread-1", p["threadId"]) assert.Equal(t, "gpt-5", p["model"]) - assert.Equal(t, "high", p["effort"]) + assert.Equal(t, "ultra", p["effort"]) input, ok := p["input"].([]map[string]any) require.True(t, ok, "input must be a slice of text elements") @@ -364,7 +428,8 @@ func TestBuildTurnStartParams(t *testing.T) { assert.Equal(t, "hi", input[0]["text"]) // No reasoning effort / empty model / nil schema should be omitted entirely. - bare := buildTurnStartParams("", req(api.Prompt{User: "hi"}), "t", nil) + bare, err := buildTurnStartParams("", req(api.Prompt{User: "hi"}), "t", nil) + require.NoError(t, err) _, hasModel := bare["model"] _, hasEffort := bare["effort"] _, hasSchema := bare["outputSchema"] @@ -376,11 +441,14 @@ func TestBuildTurnStartParams(t *testing.T) { func TestBuildTurnStartParams_OutputSchema(t *testing.T) { type answer struct { Answer string `json:"answer"` + Detail string `json:"detail,omitempty"` } - schema, err := ai.SchemaJSONFor(api.Prompt{Schema: &answer{}}) + request := req(api.Prompt{User: "solve", Schema: &answer{}}) + schema, err := codexOutputSchema(request) require.NoError(t, err) - p := buildTurnStartParams("gpt-5", req(api.Prompt{User: "solve"}), "t", schema) + p, err := buildTurnStartParams("gpt-5", request, "t", schema) + require.NoError(t, err) require.Contains(t, p, "outputSchema", "a derived schema must be sent as outputSchema") // It must serialize to a JSON Schema object describing the target struct. @@ -392,6 +460,9 @@ func TestBuildTurnStartParams_OutputSchema(t *testing.T) { props, ok := decoded["properties"].(map[string]any) require.True(t, ok) assert.Contains(t, props, "answer") + assert.Contains(t, props, "detail") + assert.Equal(t, []any{"answer", "detail"}, decoded["required"]) + assert.Equal(t, false, decoded["additionalProperties"]) } func TestBuildThreadStartParams_Safety(t *testing.T) { diff --git a/pkg/ai/provider/codex_appserver_turn.go b/pkg/ai/provider/codex_appserver_turn.go new file mode 100644 index 0000000..0f5e0ce --- /dev/null +++ b/pkg/ai/provider/codex_appserver_turn.go @@ -0,0 +1,89 @@ +package provider + +import ( + "context" + "encoding/json" + "fmt" + "sync" + + "github.com/flanksource/captain/pkg/ai" +) + +type turnState struct { + ch chan ai.Event + usage *ai.Usage + model string + streamed map[string]bool + toolOutput map[string]string + + outputSchema json.RawMessage + lastAgentMessage string + terminal chan struct{} + termOnce sync.Once + started chan struct{} + startOnce sync.Once + idMu sync.Mutex + threadID string + turnID string + sendMu sync.Mutex + closed bool +} + +func (ts *turnState) signalTerminal() { ts.termOnce.Do(func() { close(ts.terminal) }) } + +func (ts *turnState) setIDs(threadID, turnID string) { + ts.idMu.Lock() + ts.threadID = threadID + ts.turnID = turnID + ts.idMu.Unlock() + ts.signalStarted() +} + +func (ts *turnState) signalStarted() { + if ts.started != nil { + ts.startOnce.Do(func() { close(ts.started) }) + } +} + +func (ts *turnState) waitIDs(ctx context.Context) (string, string, error) { + if ts.started == nil { + return "", "", fmt.Errorf("codex app-server: turn identifiers unavailable") + } + select { + case <-ts.started: + case <-ts.terminal: + return "", "", fmt.Errorf("codex app-server: turn ended before interrupt was ready") + case <-ctx.Done(): + return "", "", ctx.Err() + } + ts.idMu.Lock() + defer ts.idMu.Unlock() + if ts.threadID == "" { + return "", "", fmt.Errorf("codex app-server: turn identifiers unavailable") + } + return ts.threadID, ts.turnID, nil +} + +func (ts *turnState) send(event ai.Event) { + ts.sendMu.Lock() + defer ts.sendMu.Unlock() + if ts.closed { + return + } + select { + case ts.ch <- event: + case <-ts.terminal: + } +} + +func (ts *turnState) finish() { + ts.signalStarted() + ts.signalTerminal() + ts.sendMu.Lock() + defer ts.sendMu.Unlock() + if ts.closed { + return + } + ts.closed = true + close(ts.ch) +} diff --git a/pkg/ai/provider/codex_cli.go b/pkg/ai/provider/codex_cli.go new file mode 100644 index 0000000..5b83913 --- /dev/null +++ b/pkg/ai/provider/codex_cli.go @@ -0,0 +1,397 @@ +package provider + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "github.com/flanksource/captain/pkg/ai" + history "github.com/flanksource/captain/pkg/ai/history" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" +) + +var codexCLICommand = "codex" + +type CodexCLI struct { + model string +} + +func NewCodexCLI(model string) *CodexCLI { + if strings.TrimSpace(model) == "" { + model = CodexCLIDefaultModel + } + model = ai.NormalizeModelForBackend(ai.BackendCodexCLI, model) + return &CodexCLI{model: model} +} + +func (c *CodexCLI) GetModel() string { return c.model } +func (c *CodexCLI) GetBackend() ai.Backend { return ai.BackendCodexCLI } + +func (c *CodexCLI) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { + start := time.Now() + events, err := c.ExecuteStream(ctx, req) + if err != nil { + return nil, err + } + resp, err := CoalesceStreamForBackend(ctx, ai.BackendCodexCLI, c.model, events, start) + if err != nil { + return nil, err + } + if req.Prompt.Schema != nil { + raw, _ := resp.StructuredData.(json.RawMessage) + if len(raw) == 0 { + raw = json.RawMessage(resp.Text) + } + if err := ai.BindStructuredOutput(req.Prompt.Schema, raw); err != nil { + return nil, err + } + resp.StructuredData = req.Prompt.Schema + resp.Text = "" + } + return resp, nil +} + +func (c *CodexCLI) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai.Event, error) { + args, cleanup, err := buildCodexCLIArgs(c.model, req) + if err != nil { + return nil, err + } + env := []string(nil) + if req.Setup != nil { + env = commandEnv(req.Setup.Env) + } + cmd, stdout, stderrBuf, err := startCLIStream(ctx, codexCLICommand, args, []byte(composePrompt(req)), req.Cwd(), env) + if err != nil { + cleanup() + return nil, err + } + out := make(chan ai.Event, 16) + go func() { + defer close(out) + defer cleanup() + defer func() { _ = stdout.Close() }() + scanner := bufio.NewScanner(stdout) + scanner.Buffer(make([]byte, 1024*1024), 10*1024*1024) + state := codexCLIState{model: c.model, cwd: req.Cwd(), pending: map[string]history.CodexEvent{}} + for scanner.Scan() { + for _, ev := range state.mapLine(scanner.Bytes()) { + emit(ctx, out, ev) + } + } + if err := scanner.Err(); err != nil { + emit(ctx, out, ai.Event{Kind: ai.EventError, Error: fmt.Sprintf("codex exec --json: %v", err), Model: c.model}) + } + if err := finishCLIStream(ctx, cmd, stderrBuf); err != nil { + emit(ctx, out, ai.Event{Kind: ai.EventError, Error: err.Error(), Model: c.model}) + } + }() + return out, nil +} + +func buildCodexCLIArgs(model string, req ai.Request) ([]string, func(), error) { + args := []string{"exec", "--json"} + cleanup := func() {} + if err := ai.ValidateAttachmentCompatibility([]api.Model{{Name: model, Backend: api.BackendCodexCLI}}, req.Prompt.Attachments); err != nil { + return nil, cleanup, err + } + for i, attachment := range req.Prompt.Attachments { + content, ok := attachment.PreparedContent() + if !ok || content.Path == "" { + return nil, cleanup, fmt.Errorf("attachment %d (%s) has no prepared local path", i+1, attachment.ID) + } + args = append(args, "--image", content.Path) + } + if req.Effort != api.EffortNone { + args = append(args, "-c", fmt.Sprintf("model_reasoning_effort=%q", req.Effort)) + } + if m := strings.TrimSpace(model); m != "" && m != "codex" { + args = append(args, "-m", m) + } + if cwd := req.Cwd(); cwd != "" { + args = append(args, "-C", cwd) + } + sandbox, _ := codexSafety(req) + if sandbox != "" { + args = append(args, "--sandbox", sandbox) + } + if req.Memory.SkipMemory || req.Memory.Bare || req.Permissions.HasPreset(api.PresetBare) { + args = append(args, "--ephemeral") + } + if req.Memory.SkipUser { + args = append(args, "--ignore-user-config") + } + if req.Memory.SkipProject || req.Memory.SkipHooks { + args = append(args, "--ignore-rules") + } + if binary, ok := captainBinary(); ok && api.MonitorHooksEnabled(req) { + notify, err := codexNotifyOverride(binary) + if err != nil { + return nil, cleanup, err + } + args = append(args, "-c", notify) + } + schema, err := ai.SchemaJSONForBackend(ai.BackendCodexCLI, req.Prompt) + if err != nil { + return nil, cleanup, fmt.Errorf("codex-cli: cannot derive structured-output schema: %w", err) + } + if len(schema) > 0 { + path, err := writeTempSchema(schema) + if err != nil { + return nil, cleanup, err + } + cleanup = func() { _ = os.Remove(path) } + args = append(args, "--output-schema", path) + } + if req.SessionID != "" { + args = append(args, "resume", req.SessionID) + } + return args, cleanup, nil +} + +func writeTempSchema(schema json.RawMessage) (string, error) { + f, err := os.CreateTemp("", "captain-codex-schema-*.json") + if err != nil { + return "", fmt.Errorf("codex-cli: create schema file: %w", err) + } + path := f.Name() + if _, err := f.Write(schema); err != nil { + _ = f.Close() + _ = os.Remove(path) + return "", fmt.Errorf("codex-cli: write schema file %s: %w", filepath.Base(path), err) + } + if err := f.Close(); err != nil { + _ = os.Remove(path) + return "", fmt.Errorf("codex-cli: close schema file %s: %w", filepath.Base(path), err) + } + return path, nil +} + +type codexCLIState struct { + model string + cwd string + sessionID string + pending map[string]history.CodexEvent + usage ai.Usage + // sawTokenCount records whether a token_count event supplied the authoritative + // (cache/reasoning-aware) usage, so turn.completed's coarser per-turn counts + // are used only as a fallback. + sawTokenCount bool +} + +func (s *codexCLIState) mapLine(line []byte) []ai.Event { + var event history.CodexEvent + if err := json.Unmarshal(line, &event); err != nil { + return []ai.Event{{Kind: ai.EventError, Error: fmt.Sprintf("codex json parse: %v", err), Model: s.model}} + } + switch event.Type { + case "thread.started": + if event.ThreadID != "" { + s.sessionID = event.ThreadID + } + ev := ai.Event{Kind: ai.EventSystem, Tool: "SessionInit", SessionID: s.sessionID, Model: s.model} + ev.Raw = codexSessionToolUse(s.sessionID, s.model) + return []ai.Event{ev} + case "item.completed": + return s.mapItemCompleted(event) + case "response_item": + return s.mapResponseItem(event) + case "event_msg": + return s.mapEventMsg(event) + case "turn.failed", "error": + msg := extractCodexErrorText(codexErrorMessage(event)) + ev := ai.Event{Kind: ai.EventError, Error: msg, Model: s.model} + ev.Raw = codexToolUse(history.NormalizeCodexToolCall(history.CodexToolCall{ + Name: "ApiError", Input: map[string]any{"error": msg}, SessionID: s.sessionID, + }), s.model) + return []ai.Event{ev} + case "turn.completed": + // turn.completed's per-turn usage lacks cache/reasoning detail; prefer the + // token_count totals when we have them and fall back to these coarse counts + // only otherwise (they cannot overlap cache, so no netting is needed). + if event.Usage != nil && !s.sawTokenCount { + s.usage.InputTokens = event.Usage.InputTokens + s.usage.OutputTokens = event.Usage.OutputTokens + } + usage := s.usage + ev := ai.Event{Kind: ai.EventResult, Tool: "Result", Success: true, SessionID: s.sessionID, Model: s.model, Usage: &usage} + ev.Raw = codexResultToolUse(ev, s.sessionID) + return []ai.Event{ev} + } + return nil +} + +func (s *codexCLIState) mapResponseItem(event history.CodexEvent) []ai.Event { + switch event.Payload.Type { + case "function_call": + s.pending[event.Payload.CallID] = event + return nil + case "function_call_output": + call, ok := s.pending[event.Payload.CallID] + if !ok { + return nil + } + delete(s.pending, event.Payload.CallID) + return s.codexToolEvents([]claude.ToolUse{s.buildFunctionToolUse(call, event)}) + case "reasoning": + text := codexReasoningText(event) + if text == "" { + return nil + } + return []ai.Event{{Kind: ai.EventThinking, Text: text, SessionID: s.sessionID, Model: s.model}} + case "message": + text := codexMessageText(event) + if text == "" { + return nil + } + return []ai.Event{{Kind: ai.EventText, Text: text, SessionID: s.sessionID, Model: s.model}} + } + return nil +} + +func (s *codexCLIState) mapEventMsg(event history.CodexEvent) []ai.Event { + if event.Payload.Type == "token_count" && event.Payload.Info != nil { + // Codex reports cumulative totals with inputTokens inclusive of the cached + // prefix and outputTokens inclusive of reasoning; net both to captain's + // disjoint buckets so pricing does not double-count (findings B1/B2). + tu := event.Payload.Info.TotalTokenUsage + s.usage.InputTokens = ai.NetInputTokens(tu.InputTokens, tu.CachedInputTokens) + s.usage.OutputTokens = ai.NetOutputTokens(tu.OutputTokens, tu.ReasoningOutputTokens) + s.usage.CacheReadTokens = tu.CachedInputTokens + s.usage.ReasoningTokens = tu.ReasoningOutputTokens + s.sawTokenCount = true + return nil + } + switch event.Payload.Type { + case "agent_reasoning": + if event.Payload.Text != "" { + return []ai.Event{{Kind: ai.EventThinking, Text: event.Payload.Text, SessionID: s.sessionID, Model: s.model}} + } + case "agent_message": + if event.Payload.Message != "" { + return []ai.Event{{Kind: ai.EventText, Text: event.Payload.Message, SessionID: s.sessionID, Model: s.model}} + } + } + return nil +} + +func (s *codexCLIState) mapItemCompleted(event history.CodexEvent) []ai.Event { + if event.Item == nil { + return nil + } + text := event.Item.Text + if text == "" { + for _, content := range event.Item.Content { + if content.Type == "output_text" && content.Text != "" { + text += content.Text + } + } + } + if text != "" { + return []ai.Event{{Kind: ai.EventText, Text: text, SessionID: s.sessionID, Model: s.model}} + } + name := firstNonEmpty(event.Item.Name, event.Item.Type) + if name == "" { + return nil + } + tu := codexToolUse(history.NormalizeCodexToolCall(history.CodexToolCall{ + Name: name, ID: event.Item.Name, SessionID: s.sessionID, + }), s.model) + return s.codexToolEvents([]claude.ToolUse{tu}) +} + +func (s *codexCLIState) codexToolEvents(toolUses []claude.ToolUse) []ai.Event { + out := make([]ai.Event, 0, len(toolUses)) + for _, tu := range toolUses { + ev := ai.Event{ + Kind: ai.EventToolUse, + Tool: tu.Tool, + Input: tu.Input, + ToolCallID: tu.ToolUseID, + SessionID: firstNonEmpty(tu.SessionID, s.sessionID), + Model: firstNonEmpty(tu.Model, s.model), + Raw: tu, + } + if tu.Tool == "Assistant" { + if text, _ := tu.Input["text"].(string); text != "" { + ev.Kind = ai.EventText + ev.Text = text + } + } + if tu.Tool == "Reasoning" { + if text, _ := tu.Input["text"].(string); text != "" { + ev.Kind = ai.EventThinking + ev.Text = text + } + } + if tu.Tool == "ApiError" { + ev.Kind = ai.EventError + ev.Error = firstString(tu.Input, "error", "message", "result") + } + out = append(out, ev) + } + return out +} + +func codexErrorMessage(event history.CodexEvent) string { + if event.Error != nil && event.Error.Message != "" { + return event.Error.Message + } + if event.Message != "" { + return event.Message + } + return event.Payload.Message +} + +func (s *codexCLIState) buildFunctionToolUse(call, output history.CodexEvent) claude.ToolUse { + return codexToolUse(history.NormalizeCodexToolCall(history.CodexToolCall{ + Name: call.Payload.Name, + Namespace: call.Payload.Namespace, + Arguments: call.Payload.Arguments, + ID: call.Payload.CallID, + SessionID: s.sessionID, + Response: history.CodexOutputText(output.Payload.Output), + RecordType: "response_item.function_call", + }), s.model) +} + +func codexReasoningText(event history.CodexEvent) string { + if len(event.Payload.Summary) == 0 { + return event.Payload.Text + } + var summaries []history.CodexReasoningSummary + if err := json.Unmarshal(event.Payload.Summary, &summaries); err == nil { + var text string + for _, summary := range summaries { + if summary.Text != "" { + text = summary.Text + } + } + if text != "" { + return text + } + } + var text string + if err := json.Unmarshal(event.Payload.Summary, &text); err == nil { + return text + } + return event.Payload.Text +} + +func codexMessageText(event history.CodexEvent) string { + if event.Payload.Role != "" && event.Payload.Role != "assistant" { + return "" + } + var text string + for _, content := range event.Payload.Content { + if content.Type == "output_text" && content.Text != "" { + text += content.Text + } + } + return text +} diff --git a/pkg/ai/provider/codex_cli_test.go b/pkg/ai/provider/codex_cli_test.go new file mode 100644 index 0000000..7778238 --- /dev/null +++ b/pkg/ai/provider/codex_cli_test.go @@ -0,0 +1,131 @@ +package provider + +import ( + "encoding/json" + "os" + "reflect" + "testing" + + "github.com/flanksource/captain/pkg/ai" + history "github.com/flanksource/captain/pkg/ai/history" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/commons-db/shell" +) + +func TestBuildCodexCLIArgs(t *testing.T) { + cwd := t.TempDir() + req := ai.Request{ + Model: api.Model{Effort: api.EffortUltra}, + Prompt: api.Prompt{ + User: "hello", + SchemaJSON: json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}}}`), + }, + Setup: &shell.Setup{Cwd: cwd}, + SessionID: "thread-1", + Memory: api.Memory{ + SkipMemory: true, + SkipUser: true, + SkipProject: true, + SkipHooks: true, + }, + Permissions: api.Permissions{Presets: []api.Preset{api.PresetEdit}}, + } + + args, cleanup, err := buildCodexCLIArgs("gpt-5.5", req) + if err != nil { + t.Fatalf("buildCodexCLIArgs: %v", err) + } + defer cleanup() + requireFlagValue(t, args, "-m", "gpt-5.5") + requireFlagValue(t, args, "-c", `model_reasoning_effort="ultra"`) + requireFlagValue(t, args, "-C", cwd) + requireFlagValue(t, args, "--sandbox", "workspace-write") + requireHasArg(t, args, "--json") + requireHasArg(t, args, "--ephemeral") + requireHasArg(t, args, "--ignore-user-config") + requireHasArg(t, args, "--ignore-rules") + if got := args[len(args)-2:]; got[0] != "resume" || got[1] != "thread-1" { + t.Fatalf("args suffix = %v, want resume thread-1 (args: %v)", got, args) + } + schemaPath := flagValue(t, args, "--output-schema") + data, err := os.ReadFile(schemaPath) + if err != nil { + t.Fatalf("read schema file: %v", err) + } + if !json.Valid(data) { + t.Fatalf("schema file is not valid JSON: %s", data) + } + var schema map[string]any + if err := json.Unmarshal(data, &schema); err != nil { + t.Fatalf("decode schema file: %v", err) + } + if got := schema["required"]; !reflect.DeepEqual(got, []any{"answer"}) { + t.Fatalf("schema required = %#v, want [answer]", got) + } + if schema["additionalProperties"] != false { + t.Fatalf("schema additionalProperties = %v, want false", schema["additionalProperties"]) + } + cleanup() + if _, err := os.Stat(schemaPath); !os.IsNotExist(err) { + t.Fatalf("schema cleanup stat err = %v, want not exist", err) + } +} + +func TestCodexCLIStateMapsJSONLEvents(t *testing.T) { + state := codexCLIState{model: "gpt-5.5", pending: map[string]history.CodexEvent{}} + + events := state.mapLine([]byte(`{"type":"thread.started","thread_id":"thread-1"}`)) + if len(events) != 1 || events[0].Kind != ai.EventSystem || events[0].SessionID != "thread-1" { + t.Fatalf("thread.started events = %+v", events) + } + + events = state.mapLine([]byte(`{"type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}}`)) + if len(events) != 1 || events[0].Kind != ai.EventText || events[0].Text != "hello" { + t.Fatalf("message events = %+v", events) + } + + events = state.mapLine([]byte(`{"type":"response_item","payload":{"type":"function_call","name":"shell","call_id":"call-1","arguments":"{\"command\":\"pwd\"}"}}`)) + if len(events) != 0 { + t.Fatalf("function_call events = %+v, want none until output", events) + } + events = state.mapLine([]byte(`{"type":"response_item","payload":{"type":"function_call_output","call_id":"call-1","output":"ok"}}`)) + if len(events) != 1 || events[0].Kind != ai.EventToolUse || events[0].Tool != "Bash" { + t.Fatalf("function_call_output events = %+v", events) + } + if got, _ := events[0].Input["command"].(string); got != "pwd" { + t.Fatalf("tool input command = %q, want pwd", got) + } + + events = state.mapLine([]byte(`{"type":"turn.completed","usage":{"input_tokens":3,"output_tokens":2}}`)) + if len(events) != 1 || events[0].Kind != ai.EventResult || !events[0].Success { + t.Fatalf("turn.completed events = %+v", events) + } + if events[0].Usage == nil || events[0].Usage.InputTokens != 3 || events[0].Usage.OutputTokens != 2 { + t.Fatalf("usage = %+v", events[0].Usage) + } +} + +// A token_count event supplies cache/reasoning-aware totals; codex reports +// input_tokens inclusive of cache and output_tokens inclusive of reasoning, so +// the emitted usage must be netted to disjoint buckets (findings B1/B2) and the +// coarser turn.completed per-turn counts must not clobber it. +func TestCodexCLIStateNetsTokenCountUsage(t *testing.T) { + state := codexCLIState{model: "gpt-5.5", pending: map[string]history.CodexEvent{}} + + events := state.mapLine([]byte(`{"type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":120,"cached_input_tokens":12,"output_tokens":40,"reasoning_output_tokens":7}}}}`)) + if len(events) != 0 { + t.Fatalf("token_count should emit no events, got %+v", events) + } + + events = state.mapLine([]byte(`{"type":"turn.completed","usage":{"input_tokens":999,"output_tokens":999}}`)) + if len(events) != 1 || events[0].Usage == nil { + t.Fatalf("turn.completed events = %+v", events) + } + u := events[0].Usage + if u.InputTokens != 108 || u.OutputTokens != 33 || u.CacheReadTokens != 12 || u.ReasoningTokens != 7 { + t.Fatalf("usage = %+v, want input=108 output=33 cache=12 reasoning=7 (netted, not clobbered by turn.completed)", u) + } + if u.InputTokens+u.CacheReadTokens != 120 || u.OutputTokens+u.ReasoningTokens != 40 { + t.Fatalf("disjoint buckets must recover raw totals: %+v", u) + } +} diff --git a/pkg/ai/provider/codex_events.go b/pkg/ai/provider/codex_events.go index 903e33d..bb5761e 100644 --- a/pkg/ai/provider/codex_events.go +++ b/pkg/ai/provider/codex_events.go @@ -1,10 +1,8 @@ package provider import ( - "encoding/json" - "strings" - "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/history" "github.com/flanksource/captain/pkg/claude" ) @@ -14,38 +12,8 @@ import ( // in favour of live `/v1/models` listings. const CodexCLIDefaultModel = "" -// codexError is the error envelope codex emits; extractCodexErrorText unwraps a -// JSON-encoded payload nested in a `message` field into this shape. -type codexError struct { - Type string `json:"type"` - Message string `json:"message"` - Status int `json:"status"` -} - -// extractCodexErrorText pulls a human-readable message out of codex's error -// envelopes. Codex sometimes nests JSON-encoded payloads inside the `message` -// field (e.g. an OpenAI 400 served back as a stringified error envelope), so we -// best-effort unwrap one level when the message itself parses as JSON. func extractCodexErrorText(raw string) string { - raw = strings.TrimSpace(raw) - if raw == "" || !strings.HasPrefix(raw, "{") { - return raw - } - var nested codexError - wrapper := struct { - Error *codexError `json:"error"` - Message string `json:"message"` - }{Error: &nested} - if err := json.Unmarshal([]byte(raw), &wrapper); err != nil { - return raw - } - if wrapper.Error != nil && wrapper.Error.Message != "" { - return wrapper.Error.Message - } - if wrapper.Message != "" { - return wrapper.Message - } - return raw + return history.NormalizeCodexError(raw) } func firstNonEmpty(values ...string) string { @@ -57,17 +25,25 @@ func firstNonEmpty(values ...string) string { return "" } -// codexToolUse builds the claude.ToolUse stand-in stashed on ai.Event.Raw so the -// shared lineRenderer in pkg/cli renders codex live streams identically to -// `captain history` output. Source is hard-coded to "codex" so -// sessionKey/sessionHeaderText pick the codex icon and label. -func codexToolUse(name string, input map[string]any, sessionID, model string) claude.ToolUse { +// codexToolUse converts the shared normalized history shape into the stand-in +// used by the live renderer. +func codexToolUse(use history.ToolUse, fallbackModel string) claude.ToolUse { return claude.ToolUse{ - Tool: name, - Input: input, - SessionID: sessionID, - Source: "codex", - Model: model, + Tool: use.Tool, + Input: use.Input, + Timestamp: use.Timestamp, + CWD: use.CWD, + SessionID: use.SessionID, + ToolUseID: use.ToolUseID, + Source: "codex", + Model: firstNonEmpty(use.Model, fallbackModel), + ReasoningEffort: use.ReasoningEffort, + InputTokens: use.InputTokens + use.CacheReadTokens, + OutputTokens: use.OutputTokens, + Response: use.Response, + AgentID: use.AgentID, + AgentType: use.AgentType, + AgentDesc: use.AgentDesc, } } diff --git a/pkg/ai/provider/gemini_cli.go b/pkg/ai/provider/gemini_cli.go index 3b00417..1cebbd9 100644 --- a/pkg/ai/provider/gemini_cli.go +++ b/pkg/ai/provider/gemini_cli.go @@ -1,6 +1,8 @@ package provider import ( + "bufio" + "bytes" "context" "encoding/json" "fmt" @@ -8,87 +10,283 @@ import ( "time" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" ) +var geminiCLICommand = "gemini" + +const geminiCLIDefaultModel = "gemini-3.5-flash" + type GeminiCLI struct { model string } +// registry.Google declares Streaming for ModeCLI, and the logging/validating +// middleware type-asserts on this interface before calling ExecuteStream. +var _ ai.StreamingProvider = (*GeminiCLI)(nil) + func NewGeminiCLI(model string) *GeminiCLI { - if model == "" { - model = "gemini-cli-pro" + if strings.TrimSpace(model) == "" { + model = geminiCLIDefaultModel } + model = ai.NormalizeModelForBackend(ai.BackendGeminiCLI, model) return &GeminiCLI{model: model} } func (g *GeminiCLI) GetModel() string { return g.model } func (g *GeminiCLI) GetBackend() ai.Backend { return ai.BackendGeminiCLI } -type geminiCLIResponse struct { - Text string `json:"text,omitempty"` - Error string `json:"error,omitempty"` - Usage *struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - } `json:"usage,omitempty"` -} - func (g *GeminiCLI) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { start := time.Now() - - timeout := 120 * time.Second - if deadline, ok := ctx.Deadline(); ok { - if ctxTimeout := time.Until(deadline); ctxTimeout < timeout { - timeout = ctxTimeout + events, err := g.ExecuteStream(ctx, req) + if err != nil { + return nil, err + } + resp, err := CoalesceStreamForBackend(ctx, ai.BackendGeminiCLI, g.model, events, start) + if err != nil { + return nil, err + } + if req.Prompt.Schema != nil { + raw, _ := resp.StructuredData.(json.RawMessage) + if err := ai.BindStructuredOutput(req.Prompt.Schema, raw); err != nil { + return nil, err } + resp.StructuredData = req.Prompt.Schema + resp.Text = "" } - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() + return resp, nil +} - cliReq := map[string]string{ - "prompt": req.Prompt.User, - "model": g.model, +func (g *GeminiCLI) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai.Event, error) { + // gemini has no structured-output flag, so the schema rides in the prompt and + // the reply is validated against it when the terminal result arrives. + req, schema, err := ai.WithSchemaPrompt(req) + if err != nil { + return nil, fmt.Errorf("gemini-cli: cannot derive structured-output schema: %w", err) } - reqBytes, err := json.Marshal(cliReq) + args, err := buildGeminiCLIArgs(g.model, req) if err != nil { - return nil, fmt.Errorf("failed to marshal request: %w", err) + return nil, err } - - var env []string + env := []string(nil) if req.Setup != nil { - env = req.Setup.Env + env = commandEnv(req.Setup.Env) } - stdoutData, _, err := runCLI(ctx, "gemini", reqBytes, req.Cwd(), env) + // gemini reads the prompt from stdin and goes headless whenever stdin or + // stdout is not a TTY, which piping both guarantees. + cmd, stdout, stderrBuf, err := startCLIStream(ctx, geminiCLICommand, args, []byte(composePrompt(req)), req.Cwd(), env) if err != nil { return nil, err } + out := make(chan ai.Event, 16) + go func() { + defer close(out) + defer func() { _ = stdout.Close() }() + scanner := bufio.NewScanner(stdout) + scanner.Buffer(make([]byte, 1024*1024), 10*1024*1024) + state := &geminiCLIState{model: g.model, schema: schema} + for scanner.Scan() { + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 { + continue + } + for _, ev := range state.mapLine(line) { + emit(ctx, out, ev) + } + } + if err := scanner.Err(); err != nil { + emit(ctx, out, ai.Event{Kind: ai.EventError, Error: fmt.Sprintf("gemini stream-json: %v", err), Model: g.model}) + } + if err := finishCLIStream(ctx, cmd, stderrBuf); err != nil { + emit(ctx, out, ai.Event{Kind: ai.EventError, Error: err.Error(), Model: g.model}) + } + }() + return out, nil +} - var cliResp geminiCLIResponse - stdoutStr := string(stdoutData) - if idx := strings.Index(stdoutStr, "{"); idx >= 0 { - stdoutStr = stdoutStr[idx:] +// buildGeminiCLIArgs maps the request onto gemini's headless flags. The knobs +// gemini has no equivalent for are deliberately absent rather than approximated: +// there is no system-prompt flag (system text is composed into the prompt), no +// per-run MCP override, no memory/hook opt-outs, and --resume takes a session +// index rather than an id — which is why the registry declares no Resume for +// this backend. +func buildGeminiCLIArgs(model string, req ai.Request) ([]string, error) { + if err := ai.ValidateAttachmentCompatibility([]api.Model{{Name: model, Backend: api.BackendGeminiCLI}}, req.Prompt.Attachments); err != nil { + return nil, err + } + args := []string{"--output-format", "stream-json"} + if m := strings.TrimSpace(model); m != "" { + args = append(args, "--model", m) + } + if mode := geminiApprovalMode(req.Permissions); mode != "" { + args = append(args, "--approval-mode", mode) } + return args, nil +} - if err := json.Unmarshal([]byte(stdoutStr), &cliResp); err != nil { - return nil, fmt.Errorf("failed to parse gemini response: %w (output: %s)", err, stdoutStr) +// geminiApprovalMode maps the permission posture onto gemini's --approval-mode +// choices (default | auto_edit | yolo | plan). The default posture emits no flag +// so gemini keeps its own default, under which a tool needing confirmation fails +// loudly instead of silently running. +func geminiApprovalMode(p api.Permissions) string { + switch { + case p.Mode == api.PermissionPlan: + return "plan" + case p.Mode == api.PermissionBypass, p.Mode == api.PermissionDontAsk: + return "yolo" + case p.Mode == api.PermissionAcceptEdits, p.Mode == api.PermissionAuto: + return "auto_edit" + case p.Mode == "" && p.HasPreset(api.PresetEdit): + return "auto_edit" + default: + return "" } +} - if cliResp.Error != "" { - return nil, fmt.Errorf("gemini CLI error: %s", cliResp.Error) +// geminiStreamEvent is one line of gemini's --output-format stream-json (the +// core StreamJsonFormatter): init | message | tool_use | tool_result | error | +// result, one JSON object per line. +type geminiStreamEvent struct { + Type string `json:"type"` + SessionID string `json:"session_id,omitempty"` + Model string `json:"model,omitempty"` + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + ToolName string `json:"tool_name,omitempty"` + ToolID string `json:"tool_id,omitempty"` + Params map[string]any `json:"parameters,omitempty"` + Status string `json:"status,omitempty"` + Output string `json:"output,omitempty"` + Severity string `json:"severity,omitempty"` + Message string `json:"message,omitempty"` + Error *geminiError `json:"error,omitempty"` + Stats *geminiRunStats `json:"stats,omitempty"` +} + +type geminiError struct { + Type string `json:"type,omitempty"` + Message string `json:"message,omitempty"` +} + +// geminiRunStats is convertToStreamStats' output. InputTokens is the gross +// prompt count (cache included) while Input is already net of Cached; captain's +// buckets are disjoint, so the netting happens here rather than trusting either +// field alone. +type geminiRunStats struct { + TotalTokens int `json:"total_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + Cached int `json:"cached"` + Input int `json:"input"` + DurationMS int `json:"duration_ms"` + ToolCalls int `json:"tool_calls"` +} + +type geminiCLIState struct { + model string + sessionID string + schema json.RawMessage + text strings.Builder +} + +func (s *geminiCLIState) mapLine(line []byte) []ai.Event { + var event geminiStreamEvent + if err := json.Unmarshal(line, &event); err != nil { + return []ai.Event{{Kind: ai.EventError, Error: fmt.Sprintf("gemini json parse: %v", err), Model: s.model}} + } + switch event.Type { + case "init": + if event.SessionID != "" { + s.sessionID = event.SessionID + } + if event.Model != "" { + s.model = event.Model + } + return []ai.Event{{Kind: ai.EventSystem, Tool: "SessionInit", SessionID: s.sessionID, Model: s.model}} + case "message": + // role=user is gemini echoing the prompt back; only assistant deltas are + // response text. + if event.Role != "assistant" || event.Content == "" { + return nil + } + s.text.WriteString(event.Content) + return []ai.Event{{Kind: ai.EventText, Text: event.Content, SessionID: s.sessionID, Model: s.model}} + case "tool_use": + return []ai.Event{{ + Kind: ai.EventToolUse, + Tool: event.ToolName, + Input: event.Params, + ToolCallID: event.ToolID, + SessionID: s.sessionID, + Model: s.model, + }} + case "tool_result": + ev := ai.Event{ + Kind: ai.EventToolResult, + Text: event.Output, + ToolCallID: event.ToolID, + Success: event.Status != "error", + SessionID: s.sessionID, + Model: s.model, + } + if !ev.Success { + ev.Error = geminiErrorMessage(event) + if ev.Text == "" { + ev.Text = ev.Error + } + } + return []ai.Event{ev} + case "error": + return []ai.Event{{Kind: ai.EventError, Error: geminiErrorMessage(event), SessionID: s.sessionID, Model: s.model}} + case "result": + return s.resultEvents(event) } + return nil +} + +func (s *geminiCLIState) resultEvents(event geminiStreamEvent) []ai.Event { + ev := ai.Event{ + Kind: ai.EventResult, + Tool: "Result", + Success: event.Status != "error", + SessionID: s.sessionID, + Model: s.model, + Usage: geminiUsage(event.Stats), + } + if !ev.Success { + ev.Error = geminiErrorMessage(event) + return []ai.Event{ev} + } + structured, err := ai.ValidatedStructuredData(s.schema, s.text.String(), nil) + if err != nil { + ev.Success = false + ev.Error = err.Error() + return []ai.Event{{Kind: ai.EventError, Error: err.Error(), SessionID: s.sessionID, Model: s.model}, ev} + } + ev.StructuredData = structured + return []ai.Event{ev} +} - usage := ai.Usage{} - if cliResp.Usage != nil { - usage.InputTokens = cliResp.Usage.InputTokens - usage.OutputTokens = cliResp.Usage.OutputTokens +func geminiUsage(stats *geminiRunStats) *ai.Usage { + if stats == nil { + return nil + } + return &ai.Usage{ + InputTokens: ai.NetInputTokens(stats.InputTokens, stats.Cached), + OutputTokens: stats.OutputTokens, + CacheReadTokens: stats.Cached, } +} - return &ai.Response{ - Text: cliResp.Text, - Model: g.model, - Backend: ai.BackendGeminiCLI, - Usage: usage, - Duration: time.Since(start), - Raw: cliResp, - }, nil +func geminiErrorMessage(event geminiStreamEvent) string { + if event.Error != nil && event.Error.Message != "" { + return event.Error.Message + } + if event.Message != "" { + return event.Message + } + if event.Output != "" { + return event.Output + } + return fmt.Sprintf("gemini reported a %s failure", firstNonEmpty(event.Type, "run")) } diff --git a/pkg/ai/provider/gemini_cli_test.go b/pkg/ai/provider/gemini_cli_test.go new file mode 100644 index 0000000..5914319 --- /dev/null +++ b/pkg/ai/provider/gemini_cli_test.go @@ -0,0 +1,183 @@ +package provider + +import ( + "strings" + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" +) + +func TestBuildGeminiCLIArgs(t *testing.T) { + args, err := buildGeminiCLIArgs("gemini-3.5-flash", ai.Request{Prompt: api.Prompt{User: "hi"}}) + if err != nil { + t.Fatalf("buildGeminiCLIArgs: %v", err) + } + // stream-json is what makes the run parseable; without it gemini prints prose. + requireFlagValue(t, args, "--output-format", "stream-json") + requireFlagValue(t, args, "--model", "gemini-3.5-flash") + for _, arg := range args { + if arg == "--approval-mode" { + t.Fatalf("default posture must not pin an approval mode, got %v", args) + } + } +} + +func TestGeminiCLIApprovalMode(t *testing.T) { + cases := []struct { + name string + perms api.Permissions + want string + }{ + {"default posture leaves gemini's own default", api.Permissions{}, ""}, + {"explicit default posture", api.Permissions{Mode: api.PermissionDefault}, ""}, + {"plan is gemini's read-only mode", api.Permissions{Mode: api.PermissionPlan}, "plan"}, + {"acceptEdits auto-approves edit tools", api.Permissions{Mode: api.PermissionAcceptEdits}, "auto_edit"}, + {"auto auto-approves edit tools", api.Permissions{Mode: api.PermissionAuto}, "auto_edit"}, + {"bypass is yolo", api.Permissions{Mode: api.PermissionBypass}, "yolo"}, + {"dontAsk is yolo", api.Permissions{Mode: api.PermissionDontAsk}, "yolo"}, + {"edit preset without a mode auto-approves edits", api.Permissions{Presets: []api.Preset{api.PresetEdit}}, "auto_edit"}, + {"edit preset with an explicit mode defers to the mode", api.Permissions{Mode: api.PermissionDefault, Presets: []api.Preset{api.PresetEdit}}, ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := geminiApprovalMode(tc.perms); got != tc.want { + t.Fatalf("geminiApprovalMode = %q, want %q", got, tc.want) + } + args, err := buildGeminiCLIArgs("gemini-3.5-flash", ai.Request{Permissions: tc.perms}) + if err != nil { + t.Fatalf("buildGeminiCLIArgs: %v", err) + } + if tc.want != "" { + requireFlagValue(t, args, "--approval-mode", tc.want) + } + }) + } +} + +func TestBuildGeminiCLIArgsRejectsAttachments(t *testing.T) { + // The registry declares no media types for gemini-cli, so an attachment must + // fail loudly rather than being silently dropped from the prompt. + req := ai.Request{Prompt: api.Prompt{ + User: "describe this", + Attachments: []api.AttachmentRef{{ID: "shot", MediaType: "image/png"}}, + }} + if _, err := buildGeminiCLIArgs("gemini-3.5-flash", req); err == nil { + t.Fatal("attachment on gemini-cli must be rejected") + } +} + +// The wire format below is gemini CLI 0.50's `--output-format stream-json` +// (StreamJsonFormatter): one JSON object per line, assistant text arriving as +// deltas and a single terminal `result` carrying session stats. +func TestGeminiCLIStateMapsStreamJSON(t *testing.T) { + state := &geminiCLIState{model: "gemini-3.5-flash"} + + init := state.mapLine([]byte(`{"type":"init","timestamp":"2026-07-26T00:00:00.000Z","session_id":"sess-1","model":"gemini-3.5-flash"}`)) + if len(init) != 1 || init[0].Kind != ai.EventSystem || init[0].SessionID != "sess-1" { + t.Fatalf("init events = %+v, want one EventSystem carrying the session id", init) + } + + // The user echo is the prompt coming back; emitting it would double-count the + // input as assistant text in the coalesced response. + if echo := state.mapLine([]byte(`{"type":"message","role":"user","content":"hi"}`)); len(echo) != 0 { + t.Fatalf("user echo produced events: %+v", echo) + } + + text := state.mapLine([]byte(`{"type":"message","role":"assistant","content":"Hello ","delta":true}`)) + if len(text) != 1 || text[0].Kind != ai.EventText || text[0].Text != "Hello " { + t.Fatalf("assistant delta events = %+v, want one EventText", text) + } + if text[0].SessionID != "sess-1" { + t.Fatalf("assistant delta lost the session id: %+v", text[0]) + } + + use := state.mapLine([]byte(`{"type":"tool_use","tool_name":"read_file","tool_id":"call-1","parameters":{"absolute_path":"/tmp/x"}}`)) + if len(use) != 1 || use[0].Kind != ai.EventToolUse || use[0].Tool != "read_file" || use[0].ToolCallID != "call-1" { + t.Fatalf("tool_use events = %+v", use) + } + if got, _ := use[0].Input["absolute_path"].(string); got != "/tmp/x" { + t.Fatalf("tool_use input = %+v, want absolute_path=/tmp/x", use[0].Input) + } + + ok := state.mapLine([]byte(`{"type":"tool_result","tool_id":"call-1","status":"success","output":"file body"}`)) + if len(ok) != 1 || ok[0].Kind != ai.EventToolResult || !ok[0].Success || ok[0].Text != "file body" { + t.Fatalf("tool_result events = %+v", ok) + } + + failed := state.mapLine([]byte(`{"type":"tool_result","tool_id":"call-2","status":"error","output":"","error":{"type":"TOOL_EXECUTION_ERROR","message":"no such file"}}`)) + if len(failed) != 1 || failed[0].Success || failed[0].Text != "no such file" { + t.Fatalf("errored tool_result events = %+v", failed) + } + + // A warning-severity error is advisory: it must not be swallowed, but the run + // still ends on the terminal result. + warn := state.mapLine([]byte(`{"type":"error","severity":"warning","message":"rate limited, retrying"}`)) + if len(warn) != 1 || warn[0].Kind != ai.EventError || warn[0].Error != "rate limited, retrying" { + t.Fatalf("warning events = %+v", warn) + } + + done := state.mapLine([]byte(`{"type":"result","status":"success","stats":{"total_tokens":150,"input_tokens":120,"output_tokens":30,"cached":100,"input":20,"duration_ms":900,"tool_calls":1}}`)) + if len(done) != 1 || done[0].Kind != ai.EventResult || !done[0].Success { + t.Fatalf("result events = %+v", done) + } + usage := done[0].Usage + if usage == nil { + t.Fatal("terminal result carried no usage") + } + // gemini's input_tokens is gross (cache included); captain's buckets are + // disjoint, so the cached prefix must be netted out of InputTokens. + if usage.InputTokens != 20 || usage.CacheReadTokens != 100 || usage.OutputTokens != 30 { + t.Fatalf("usage = %+v, want input=20 cacheRead=100 output=30", *usage) + } + if usage.InputTokens+usage.CacheReadTokens != 120 { + t.Fatalf("netted input buckets = %d, want the reported gross 120", usage.InputTokens+usage.CacheReadTokens) + } + if usage.TotalTokens() != 150 { + t.Fatalf("total tokens = %d, want the reported 150", usage.TotalTokens()) + } +} + +func TestGeminiCLIStateMapsFailedResult(t *testing.T) { + state := &geminiCLIState{model: "gemini-3.5-flash"} + events := state.mapLine([]byte(`{"type":"result","status":"error","error":{"type":"IneligibleTierError","message":"account not eligible"},"stats":{"input_tokens":5,"output_tokens":0}}`)) + if len(events) != 1 || events[0].Kind != ai.EventResult { + t.Fatalf("events = %+v, want one EventResult", events) + } + if events[0].Success || events[0].Error != "account not eligible" { + t.Fatalf("failed result = %+v, want Success=false carrying the error message", events[0]) + } +} + +func TestGeminiCLIStateValidatesPromptAppendedSchema(t *testing.T) { + schema := []byte(`{"type":"object","required":["pass"],"properties":{"pass":{"type":"boolean"}}}`) + + state := &geminiCLIState{model: "gemini-3.5-flash", schema: schema} + state.mapLine([]byte(`{"type":"message","role":"assistant","content":"{\"pass\":true}","delta":true}`)) + events := state.mapLine([]byte(`{"type":"result","status":"success","stats":{"input_tokens":1,"output_tokens":1}}`)) + if len(events) != 1 || !events[0].Success { + t.Fatalf("events = %+v, want one successful EventResult", events) + } + if string(events[0].StructuredData) != `{"pass":true}` { + t.Fatalf("structured data = %s", events[0].StructuredData) + } + + // A reply that ignores the schema must fail the run, not return empty output. + bad := &geminiCLIState{model: "gemini-3.5-flash", schema: schema} + bad.mapLine([]byte(`{"type":"message","role":"assistant","content":"sure thing!","delta":true}`)) + events = bad.mapLine([]byte(`{"type":"result","status":"success","stats":{}}`)) + if len(events) != 2 || events[0].Kind != ai.EventError || events[1].Kind != ai.EventResult || events[1].Success { + t.Fatalf("events = %+v, want an EventError then a failed EventResult", events) + } + if !strings.Contains(events[1].Error, "no JSON object") { + t.Fatalf("failed result error = %q, want the schema-extraction failure", events[1].Error) + } +} + +func TestGeminiCLIStateReportsUnparseableLine(t *testing.T) { + state := &geminiCLIState{model: "gemini-3.5-flash"} + events := state.mapLine([]byte(`not json`)) + if len(events) != 1 || events[0].Kind != ai.EventError { + t.Fatalf("events = %+v, want one EventError", events) + } +} diff --git a/pkg/ai/provider/genkit/approval.go b/pkg/ai/provider/genkit/approval.go new file mode 100644 index 0000000..3c3ffdb --- /dev/null +++ b/pkg/ai/provider/genkit/approval.go @@ -0,0 +1,220 @@ +package genkit + +import ( + "encoding/json" + "fmt" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" +) + +func toolApprovalState(req ai.Request, response *gkai.ModelResponse) (*api.ToolApprovalState, error) { + if response == nil || response.FinishReason != gkai.FinishReasonInterrupted || response.Message == nil { + return nil, fmt.Errorf("genkit approval state requires an interrupted model response") + } + messages, err := approvalRequestMessages(req) + if err != nil { + return nil, err + } + assistant, calls, err := interruptedAssistantMessage(response.Message) + if err != nil { + return nil, err + } + state := &api.ToolApprovalState{Messages: append(messages, assistant), Calls: calls} + if err := state.Validate(); err != nil { + return nil, fmt.Errorf("genkit approval state: %w", err) + } + return state, nil +} + +func approvalRequestMessages(req ai.Request) ([]api.Message, error) { + if len(req.Messages) > 0 { + if err := api.ValidateMessages(req.Messages); err != nil { + return nil, fmt.Errorf("canonical messages: %w", err) + } + return append([]api.Message(nil), req.Messages...), nil + } + messages := make([]api.Message, 0, 2) + if req.Prompt.System != "" { + messages = append(messages, api.Message{Role: api.RoleSystem, Parts: []api.Part{{Type: api.PartText, Text: req.Prompt.System}}}) + } + parts := make([]api.Part, 0, len(req.Prompt.Attachments)) + if req.Prompt.User != "" { + parts = append(parts, api.Part{Type: api.PartText, Text: req.Prompt.User}) + } + for i := range req.Prompt.Attachments { + attachment := req.Prompt.Attachments[i] + parts = append(parts, api.Part{Type: api.PartAttachment, Attachment: &attachment}) + } + if len(parts) == 0 { + return nil, fmt.Errorf("genkit approval state has no user prompt") + } + return append(messages, api.Message{Role: api.RoleUser, Parts: parts}), nil +} + +func interruptedAssistantMessage(message *gkai.Message) (api.Message, []api.ToolApprovalCall, error) { + if message.Role != gkai.RoleModel { + return api.Message{}, nil, fmt.Errorf("genkit interrupted message has role %q, expected model", message.Role) + } + assistant := api.Message{Role: api.RoleAssistant, Parts: make([]api.Part, 0, len(message.Content))} + calls := make([]api.ToolApprovalCall, 0, len(message.Content)) + for i, part := range message.Content { + switch { + case part.IsText(): + if part.Text != "" { + assistant.Parts = append(assistant.Parts, api.Part{Type: api.PartText, Text: part.Text}) + } + case part.IsReasoning(): + if part.Text != "" { + assistant.Parts = append(assistant.Parts, api.Part{Type: api.PartReasoning, Text: part.Text}) + } + case part.IsToolRequest() && part.ToolRequest != nil: + call, err := interruptedApprovalCall(part) + if err != nil { + return api.Message{}, nil, fmt.Errorf("interrupted message part %d: %w", i+1, err) + } + calls = append(calls, call) + request := call.Request + assistant.Parts = append(assistant.Parts, api.Part{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: request.ToolCallID, Name: request.Tool, Input: request.Input, + }}) + default: + return api.Message{}, nil, fmt.Errorf("interrupted message part %d has unsupported content", i+1) + } + } + return assistant, calls, nil +} + +func interruptedApprovalCall(part *gkai.Part) (api.ToolApprovalCall, error) { + request := part.ToolRequest + input, err := json.Marshal(request.Input) + if err != nil { + return api.ToolApprovalCall{}, fmt.Errorf("marshal tool %q input: %w", request.Name, err) + } + call := api.ToolApprovalCall{Request: api.ToolApprovalRequest{ToolCallID: request.Ref, Tool: request.Name, Input: input}} + pendingOutput, resolved := part.Metadata["pendingOutput"] + _, interrupted := part.Metadata["interrupt"] + if resolved == interrupted { + return api.ToolApprovalCall{}, fmt.Errorf("tool call %q must be either interrupted or resolved", request.Ref) + } + if resolved { + output, err := json.Marshal(pendingOutput) + if err != nil { + return api.ToolApprovalCall{}, fmt.Errorf("marshal completed tool %q output: %w", request.Name, err) + } + call.Result = &api.ToolResult{ToolCallID: request.Ref, Output: output} + } + return call, nil +} + +func prepareToolApprovalResume(resume *api.ToolApprovalResume) ([]*gkai.Message, []*gkai.Part, []*gkai.Part, error) { + if resume == nil { + return nil, nil, nil, fmt.Errorf("tool approval resume is required") + } + if err := resume.Validate(); err != nil { + return nil, nil, nil, err + } + messages, err := conversationMessages(resume.State.Messages) + if err != nil { + return nil, nil, nil, err + } + requests := genkitApprovalRequests(messages[len(messages)-1]) + decisions := make(map[string]api.ToolApprovalDecision, len(resume.Decisions)) + for _, decision := range resume.Decisions { + decisions[decision.ToolCallID] = decision + } + restarts := make([]*gkai.Part, 0, len(resume.Decisions)) + responses := make([]*gkai.Part, 0, len(resume.Decisions)) + for _, call := range resume.State.Calls { + part := requests[call.Request.ToolCallID] + if call.Result != nil { + output, err := approvalResultOutput(call.Result) + if err != nil { + return nil, nil, nil, err + } + part.Metadata = map[string]any{"pendingOutput": output} + continue + } + decision := decisions[call.Request.ToolCallID] + switch decision.Action { + case api.ToolApprovalApprove: + input, err := approvalRestartInput(call.Request, decision) + if err != nil { + return nil, nil, nil, err + } + restart := gkai.NewToolRequestPart(&gkai.ToolRequest{Name: call.Request.Tool, Ref: call.Request.ToolCallID, Input: input}) + restart.Metadata = map[string]any{"resumed": map[string]any{"captainApproval": true}} + restarts = append(restarts, restart) + case api.ToolApprovalDeny: + reason := decision.Message + if reason == "" { + reason = "tool call denied" + } + responses = append(responses, gkai.NewToolResponsePart(&gkai.ToolResponse{ + Name: call.Request.Tool, Ref: call.Request.ToolCallID, Output: map[string]any{"denied": true, "reason": reason}, + })) + case api.ToolApprovalRespond: + output, err := approvalResultOutput(decision.Result) + if err != nil { + return nil, nil, nil, err + } + responses = append(responses, gkai.NewToolResponsePart(&gkai.ToolResponse{Name: call.Request.Tool, Ref: call.Request.ToolCallID, Output: output})) + } + } + return messages, restarts, responses, nil +} + +func genkitApprovalRequests(message *gkai.Message) map[string]*gkai.Part { + requests := make(map[string]*gkai.Part) + for _, part := range message.Content { + if part.IsToolRequest() && part.ToolRequest != nil { + requests[part.ToolRequest.Ref] = part + } + } + return requests +} + +func approvalRestartInput(request api.ToolApprovalRequest, decision api.ToolApprovalDecision) (any, error) { + if len(decision.Input) > 0 { + return decodePartJSON(decision.Input) + } + return decodePartJSON(request.Input) +} + +func approvalResultOutput(result *api.ToolResult) (any, error) { + if result.Error != "" { + return map[string]any{"error": result.Error}, nil + } + return decodePartJSON(result.Output) +} + +func seedToolApprovalCorrelation(resume *api.ToolApprovalResume, correlation *toolEventCorrelation) error { + if resume == nil || correlation == nil { + return nil + } + decisions := make(map[string]api.ToolApprovalDecision, len(resume.Decisions)) + for _, decision := range resume.Decisions { + decisions[decision.ToolCallID] = decision + } + for _, call := range resume.State.Calls { + input, err := decodePartJSON(call.Request.Input) + if err != nil { + return fmt.Errorf("seed approval call %q: %w", call.Request.ToolCallID, err) + } + if decision, ok := decisions[call.Request.ToolCallID]; ok && len(decision.Input) > 0 { + input, err = decodePartJSON(decision.Input) + if err != nil { + return fmt.Errorf("seed approval decision %q: %w", call.Request.ToolCallID, err) + } + } + request := &gkai.ToolRequest{Name: call.Request.Tool, Ref: call.Request.ToolCallID, Input: input} + if call.Result == nil && decisions[call.Request.ToolCallID].Action == api.ToolApprovalApprove { + correlation.observeRequest(request) + continue + } + correlation.seedResolved(request) + } + return nil +} diff --git a/pkg/ai/provider/genkit/attachments_ginkgo_test.go b/pkg/ai/provider/genkit/attachments_ginkgo_test.go new file mode 100644 index 0000000..4640bc7 --- /dev/null +++ b/pkg/ai/provider/genkit/attachments_ginkgo_test.go @@ -0,0 +1,45 @@ +package genkit + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" +) + +var _ = Describe("promptParts", func() { + It("preserves text followed by ordered prepared media parts", func() { + first := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"}. + WithPreparedContent(api.AttachmentContent{Bytes: []byte("first")}) + second := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "application/pdf"}. + WithPreparedContent(api.AttachmentContent{Bytes: []byte("second")}) + + parts, err := promptParts(ai.Request{Prompt: api.Prompt{User: "compare", Attachments: []api.AttachmentRef{first, second}}}) + Expect(err).NotTo(HaveOccurred()) + Expect(parts).To(HaveLen(3)) + Expect(parts[0].Kind).To(Equal(gkai.PartText)) + Expect(parts[0].Text).To(Equal("compare")) + Expect(parts[1].ContentType).To(Equal("image/png")) + Expect(parts[1].Text).To(Equal("data:image/png;base64,Zmlyc3Q=")) + Expect(parts[2].ContentType).To(Equal("application/pdf")) + Expect(parts[2].Text).To(Equal("data:application/pdf;base64,c2Vjb25k")) + }) + + It("supports a file-only prompt", func() { + attachment := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"}. + WithPreparedContent(api.AttachmentContent{Bytes: []byte("image")}) + parts, err := promptParts(ai.Request{Prompt: api.Prompt{Attachments: []api.AttachmentRef{attachment}}}) + Expect(err).NotTo(HaveOccurred()) + Expect(parts).To(HaveLen(1)) + Expect(parts[0].Kind).To(Equal(gkai.PartMedia)) + }) + + It("fails loudly when resolution was skipped", func() { + attachment := api.AttachmentRef{ID: api.AttachmentIDPrefix + string(make([]byte, 64)), MediaType: "image/png"} + _, err := promptParts(ai.Request{Prompt: api.Prompt{Attachments: []api.AttachmentRef{attachment}}}) + Expect(err).To(MatchError(ContainSubstring("is not prepared"))) + }) +}) diff --git a/pkg/ai/provider/genkit/attachments_suite_test.go b/pkg/ai/provider/genkit/attachments_suite_test.go new file mode 100644 index 0000000..66ba0ba --- /dev/null +++ b/pkg/ai/provider/genkit/attachments_suite_test.go @@ -0,0 +1,13 @@ +package genkit + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAttachments(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Genkit Attachments Suite") +} diff --git a/pkg/ai/provider/genkit/genkit.go b/pkg/ai/provider/genkit/genkit.go index c6301d7..e61bb07 100644 --- a/pkg/ai/provider/genkit/genkit.go +++ b/pkg/ai/provider/genkit/genkit.go @@ -12,10 +12,14 @@ import ( "context" "encoding/json" "fmt" + "strings" + "sync" "time" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/ai/pricing" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" "github.com/flanksource/commons/logger" gkai "github.com/firebase/genkit/go/ai" @@ -28,10 +32,12 @@ var log = logger.GetLogger("ai") // Provider is a genkit-backed ai.StreamingProvider for one API backend. type Provider struct { - cfg ai.Config - backend ai.Backend - g *gk.Genkit - modelRef string + cfg ai.Config + backend ai.Backend + g *gk.Genkit + modelRef string + toolOptionsMu sync.Mutex + toolCorrelation *toolEventCorrelation } var _ ai.StreamingProvider = (*Provider)(nil) @@ -57,12 +63,19 @@ func New(cfg ai.Config) (*Provider, error) { apiKey := cfg.APIKey if apiKey == "" { - apiKey = ai.GetAPIKeyFromEnv(backend) + resolved, err := ai.ResolveAPIKey(backend) + if err != nil { + return nil, err + } + apiKey = resolved.Token } if apiKey == "" { - return nil, fmt.Errorf("genkit provider: no API key for backend %q (set the provider's API key, e.g. ANTHROPIC_API_KEY/OPENAI_API_KEY/GEMINI_API_KEY/DEEPSEEK_API_KEY)", backend) + return nil, fmt.Errorf("%w: genkit provider has no API key for backend %q (set the provider's API key, e.g. ANTHROPIC_API_KEY/OPENAI_API_KEY/GEMINI_API_KEY/DEEPSEEK_API_KEY)", ai.ErrNoAPIKey, backend) } + cfg.Model.Backend = backend + cfg.Model.Name = ai.NormalizeModelForBackend(backend, cfg.Model.Name) + ref, err := modelRef(backend, cfg.Model.Name) if err != nil { return nil, err @@ -73,7 +86,6 @@ func New(cfg ai.Config) (*Provider, error) { return nil, err } - cfg.Model.Backend = backend return &Provider{cfg: cfg, backend: backend, g: g, modelRef: ref}, nil } @@ -84,7 +96,7 @@ func (p *Provider) GetBackend() ai.Backend { return p.backend } func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, error) { start := time.Now() - opts, err := generateOptions(p, req, nil) + opts, err := p.correlatedGenerateOptions(req, nil, nil, nil) if err != nil { return nil, err } @@ -93,11 +105,27 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e if ctx.Err() != nil { return nil, fmt.Errorf("%w: %v", ai.ErrTimeout, ctx.Err()) } + // genkit validates the model's constrained output against the request + // schema during generation; a rejection is recoverable by re-asking the + // model with the errors, so classify it as ErrSchemaValidation (preserving + // genkit's detail lines) for the schema-validation middleware to act on. + if isSchemaMismatch(err) { + return nil, fmt.Errorf("%w: %v", ai.ErrSchemaValidation, err) + } return nil, fmt.Errorf("genkit %s generate: %w", p.backend, err) } out := responseToResponse(resp, p.backend, p.cfg.Model.Name, start) + if resp.FinishReason == gkai.FinishReasonInterrupted { + out.ToolApproval, err = toolApprovalState(req, resp) + if err != nil { + return nil, err + } + } + if out.ToolApproval != nil { + return out, nil + } if req.Prompt.Schema != nil { if err := resp.Output(req.Prompt.Schema); err != nil { return nil, fmt.Errorf("%w: %v", ai.ErrSchemaValidation, err) @@ -113,12 +141,22 @@ func (p *Provider) Execute(ctx context.Context, req ai.Request) (*ai.Response, e } if cost := p.costUSD(out.Usage); cost > 0 { + out.CostUSD = cost log.Debugf("genkit %s cost: $%.6f (model=%s)", p.backend, cost, p.cfg.Model.Name) } return out, nil } +// isSchemaMismatch reports whether a genkit Generate error is the library's +// constrained-output validation failure (vs a transport/model error), so the +// middleware can re-ask the model with the errors instead of failing. +func isSchemaMismatch(err error) bool { + msg := err.Error() + return strings.Contains(msg, "did not match expected schema") || + strings.Contains(msg, "output matching expected schema") +} + // ExecuteStream runs a streaming generation, publishing each chunk as ai.Events // and a terminal EventResult carrying usage + best-effort cost. Structured // output is unsupported in stream mode (mirrors the claude_cli stream provider). @@ -128,8 +166,13 @@ func (p *Provider) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai } ch := make(chan ai.Event, 16) + correlation := newToolEventCorrelation() cb := func(_ context.Context, chunk *gkai.ModelResponseChunk) error { - for _, ev := range chunkToEvents(chunk, p.cfg.Model.Name) { + events, err := chunkToEvents(chunk, p.cfg.Model.Name, correlation) + if err != nil { + return err + } + for _, ev := range events { select { case ch <- ev: case <-ctx.Done(): @@ -138,8 +181,16 @@ func (p *Provider) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai } return nil } + // emit lets in-process caller-tool execution publish tool_use/permission/ + // tool_result events onto the same stream as the model's text chunks. + emit := func(ev ai.Event) { + select { + case ch <- ev: + case <-ctx.Done(): + } + } - opts, err := generateOptions(p, req, cb) + opts, err := p.correlatedGenerateOptions(req, cb, emit, correlation) if err != nil { return nil, err } @@ -150,35 +201,55 @@ func (p *Provider) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai ch <- ai.Event{Kind: ai.EventError, Error: err.Error(), Model: p.cfg.Model.Name} return } - usage := mapUsage(resp.Usage) + usage := mapUsage(resp.Usage, p.backend) + var approval *api.ToolApprovalState + if resp.FinishReason == gkai.FinishReasonInterrupted { + approval, err = toolApprovalState(req, resp) + if err != nil { + ch <- ai.Event{Kind: ai.EventError, Error: err.Error(), Model: p.cfg.Model.Name} + return + } + } ch <- ai.Event{ - Kind: ai.EventResult, - Success: true, - Usage: &usage, - CostUSD: p.costUSD(usage), - Model: p.cfg.Model.Name, + Kind: ai.EventResult, + Success: true, + Usage: &usage, + CostUSD: p.costUSD(usage), + Model: p.cfg.Model.Name, + ToolApproval: approval, } }() return ch, nil } +func (p *Provider) correlatedGenerateOptions( + req ai.Request, + stream gkai.ModelStreamCallback, + emit func(ai.Event), + correlation *toolEventCorrelation, +) ([]gkai.GenerateOption, error) { + p.toolOptionsMu.Lock() + defer p.toolOptionsMu.Unlock() + p.toolCorrelation = correlation + defer func() { p.toolCorrelation = nil }() + if err := seedToolApprovalCorrelation(req.ToolApproval, correlation); err != nil { + return nil, err + } + return generateOptions(p, req, stream, emit) +} + // pricingModelID maps a backend+model onto the OpenRouter-style id the pricing // registry is keyed on (note: Gemini is google/, not googleai/). +// pricingModelID is the OpenRouter pricing key for a backend+model. It uses the +// provider's PricingPrefix — "google" for Gemini, whose catalog namespace is +// "googleai" (see modelRef). The two must not be derived from one another. func pricingModelID(backend ai.Backend, model string) string { - bare := bareModel(model) - switch backend { - case ai.BackendAnthropic: - return "anthropic/" + bare - case ai.BackendOpenAI: - return "openai/" + bare - case ai.BackendGemini: - return "google/" + bare - case ai.BackendDeepSeek: - return "deepseek/" + bare - default: + p, _, ok := registry.ProviderFor(backend) + if !ok { return model } + return p.PricingIDs(model)[0] } // costUSD prices a generation. Genkit usage carries no cost, so it is computed diff --git a/pkg/ai/provider/genkit/genkit_test.go b/pkg/ai/provider/genkit/genkit_test.go index d206589..3f02a25 100644 --- a/pkg/ai/provider/genkit/genkit_test.go +++ b/pkg/ai/provider/genkit/genkit_test.go @@ -1,103 +1,68 @@ package genkit import ( + "fmt" + "path/filepath" "testing" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/credentials" gkai "github.com/firebase/genkit/go/ai" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -func TestEffortConfig(t *testing.T) { - tests := []struct { - name string - backend ai.Backend - req ai.Request - want map[string]any - }{ - { - name: "anthropic no effort defaults max_tokens only", - backend: ai.BackendAnthropic, - req: ai.Request{}, - want: map[string]any{"max_tokens": 4096}, - }, - { - name: "anthropic high effort adds thinking budget on top of base", - backend: ai.BackendAnthropic, - req: ai.Request{Model: api.Model{Effort: api.EffortHigh}}, - want: map[string]any{ - "max_tokens": 24576 + 4096, - "thinking": map[string]any{"type": "enabled", "budget_tokens": 24576}, - }, - }, - { - name: "anthropic medium honours explicit max tokens as base", - backend: ai.BackendAnthropic, - req: ai.Request{Model: api.Model{Effort: api.EffortMedium}, Budget: api.Budget{MaxTokens: 1000}}, - want: map[string]any{ - "max_tokens": 8192 + 1000, - "thinking": map[string]any{"type": "enabled", "budget_tokens": 8192}, - }, - }, - { - name: "openai high effort sets reasoning_effort", - backend: ai.BackendOpenAI, - req: ai.Request{Model: api.Model{Effort: api.EffortHigh}}, - want: map[string]any{"reasoning_effort": "high"}, - }, - { - name: "openai no effort omits config", - backend: ai.BackendOpenAI, - req: ai.Request{}, - want: nil, - }, - { - name: "gemini low effort sets thinkingBudget", - backend: ai.BackendGemini, - req: ai.Request{Model: api.Model{Effort: api.EffortLow}}, - want: map[string]any{"thinkingConfig": map[string]any{"thinkingBudget": 2048}}, - }, - { - name: "gemini no effort omits config", - backend: ai.BackendGemini, - req: ai.Request{}, - want: nil, - }, - { - name: "deepseek omits config (reasoning is selected by model)", - backend: ai.BackendDeepSeek, - req: ai.Request{Model: api.Model{Effort: api.EffortHigh}}, - want: nil, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - assert.Equal(t, tt.want, effortConfig(tt.backend, tt.req)) - }) - } -} - func TestMapUsage(t *testing.T) { - got := mapUsage(&gkai.GenerationUsage{ + // genkit reports overlapping buckets differently per backend; mapUsage must + // normalize to captain's disjoint contract (Input excludes cache, Output + // excludes reasoning). + raw := &gkai.GenerationUsage{ InputTokens: 120, OutputTokens: 45, ThoughtsTokens: 30, CachedContentTokens: 17, TotalTokens: 212, - }) + } + // Anthropic: input_tokens already excludes cache and there is no reasoning + // fold, so both buckets pass through unchanged. assert.Equal(t, ai.Usage{ InputTokens: 120, OutputTokens: 45, - ReasoningTokens: 30, // ThoughtsTokens -> ReasoningTokens - CacheReadTokens: 17, // CachedContentTokens -> CacheReadTokens - }, got) + ReasoningTokens: 30, + CacheReadTokens: 17, + }, mapUsage(raw, ai.BackendAnthropic)) - assert.Equal(t, ai.Usage{}, mapUsage(nil)) + // Gemini: PromptTokenCount folds in cache → net input; CandidatesTokenCount + // excludes thoughts → output passes through. + assert.Equal(t, ai.Usage{ + InputTokens: 103, // 120 - 17 + OutputTokens: 45, + ReasoningTokens: 30, + CacheReadTokens: 17, + }, mapUsage(raw, ai.BackendGemini)) + + // OpenAI/DeepSeek (compat_oai): prompt_tokens folds in cache AND + // completion_tokens folds in reasoning → net both. + openaiWant := ai.Usage{ + InputTokens: 103, // 120 - 17 + OutputTokens: 15, // 45 - 30 + ReasoningTokens: 30, + CacheReadTokens: 17, + } + assert.Equal(t, openaiWant, mapUsage(raw, ai.BackendOpenAI)) + assert.Equal(t, openaiWant, mapUsage(raw, ai.BackendDeepSeek)) + + // Disjoint invariant: for cache-folding backends InputTokens no longer + // overlaps CacheReadTokens, so pricing cannot bill the cached prefix twice. + for _, backend := range []ai.Backend{ai.BackendGemini, ai.BackendOpenAI, ai.BackendDeepSeek} { + got := mapUsage(raw, backend) + assert.Equal(t, raw.InputTokens, got.InputTokens+got.CacheReadTokens, "backend %s input+cache", backend) + } + + assert.Equal(t, ai.Usage{}, mapUsage(nil, ai.BackendAnthropic)) } func TestModelRef(t *testing.T) { @@ -132,6 +97,17 @@ func TestModelRef(t *testing.T) { } } +func TestNewNormalizesBackendModelName(t *testing.T) { + p, err := New(ai.Config{ + Model: api.Model{Backend: ai.BackendAnthropic, Name: "opus-4-8"}, + APIKey: "test-anthropic-key", + }) + require.NoError(t, err) + + assert.Equal(t, "claude-opus-4-8", p.GetModel()) + assert.Equal(t, "anthropic/claude-opus-4-8", p.modelRef) +} + func TestPricingModelID(t *testing.T) { assert.Equal(t, "anthropic/claude-sonnet-4", pricingModelID(ai.BackendAnthropic, "claude-sonnet-4")) assert.Equal(t, "openai/gpt-4o", pricingModelID(ai.BackendOpenAI, "gpt-4o")) @@ -141,6 +117,8 @@ func TestPricingModelID(t *testing.T) { } func TestNewMissingAPIKey(t *testing.T) { + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) // Ensure no provider key leaks in from the environment. for _, env := range []string{"ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY", "GOOGLE_API_KEY", "DEEPSEEK_API_KEY"} { t.Setenv(env, "") @@ -151,6 +129,7 @@ func TestNewMissingAPIKey(t *testing.T) { _, err := New(ai.Config{Model: api.Model{Backend: backend, Name: "some-model"}}) require.Error(t, err) assert.Contains(t, err.Error(), "no API key") + assert.ErrorIs(t, err, ai.ErrNoAPIKey) }) } } @@ -160,3 +139,55 @@ func TestNewUnsupportedBackend(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "does not support backend") } + +func TestIsSchemaMismatch(t *testing.T) { + tests := []struct { + name string + err error + want bool + }{ + { + name: "genkit constrained-output rejection", + err: fmt.Errorf("model failed to generate output matching expected schema: data did not match expected schema:\n- title: String length must be less than or equal to 40"), + want: true, + }, + { + name: "genkit inner detail phrasing", + err: fmt.Errorf("data did not match expected schema:\n- branch: String length must be less than or equal to 40"), + want: true, + }, + { + name: "unrelated transport error", + err: fmt.Errorf("dial tcp: connection refused"), + want: false, + }, + { + name: "rate-limit error is not a schema mismatch", + err: fmt.Errorf("429 rate limit exceeded"), + want: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, isSchemaMismatch(tt.err)) + }) + } +} + +func TestBackendOutputSchema_NormalizesOpenAIStructuredOutput(t *testing.T) { + type answer struct { + Answer string `json:"answer"` + Detail string `json:"detail,omitempty"` + } + req := ai.Request{Prompt: api.Prompt{Schema: &answer{}}} + + schema, handled, err := backendOutputSchema(ai.BackendOpenAI, req) + require.NoError(t, err) + require.True(t, handled) + assert.Equal(t, []any{"answer", "detail"}, schema["required"]) + assert.Equal(t, false, schema["additionalProperties"]) + + _, handled, err = backendOutputSchema(ai.BackendGemini, req) + require.NoError(t, err) + assert.False(t, handled, "Gemini should retain Genkit's WithOutputType path") +} diff --git a/pkg/ai/provider/genkit/instance.go b/pkg/ai/provider/genkit/instance.go index ec5a0bf..204f343 100644 --- a/pkg/ai/provider/genkit/instance.go +++ b/pkg/ai/provider/genkit/instance.go @@ -40,7 +40,7 @@ var instances sync.Map // instanceKey -> *instanceEntry // it on first use. A missing API key is a loud error. func getInstance(ctx context.Context, backend ai.Backend, apiKey string) (*gk.Genkit, error) { if apiKey == "" { - return nil, fmt.Errorf("genkit provider: missing API key for backend %q", backend) + return nil, fmt.Errorf("%w: genkit provider is missing an API key for backend %q", ai.ErrNoAPIKey, backend) } v, _ := instances.LoadOrStore(instanceKey{backend: backend, apiKey: apiKey}, &instanceEntry{}) diff --git a/pkg/ai/provider/genkit/mapping.go b/pkg/ai/provider/genkit/mapping.go index deda376..a7f017b 100644 --- a/pkg/ai/provider/genkit/mapping.go +++ b/pkg/ai/provider/genkit/mapping.go @@ -2,6 +2,9 @@ package genkit import ( "encoding/json" + "fmt" + "reflect" + "sync" "time" "github.com/flanksource/captain/pkg/ai" @@ -9,12 +12,26 @@ import ( gkai "github.com/firebase/genkit/go/ai" ) -// chunkToEvents translates a streamed genkit chunk into captain stream events: -// text parts -> EventText, reasoning parts -> EventThinking, and completed -// (non-partial) tool requests -> EventToolUse. -func chunkToEvents(chunk *gkai.ModelResponseChunk, model string) []ai.Event { +type toolEventCorrelation struct { + mu sync.Mutex + pending []*gkai.ToolRequest + started map[string]string + finished map[string]bool +} + +func newToolEventCorrelation() *toolEventCorrelation { + return &toolEventCorrelation{ + started: make(map[string]string), + finished: make(map[string]bool), + } +} + +// chunkToEvents translates text and reasoning while retaining provider tool +// requests for the in-process handler. Tool lifecycle events are emitted only +// by runTool, after the handler can correlate the call to its provider ref. +func chunkToEvents(chunk *gkai.ModelResponseChunk, model string, correlation *toolEventCorrelation) ([]ai.Event, error) { if chunk == nil { - return nil + return nil, nil } var events []ai.Event for _, p := range chunk.Content { @@ -23,17 +40,139 @@ func chunkToEvents(chunk *gkai.ModelResponseChunk, model string) []ai.Event { events = append(events, ai.Event{Kind: ai.EventText, Text: p.Text, Model: model}) case p.IsReasoning() && p.Text != "": events = append(events, ai.Event{Kind: ai.EventThinking, Text: p.Text, Model: model}) - case p.IsToolRequest() && p.ToolRequest != nil && !p.ToolRequest.Partial: - tr := p.ToolRequest - events = append(events, ai.Event{ - Kind: ai.EventToolUse, - Tool: tr.Name, - Input: toInputMap(tr.Input), - Model: model, - }) + case p.IsToolRequest() && p.ToolRequest != nil: + if correlation != nil { + correlation.observeRequest(p.ToolRequest) + } + case p.IsToolResponse() && p.ToolResponse != nil && !p.IsPartial(): + if correlation == nil { + return nil, fmt.Errorf("genkit tool response %q has no correlation state", p.ToolResponse.Ref) + } + if err := correlation.observeResponse(p.ToolResponse); err != nil { + return nil, err + } } } - return events + return events, nil +} + +func (c *toolEventCorrelation) observeRequest(request *gkai.ToolRequest) { + if request == nil || request.Partial || request.Name == "" { + return + } + c.mu.Lock() + defer c.mu.Unlock() + for _, pending := range c.pending { + if pending == request || (request.Ref != "" && pending.Ref == request.Ref) { + return + } + } + if request.Ref != "" { + if _, ok := c.started[request.Ref]; ok { + return + } + } + c.pending = append(c.pending, request) +} + +func (c *toolEventCorrelation) begin(name string, input map[string]any) (string, error) { + c.mu.Lock() + defer c.mu.Unlock() + + matchingName := make([]int, 0, 1) + exact := make([]int, 0, 1) + for i, request := range c.pending { + if request.Name != name { + continue + } + matchingName = append(matchingName, i) + if requestInput, ok := toolRequestInput(request.Input); ok && reflect.DeepEqual(requestInput, input) { + exact = append(exact, i) + } + } + + index, err := correlatedRequestIndex(name, matchingName, exact) + if err != nil { + return "", err + } + request := c.pending[index] + if request.Ref == "" { + return "", fmt.Errorf("genkit tool %q provider request has no call reference", name) + } + if _, exists := c.started[request.Ref]; exists { + return "", fmt.Errorf("genkit tool request %q started more than once", request.Ref) + } + c.pending = append(c.pending[:index], c.pending[index+1:]...) + c.started[request.Ref] = name + return request.Ref, nil +} + +func correlatedRequestIndex(name string, matchingName, exact []int) (int, error) { + switch { + case len(exact) == 1: + return exact[0], nil + case len(exact) > 1: + return 0, fmt.Errorf("genkit tool %q has multiple provider requests with the same input", name) + case len(matchingName) == 1: + return matchingName[0], nil + case len(matchingName) > 1: + return 0, fmt.Errorf("genkit tool %q has multiple provider requests that cannot be correlated by input", name) + default: + return 0, fmt.Errorf("genkit tool %q execution has no correlated provider request", name) + } +} + +func (c *toolEventCorrelation) finish(ref string) error { + c.mu.Lock() + defer c.mu.Unlock() + if _, ok := c.started[ref]; !ok { + return fmt.Errorf("genkit tool result %q has no correlated provider request", ref) + } + if c.finished[ref] { + return fmt.Errorf("genkit tool result %q emitted more than once", ref) + } + c.finished[ref] = true + return nil +} + +func (c *toolEventCorrelation) seedResolved(request *gkai.ToolRequest) { + c.mu.Lock() + defer c.mu.Unlock() + c.started[request.Ref] = request.Name + c.finished[request.Ref] = true +} + +func (c *toolEventCorrelation) observeResponse(response *gkai.ToolResponse) error { + c.mu.Lock() + defer c.mu.Unlock() + if response.Ref == "" { + return fmt.Errorf("genkit tool response has no provider call reference") + } + name, ok := c.started[response.Ref] + if !ok { + return fmt.Errorf("genkit received uncorrelated tool response %q", response.Ref) + } + if name != response.Name { + return fmt.Errorf("genkit tool response %q names %q, expected %q", response.Ref, response.Name, name) + } + if !c.finished[response.Ref] { + return fmt.Errorf("genkit tool response %q arrived before its result", response.Ref) + } + delete(c.started, response.Ref) + delete(c.finished, response.Ref) + return nil +} + +func toolRequestInput(input any) (map[string]any, bool) { + if text, ok := input.(string); ok { + var decoded map[string]any + if err := json.Unmarshal([]byte(text), &decoded); err != nil { + return nil, false + } + return decoded, true + } + decoded := toInputMap(input) + return decoded, decoded != nil } // toInputMap normalizes a genkit tool-request input (typed any) into the @@ -56,16 +195,30 @@ func toInputMap(v any) map[string]any { return m } -// mapUsage maps genkit's GenerationUsage onto captain's Usage. Genkit reports -// reasoning via ThoughtsTokens and cache reads via CachedContentTokens; it does -// not expose cache writes. -func mapUsage(u *gkai.GenerationUsage) ai.Usage { +// mapUsage maps genkit's GenerationUsage onto captain's disjoint-bucket Usage. +// genkit folds cache reads into InputTokens for Gemini and the OpenAI-compatible +// backends (OpenAI/DeepSeek), and folds reasoning into OutputTokens for the +// OpenAI-compatible backends; Anthropic reports InputTokens already net of cache, +// and both Anthropic and Gemini report OutputTokens without reasoning. Normalize +// to the disjoint contract so the pricing registry and TotalTokens do not +// double-count. genkit's GenerationUsage has no cache-write field, so +// CacheWriteTokens is always zero here — cache-write spend on the API path is +// invisible upstream (finding C4). +func mapUsage(u *gkai.GenerationUsage, backend ai.Backend) ai.Usage { if u == nil { return ai.Usage{} } + input := u.InputTokens + if backend != ai.BackendAnthropic { + input = ai.NetInputTokens(u.InputTokens, u.CachedContentTokens) + } + output := u.OutputTokens + if backend == ai.BackendOpenAI || backend == ai.BackendDeepSeek { + output = ai.NetOutputTokens(u.OutputTokens, u.ThoughtsTokens) + } return ai.Usage{ - InputTokens: u.InputTokens, - OutputTokens: u.OutputTokens, + InputTokens: input, + OutputTokens: output, ReasoningTokens: u.ThoughtsTokens, CacheReadTokens: u.CachedContentTokens, } @@ -78,7 +231,7 @@ func responseToResponse(resp *gkai.ModelResponse, backend ai.Backend, model stri Text: resp.Text(), Model: model, Backend: backend, - Usage: mapUsage(resp.Usage), + Usage: mapUsage(resp.Usage, backend), Duration: time.Since(start), Raw: resp, } diff --git a/pkg/ai/provider/genkit/messages_ginkgo_test.go b/pkg/ai/provider/genkit/messages_ginkgo_test.go new file mode 100644 index 0000000..fabab57 --- /dev/null +++ b/pkg/ai/provider/genkit/messages_ginkgo_test.go @@ -0,0 +1,65 @@ +package genkit + +import ( + "encoding/json" + "strings" + + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("canonical message projection", func() { + It("projects multimodal history without flattening typed parts", func() { + attachment := api.AttachmentRef{ID: api.AttachmentIDPrefix + strings.Repeat("a", 64), MediaType: "image/png"}. + WithPreparedContent(api.AttachmentContent{Bytes: []byte("image")}) + messages, err := conversationMessages([]api.Message{ + {Role: api.RoleSystem, Parts: []api.Part{{Type: api.PartText, Text: "Be precise."}}}, + {Role: api.RoleUser, Parts: []api.Part{ + {Type: api.PartText, Text: "Inspect this."}, + {Type: api.PartAttachment, Attachment: &attachment}, + }}, + {Role: api.RoleAssistant, Parts: []api.Part{ + {Type: api.PartReasoning, Text: "I should inspect it."}, + {Type: api.PartText, Text: "Done."}, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(messages).To(HaveLen(3)) + Expect(messages[0].Role).To(Equal(gkai.RoleSystem)) + Expect(messages[1].Role).To(Equal(gkai.RoleUser)) + Expect(messages[1].Content).To(HaveLen(2)) + Expect(messages[1].Content[1]).To(SatisfyAll( + HaveField("Kind", gkai.PartMedia), + HaveField("ContentType", "image/png"), + HaveField("Text", "data:image/png;base64,aW1hZ2U="), + )) + Expect(messages[2].Role).To(Equal(gkai.RoleModel)) + Expect(messages[2].Content[0].Kind).To(Equal(gkai.PartReasoning)) + }) + + It("correlates canonical tool requests and results", func() { + messages, err := conversationMessages([]api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "List stacks."}}}, + {Role: api.RoleAssistant, Parts: []api.Part{{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: "call-1", Name: "stack_list", Input: json.RawMessage(`{"limit":2}`), + }}}}, + {Role: api.RoleTool, Parts: []api.Part{{Type: api.PartToolResult, ToolResult: &api.ToolResult{ + ToolCallID: "call-1", Output: json.RawMessage(`{"items":["a","b"]}`), + }}}}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(messages).To(HaveLen(3)) + request := messages[1].Content[0].ToolRequest + Expect(request).To(Equal(&gkai.ToolRequest{Ref: "call-1", Name: "stack_list", Input: map[string]any{"limit": float64(2)}})) + result := messages[2].Content[0].ToolResponse + Expect(result).To(Equal(&gkai.ToolResponse{Ref: "call-1", Name: "stack_list", Output: map[string]any{"items": []any{"a", "b"}}})) + }) + + It("does not fall back to the single-turn prompt for an invalid conversation", func() { + _, err := conversationMessages([]api.Message{{Role: api.RoleUser, Parts: []api.Part{{Type: api.PartReasoning, Text: "invalid"}}}}) + Expect(err).To(MatchError(ContainSubstring("user message cannot contain reasoning"))) + }) +}) diff --git a/pkg/ai/provider/genkit/options.go b/pkg/ai/provider/genkit/options.go index ab2828a..26564af 100644 --- a/pkg/ai/provider/genkit/options.go +++ b/pkg/ai/provider/genkit/options.go @@ -1,147 +1,125 @@ package genkit import ( + "encoding/base64" "encoding/json" "fmt" - "strings" + "os" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" gkai "github.com/firebase/genkit/go/ai" ) -// defaultMaxOutputTokens is the visible-answer budget Anthropic requires on every -// request; it is added on top of any extended-thinking budget. -const defaultMaxOutputTokens = 4096 - // bareModel strips a leading provider prefix so a model id can be re-prefixed for -// either a genkit ref (anthropic/openai/googleai) or an OpenRouter pricing key. +// either a genkit ref or an OpenRouter pricing key. func bareModel(model string) string { - for _, prefix := range []string{"anthropic/", "openai/", "googleai/", "google/", "deepseek/", "models/"} { - if strings.HasPrefix(model, prefix) { - return strings.TrimPrefix(model, prefix) - } - } - return model + return registry.StripProviderPrefix(model) } // modelRef produces the genkit model reference for a backend+model // (anthropic/, openai/, googleai/). +// +// The namespace is the provider's CatalogPrefix — googleai for Gemini — which is +// deliberately not the same field pricing uses (google). Keeping them as one +// hand-written switch each is how the two drifted apart. func modelRef(backend ai.Backend, model string) (string, error) { if model == "" { return "", fmt.Errorf("genkit provider: model cannot be empty") } - bare := bareModel(model) - switch backend { - case ai.BackendAnthropic: - return "anthropic/" + bare, nil - case ai.BackendOpenAI: - return "openai/" + bare, nil - case ai.BackendGemini: - return "googleai/" + bare, nil - case ai.BackendDeepSeek: - return "deepseek/" + bare, nil - default: + // genkit serves the API modes only; the CLI/agent/cmux backends are driven by + // their own providers. + p, mode, ok := registry.ProviderFor(backend) + if !ok || mode != registry.ModeAPI { return "", fmt.Errorf("genkit provider: unsupported backend %q", backend) } + return p.CatalogPrefix + "/" + bareModel(model), nil } -// thinkingBudget maps a reasoning effort to an extended-thinking token budget -// (shared by Anthropic and Gemini). xhigh is captain's top tier. -func thinkingBudget(e api.Effort) int { - switch e { - case api.EffortLow: - return 2048 - case api.EffortMedium: - return 8192 - case api.EffortHigh: - return 24576 - case api.EffortXHigh: - return 32768 - default: - return 0 +// generateOptions assembles the genkit Generate options for one turn: model, +// system prompt, user prompt, effort config, and (when streaming) the callback. +// WithOutputType is added only for the non-streaming structured-output path; +// ExecuteStream rejects structured output before calling this. +func generateOptions(p *Provider, req ai.Request, stream gkai.ModelStreamCallback, emit func(ai.Event)) ([]gkai.GenerateOption, error) { + if err := req.ValidateRequestMode(); err != nil { + return nil, err } -} - -// openaiReasoningEffort maps captain's effort onto OpenAI's reasoning_effort -// values; OpenAI tops out at "high", so xhigh clamps to it. -func openaiReasoningEffort(e api.Effort) string { - if e == api.EffortXHigh { - return string(api.EffortHigh) + opts := []gkai.GenerateOption{gkai.WithModelName(p.modelRef)} + toolOptions, err := p.toolOptions(req.ToolPreferences, emit) + if err != nil { + return nil, err + } + opts = append(opts, toolOptions...) + if p.cfg.Model.Name != "" { + req.Name = p.cfg.Model.Name + } + if p.cfg.Model.ID != "" && req.ID == "" { + req.ID = p.cfg.Model.ID } - return string(e) -} - -// anthropicMaxTokens returns max_tokens for an Anthropic request. The thinking -// budget is counted inside max_tokens, so leave room for the visible answer on -// top of it; honour an explicit budget MaxTokens as the visible-answer base. -func anthropicMaxTokens(req ai.Request, e api.Effort) int { - base := req.Budget.MaxTokens - if base <= 0 { - base = defaultMaxOutputTokens - } - budget := thinkingBudget(e) - if budget == 0 { - return base - } - return budget + base -} -// effortConfig builds the provider-specific generation config translating the -// request's reasoning effort into each backend's native control. Anthropic also -// requires max_tokens on every request, so it always returns a config. -func effortConfig(backend ai.Backend, req ai.Request) map[string]any { - e := req.Effort - switch backend { - case ai.BackendOpenAI: - if e == api.EffortNone { - return nil + if req.ToolApproval != nil { + if err := ai.ValidateAttachmentCompatibility([]api.Model{req.Model}, messageAttachments(req.ToolApproval.State.Messages)); err != nil { + return nil, err } - return map[string]any{"reasoning_effort": openaiReasoningEffort(e)} - case ai.BackendGemini: - if e == api.EffortNone { - return nil + messages, restarts, responses, err := prepareToolApprovalResume(req.ToolApproval) + if err != nil { + return nil, fmt.Errorf("genkit tool approval resume: %w", err) } - return map[string]any{"thinkingConfig": map[string]any{"thinkingBudget": thinkingBudget(e)}} - case ai.BackendDeepSeek: - // DeepSeek selects reasoning by model (deepseek-reasoner vs deepseek-chat), - // not a per-request effort knob, so there is no effort config to send. - return nil - case ai.BackendAnthropic: - cfg := map[string]any{"max_tokens": anthropicMaxTokens(req, e)} - if e != api.EffortNone { - cfg["thinking"] = map[string]any{ - "type": "enabled", - "budget_tokens": thinkingBudget(e), + opts = append(opts, gkai.WithMessages(messages...)) + if len(restarts) > 0 { + opts = append(opts, gkai.WithToolRestarts(restarts...)) + } + if len(responses) > 0 { + opts = append(opts, gkai.WithToolResponses(responses...)) + } + } else if len(req.Messages) > 0 { + if err := req.ValidateRequestMode(); err != nil { + return nil, err + } + if err := ai.ValidateAttachmentCompatibility([]api.Model{req.Model}, messageAttachments(req.Messages)); err != nil { + return nil, err + } + messages, err := conversationMessages(req.Messages) + if err != nil { + return nil, err + } + opts = append(opts, gkai.WithMessages(messages...)) + } else { + if req.Prompt.System != "" { + opts = append(opts, gkai.WithSystem(req.Prompt.System)) + } + if len(req.Prompt.Attachments) == 0 { + opts = append(opts, gkai.WithPrompt(req.Prompt.User)) + } else { + if err := ai.ValidateAttachmentCompatibility([]api.Model{req.Model}, req.Prompt.Attachments); err != nil { + return nil, err + } + parts, err := promptParts(req) + if err != nil { + return nil, err } + opts = append(opts, gkai.WithMessages(gkai.NewUserMessage(parts...))) } - return cfg - default: - return nil } -} -// generateOptions assembles the genkit Generate options for one turn: model, -// system prompt, user prompt, effort config, and (when streaming) the callback. -// WithOutputType is added only for the non-streaming structured-output path; -// ExecuteStream rejects structured output before calling this. -func generateOptions(p *Provider, req ai.Request, stream gkai.ModelStreamCallback) ([]gkai.GenerateOption, error) { - opts := []gkai.GenerateOption{gkai.WithModelName(p.modelRef)} - - if req.Prompt.System != "" { - opts = append(opts, gkai.WithSystem(req.Prompt.System)) + modelToken := req.Name + if modelToken == "" { + modelToken = req.ID } - opts = append(opts, gkai.WithPrompt(req.Prompt.User)) - - if cfg := effortConfig(p.backend, req); cfg != nil { + if cfg := ai.EffortConfig(p.backend, modelToken, req.Effort, req.Budget.MaxTokens, req.Temperature); cfg != nil { opts = append(opts, gkai.WithConfig(cfg)) } if stream != nil { opts = append(opts, gkai.WithStreaming(stream)) } if stream == nil { - if len(req.Prompt.SchemaJSON) > 0 { + if schema, handled, err := backendOutputSchema(p.backend, req); err != nil { + return nil, err + } else if handled { + opts = append(opts, gkai.WithOutputSchema(schema)) + } else if len(req.Prompt.SchemaJSON) > 0 { var schema map[string]any if err := json.Unmarshal(req.Prompt.SchemaJSON, &schema); err != nil { return nil, fmt.Errorf("genkit %s: invalid Prompt.SchemaJSON: %w", p.backend, err) @@ -153,3 +131,135 @@ func generateOptions(p *Provider, req ai.Request, stream gkai.ModelStreamCallbac } return opts, nil } + +func promptParts(req ai.Request) ([]*gkai.Part, error) { + parts := make([]*gkai.Part, 0, len(req.Prompt.Attachments)) + if req.Prompt.User != "" { + parts = append(parts, gkai.NewTextPart(req.Prompt.User)) + } + for i, attachment := range req.Prompt.Attachments { + part, err := preparedAttachmentPart(attachment, fmt.Sprintf("attachment %d", i+1)) + if err != nil { + return nil, err + } + parts = append(parts, part) + } + return parts, nil +} + +func conversationMessages(messages []api.Message) ([]*gkai.Message, error) { + if err := api.ValidateMessages(messages); err != nil { + return nil, fmt.Errorf("canonical messages: %w", err) + } + toolNames := map[string]string{} + out := make([]*gkai.Message, 0, len(messages)) + for i, message := range messages { + parts := make([]*gkai.Part, 0, len(message.Parts)) + for j, part := range message.Parts { + switch part.Type { + case api.PartText: + parts = append(parts, gkai.NewTextPart(part.Text)) + case api.PartReasoning: + parts = append(parts, gkai.NewReasoningPart(part.Text, nil)) + case api.PartAttachment: + media, err := preparedAttachmentPart(*part.Attachment, fmt.Sprintf("message %d part %d attachment", i+1, j+1)) + if err != nil { + return nil, err + } + parts = append(parts, media) + case api.PartToolRequest: + input, err := decodePartJSON(part.ToolRequest.Input) + if err != nil { + return nil, fmt.Errorf("message %d part %d tool request input: %w", i+1, j+1, err) + } + toolNames[part.ToolRequest.ToolCallID] = part.ToolRequest.Name + parts = append(parts, gkai.NewToolRequestPart(&gkai.ToolRequest{Ref: part.ToolRequest.ToolCallID, Name: part.ToolRequest.Name, Input: input})) + case api.PartToolResult: + output, err := decodePartJSON(part.ToolResult.Output) + if err != nil { + return nil, fmt.Errorf("message %d part %d tool result output: %w", i+1, j+1, err) + } + if part.ToolResult.Error != "" { + output = map[string]any{"error": part.ToolResult.Error} + } + parts = append(parts, gkai.NewToolResponsePart(&gkai.ToolResponse{Ref: part.ToolResult.ToolCallID, Name: toolNames[part.ToolResult.ToolCallID], Output: output})) + } + } + out = append(out, gkai.NewMessage(genkitRole(message.Role), nil, parts...)) + } + return out, nil +} + +func genkitRole(role api.MessageRole) gkai.Role { + switch role { + case api.RoleSystem: + return gkai.RoleSystem + case api.RoleUser: + return gkai.RoleUser + case api.RoleAssistant: + return gkai.RoleModel + case api.RoleTool: + return gkai.RoleTool + default: + return "" + } +} + +func decodePartJSON(raw json.RawMessage) (any, error) { + if len(raw) == 0 { + return nil, nil + } + var value any + if err := json.Unmarshal(raw, &value); err != nil { + return nil, err + } + return value, nil +} + +func messageAttachments(messages []api.Message) []api.AttachmentRef { + var attachments []api.AttachmentRef + for _, message := range messages { + for _, part := range message.Parts { + if part.Type == api.PartAttachment && part.Attachment != nil { + attachments = append(attachments, *part.Attachment) + } + } + } + return attachments +} + +func preparedAttachmentPart(attachment api.AttachmentRef, label string) (*gkai.Part, error) { + content, ok := attachment.PreparedContent() + if !ok { + return nil, fmt.Errorf("%s (%s) is not prepared", label, attachment.ID) + } + data := content.Bytes + if data == nil && content.Path != "" { + var err error + data, err = os.ReadFile(content.Path) + if err != nil { + return nil, fmt.Errorf("read prepared attachment %s: %w", attachment.ID, err) + } + } + uri := "data:" + attachment.MediaType + ";base64," + base64.StdEncoding.EncodeToString(data) + return gkai.NewMediaPart(attachment.MediaType, uri), nil +} + +// backendOutputSchema resolves schemas for native providers whose supported +// JSON Schema subset differs from Captain's caller-facing schema. The bool is +// false for backends that should retain Genkit's existing WithOutputType or raw +// SchemaJSON behavior. +func backendOutputSchema(backend ai.Backend, req ai.Request) (map[string]any, bool, error) { + if !req.Prompt.HasSchema() || (!ai.UsesAnthropicSchemaSubset(backend) && !ai.UsesOpenAISchemaSubset(backend)) { + return nil, false, nil + } + raw, err := ai.SchemaJSONForBackend(backend, req.Prompt) + if err != nil { + return nil, true, fmt.Errorf("genkit %s: cannot derive Prompt schema: %w", backend, err) + } + var schema map[string]any + if err := json.Unmarshal(raw, &schema); err != nil { + return nil, true, fmt.Errorf("genkit %s: invalid Prompt schema: %w", backend, err) + } + return schema, true, nil +} diff --git a/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go b/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go new file mode 100644 index 0000000..16f545b --- /dev/null +++ b/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go @@ -0,0 +1,175 @@ +package genkit + +import ( + "context" + "encoding/json" + "sync/atomic" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" + gk "github.com/firebase/genkit/go/genkit" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Genkit resumable tool approval", func() { + It("suspends an ask-mode tool instead of executing it without a broker", func() { + provider := newToolProvider(nil) + correlation := newToolEventCorrelation() + correlation.observeRequest(&gkai.ToolRequest{Name: "invoice_update", Ref: "call-update", Input: map[string]any{"amount": 10}}) + events := make([]ai.Event, 0, 2) + ran := false + + _, err := provider.runTool(context.Background(), api.ToolDefinition{ + Name: "invoice_update", + DefaultPermission: api.ToolModeAsk, + Handler: func(context.Context, map[string]any) (any, error) { + ran = true + return "updated", nil + }, + }, map[string]any{"amount": 10}, func(event ai.Event) { events = append(events, event) }, correlation) + + interrupted, metadata := gkai.IsToolInterruptError(err) + Expect(interrupted).To(BeTrue()) + Expect(metadata).To(Equal(map[string]any{"approvalRequired": true})) + Expect(ran).To(BeFalse()) + Expect(events).To(HaveLen(2)) + Expect(events[0].Kind).To(Equal(ai.EventToolUse)) + Expect(events[1].Kind).To(Equal(ai.EventPermission)) + Expect(events[1].ToolCallID).To(Equal("call-update")) + }) + + It("serializes an interrupted turn with pending and completed sibling calls", func() { + request := api.Spec{Prompt: api.Prompt{User: "Update then inspect."}} + pending := gkai.NewToolRequestPart(&gkai.ToolRequest{Name: "invoice_update", Ref: "call-update", Input: map[string]any{"amount": 10}}) + pending.Metadata = map[string]any{"interrupt": map[string]any{"approvalRequired": true}} + completed := gkai.NewToolRequestPart(&gkai.ToolRequest{Name: "invoice_get", Ref: "call-read", Input: map[string]any{"id": "inv-1"}}) + completed.Metadata = map[string]any{"pendingOutput": map[string]any{"amount": 10}} + response := &gkai.ModelResponse{Message: gkai.NewModelMessage(pending, completed), FinishReason: gkai.FinishReasonInterrupted} + + state, err := toolApprovalState(request, response) + Expect(err).NotTo(HaveOccurred()) + Expect(state.Validate()).To(Succeed()) + Expect(state.Messages).To(HaveLen(2)) + Expect(state.Pending()).To(Equal([]api.ToolApprovalRequest{{ + ToolCallID: "call-update", Tool: "invoice_update", Input: json.RawMessage(`{"amount":10}`), + }})) + Expect(state.Calls[1].Result).NotTo(BeNil()) + Expect(state.Calls[1].Result.Output).To(MatchJSON(`{"amount":10}`)) + }) + + It("restarts an approved call with edited input and never replays completed siblings", func(ctx SpecContext) { + var updateRuns atomic.Int32 + var readRuns atomic.Int32 + var updatedAmount atomic.Int32 + genkit := gk.Init(ctx) + modelRef := "test/resumable-approval" + gk.DefineModel(genkit, modelRef, &gkai.ModelOptions{Supports: &gkai.ModelSupports{Tools: true, Multiturn: true}}, + func(_ context.Context, request *gkai.ModelRequest, stream gkai.ModelStreamCallback) (*gkai.ModelResponse, error) { + last := request.Messages[len(request.Messages)-1] + if last.Role == gkai.RoleTool { + if stream != nil { + Expect(stream(ctx, &gkai.ModelResponseChunk{Role: gkai.RoleModel, Content: []*gkai.Part{gkai.NewTextPart("done")}})).To(Succeed()) + } + return &gkai.ModelResponse{Message: gkai.NewModelTextMessage("done"), FinishReason: gkai.FinishReasonStop}, nil + } + message := gkai.NewModelMessage( + gkai.NewToolRequestPart(&gkai.ToolRequest{Name: "invoice_get", Ref: "call-read", Input: map[string]any{"id": "inv-1"}}), + gkai.NewToolRequestPart(&gkai.ToolRequest{Name: "invoice_update", Ref: "call-update", Input: map[string]any{"amount": 10}}), + ) + if stream != nil { + Expect(stream(ctx, &gkai.ModelResponseChunk{Role: gkai.RoleModel, Content: message.Content})).To(Succeed()) + } + return &gkai.ModelResponse{Message: message, FinishReason: gkai.FinishReasonStop}, nil + }) + + provider := &Provider{ + cfg: ai.Config{ + Model: api.Model{Name: "resumable-approval", Backend: api.BackendOpenAI}, + Tools: []api.ToolDefinition{ + {Name: "invoice_get", DefaultPermission: api.ToolModeOn, Handler: func(context.Context, map[string]any) (any, error) { + readRuns.Add(1) + return map[string]any{"amount": 10}, nil + }}, + {Name: "invoice_update", DefaultPermission: api.ToolModeAsk, Handler: func(_ context.Context, input map[string]any) (any, error) { + updateRuns.Add(1) + updatedAmount.Store(int32(input["amount"].(float64))) + return map[string]any{"updated": true}, nil + }}, + }, + }, + backend: api.BackendOpenAI, + g: genkit, modelRef: modelRef, + } + + firstStream, err := provider.ExecuteStream(ctx, api.Spec{Prompt: api.Prompt{User: "Update then inspect."}}) + Expect(err).NotTo(HaveOccurred()) + var state *api.ToolApprovalState + for event := range firstStream { + Expect(event.Kind).NotTo(Equal(ai.EventError), event.Error) + if event.Kind == ai.EventResult { + state = event.ToolApproval + } + } + Expect(state).NotTo(BeNil()) + Expect(state.Pending()).To(HaveLen(1)) + Expect(readRuns.Load()).To(Equal(int32(1))) + Expect(updateRuns.Load()).To(BeZero()) + + resumedStream, err := provider.ExecuteStream(ctx, api.Spec{ToolApproval: &api.ToolApprovalResume{ + State: *state, + Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-update", Tool: "invoice_update", Action: api.ToolApprovalApprove, + Input: json.RawMessage(`{"amount":12}`), + }}, + }}) + Expect(err).NotTo(HaveOccurred()) + var resumed []ai.Event + sawDone := false + for event := range resumedStream { + resumed = append(resumed, event) + Expect(event.Kind).NotTo(Equal(ai.EventError), event.Error) + if event.Kind == ai.EventText && event.Text == "done" { + sawDone = true + } + } + Expect(sawDone).To(BeTrue()) + Expect(resumed[len(resumed)-1].Kind).To(Equal(ai.EventResult)) + Expect(resumed[len(resumed)-1].ToolApproval).To(BeNil()) + Expect(readRuns.Load()).To(Equal(int32(1))) + Expect(updateRuns.Load()).To(Equal(int32(1))) + Expect(updatedAmount.Load()).To(Equal(int32(12))) + }) + + It("maps deny and externally-resolved calls to native responses", func() { + state := api.ToolApprovalState{ + Messages: []api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Change it."}}}, + {Role: api.RoleAssistant, Parts: []api.Part{ + {Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ToolCallID: "call-deny", Name: "invoice_delete", Input: json.RawMessage(`{"id":"inv-1"}`)}}, + {Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ToolCallID: "call-respond", Name: "invoice_update", Input: json.RawMessage(`{"amount":10}`)}}, + }}, + }, + Calls: []api.ToolApprovalCall{ + {Request: api.ToolApprovalRequest{ToolCallID: "call-deny", Tool: "invoice_delete", Input: json.RawMessage(`{"id":"inv-1"}`)}}, + {Request: api.ToolApprovalRequest{ToolCallID: "call-respond", Tool: "invoice_update", Input: json.RawMessage(`{"amount":10}`)}}, + }, + } + resume := &api.ToolApprovalResume{State: state, Decisions: []api.ToolApprovalDecision{ + {ToolCallID: "call-deny", Tool: "invoice_delete", Action: api.ToolApprovalDeny, Message: "keep it"}, + {ToolCallID: "call-respond", Tool: "invoice_update", Action: api.ToolApprovalRespond, Result: &api.ToolResult{ + ToolCallID: "call-respond", Output: json.RawMessage(`{"updated":true}`), + }}, + }} + + messages, restarts, responses, err := prepareToolApprovalResume(resume) + Expect(err).NotTo(HaveOccurred()) + Expect(restarts).To(BeEmpty()) + Expect(responses).To(HaveLen(2)) + Expect(messages[len(messages)-1].Role).To(Equal(gkai.RoleModel)) + Expect(responses[0].ToolResponse.Output).To(Equal(map[string]any{"denied": true, "reason": "keep it"})) + Expect(responses[1].ToolResponse.Output).To(Equal(map[string]any{"updated": true})) + }) +}) diff --git a/pkg/ai/provider/genkit/tool_lifecycle_ginkgo_test.go b/pkg/ai/provider/genkit/tool_lifecycle_ginkgo_test.go new file mode 100644 index 0000000..7d77327 --- /dev/null +++ b/pkg/ai/provider/genkit/tool_lifecycle_ginkgo_test.go @@ -0,0 +1,130 @@ +package genkit + +import ( + "context" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" +) + +var _ = Describe("Genkit tool event correlation", func() { + var ( + provider *Provider + events []ai.Event + emit func(ai.Event) + tool api.ToolDefinition + ) + + BeforeEach(func() { + provider = newToolProvider(nil) + events = nil + emit = func(event ai.Event) { events = append(events, event) } + tool = api.ToolDefinition{ + Name: "lookup", + DefaultPermission: api.ToolModeOn, + Handler: func(_ context.Context, input map[string]any) (any, error) { + return map[string]any{"city": input["city"]}, nil + }, + } + }) + + It("uses the Anthropic provider reference for one lifecycle", func() { + correlation := newToolEventCorrelation() + request := &gkai.ToolRequest{Name: tool.Name, Ref: "toolu_123", Input: map[string]any{"city": "Cape Town"}} + + mapped, err := chunkToEvents(toolRequestChunk(request), provider.GetModel(), correlation) + Expect(err).NotTo(HaveOccurred()) + Expect(mapped).To(BeEmpty()) + mapped, err = chunkToEvents(toolRequestChunk(request), provider.GetModel(), correlation) + Expect(err).NotTo(HaveOccurred()) + Expect(mapped).To(BeEmpty()) + + _, err = provider.runTool(context.Background(), tool, request.Input, emit, correlation) + Expect(err).NotTo(HaveOccurred()) + Expect(events).To(HaveLen(2)) + Expect(events[0].Kind).To(Equal(ai.EventToolUse)) + Expect(events[0].ToolCallID).To(Equal(request.Ref)) + Expect(events[1].Kind).To(Equal(ai.EventToolResult)) + Expect(events[1].ToolCallID).To(Equal(request.Ref)) + + mapped, err = chunkToEvents(toolResponseChunk(tool.Name, request.Ref), provider.GetModel(), correlation) + Expect(err).NotTo(HaveOccurred()) + Expect(mapped).To(BeEmpty()) + }) + + It("ignores OpenAI argument deltas after retaining the provider reference", func() { + correlation := newToolEventCorrelation() + first := &gkai.ToolRequest{Name: tool.Name, Ref: "call_123", Input: `{"city"`} + delta := &gkai.ToolRequest{Input: `:"Cape Town"}`} + + _, err := chunkToEvents(toolRequestChunk(first), provider.GetModel(), correlation) + Expect(err).NotTo(HaveOccurred()) + _, err = chunkToEvents(toolRequestChunk(delta), provider.GetModel(), correlation) + Expect(err).NotTo(HaveOccurred()) + + _, err = provider.runTool(context.Background(), tool, map[string]any{"city": "Cape Town"}, emit, correlation) + Expect(err).NotTo(HaveOccurred()) + Expect(events).To(HaveLen(2)) + Expect(events[0].ToolCallID).To(Equal(first.Ref)) + Expect(events[1].ToolCallID).To(Equal(first.Ref)) + }) + + It("preserves a Gemini reference assigned after its chunk was observed", func() { + correlation := newToolEventCorrelation() + request := &gkai.ToolRequest{Name: tool.Name, Input: map[string]any{"city": "Cape Town"}} + + _, err := chunkToEvents(toolRequestChunk(request), provider.GetModel(), correlation) + Expect(err).NotTo(HaveOccurred()) + request.Ref = "genkit-assigned-ref" + + _, err = provider.runTool(context.Background(), tool, request.Input, emit, correlation) + Expect(err).NotTo(HaveOccurred()) + Expect(events).To(HaveLen(2)) + Expect(events[0].ToolCallID).To(Equal(request.Ref)) + Expect(events[1].ToolCallID).To(Equal(request.Ref)) + }) + + It("keeps permission events on the provider reference", func() { + provider = newToolProvider(func(_ context.Context, request api.PermissionRequest) (api.PermissionDecision, error) { + Expect(request.ToolUseID).To(Equal("toolu_approved")) + return api.PermissionDecision{Allow: true}, nil + }) + tool.DefaultPermission = api.ToolModeAsk + correlation := newToolEventCorrelation() + request := &gkai.ToolRequest{Name: tool.Name, Ref: "toolu_approved", Input: map[string]any{"city": "Cape Town"}} + _, err := chunkToEvents(toolRequestChunk(request), provider.GetModel(), correlation) + Expect(err).NotTo(HaveOccurred()) + + _, err = provider.runTool(context.Background(), tool, request.Input, emit, correlation) + Expect(err).NotTo(HaveOccurred()) + Expect(events).To(HaveLen(3)) + Expect(events[1].Kind).To(Equal(ai.EventPermission)) + Expect(events[1].ToolCallID).To(Equal(request.Ref)) + }) + + It("fails loudly when a tool response has no correlated request", func() { + correlation := newToolEventCorrelation() + _, err := chunkToEvents(toolResponseChunk(tool.Name, "missing"), provider.GetModel(), correlation) + Expect(err).To(MatchError(ContainSubstring(`uncorrelated tool response "missing"`))) + }) + + It("fails loudly when execution has no provider request", func() { + correlation := newToolEventCorrelation() + _, err := provider.runTool(context.Background(), tool, map[string]any{"city": "Cape Town"}, emit, correlation) + Expect(err).To(MatchError(ContainSubstring(`no correlated provider request`))) + Expect(events).To(BeEmpty()) + }) +}) + +func toolRequestChunk(request *gkai.ToolRequest) *gkai.ModelResponseChunk { + return &gkai.ModelResponseChunk{Role: gkai.RoleModel, Content: []*gkai.Part{gkai.NewToolRequestPart(request)}} +} + +func toolResponseChunk(name, ref string) *gkai.ModelResponseChunk { + return &gkai.ModelResponseChunk{Role: gkai.RoleTool, Content: []*gkai.Part{gkai.NewToolResponsePart(&gkai.ToolResponse{Name: name, Ref: ref, Output: "ok"})}} +} diff --git a/pkg/ai/provider/genkit/tool_preferences_ginkgo_test.go b/pkg/ai/provider/genkit/tool_preferences_ginkgo_test.go new file mode 100644 index 0000000..14f8220 --- /dev/null +++ b/pkg/ai/provider/genkit/tool_preferences_ginkgo_test.go @@ -0,0 +1,116 @@ +package genkit + +import ( + "context" + "fmt" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Genkit tool policy", func() { + noop := func(context.Context, map[string]any) (any, error) { return "ok", nil } + + It("resolves tool preferences before exposing tools", func() { + defs := []api.ToolDefinition{ + {Name: "invoice_list", Group: "billing", DefaultPermission: api.ToolModeAsk, Handler: noop}, + {Name: "invoice_delete", Group: "billing", DefaultPermission: api.ToolModeOn, Handler: noop}, + {Name: "search", DefaultPermission: api.ToolModeOff, Handler: noop}, + } + + selected, err := resolveToolDefinitions(defs, api.ToolPreferences{ + "billing": api.ToolModeOff, + "invoice_list": api.ToolModeOn, + "search": api.ToolModeAsk, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(selected).To(HaveLen(2)) + Expect(selected[0].Name).To(Equal("invoice_list")) + Expect(selected[0].NeedsApproval()).To(BeFalse()) + Expect(selected[1].Name).To(Equal("search")) + Expect(selected[1].NeedsApproval()).To(BeTrue()) + }) + + It("lets tool-level ask override group-level on and invokes Captain approval", func() { + approved := 0 + provider := newToolProvider(func(_ context.Context, request api.PermissionRequest) (api.PermissionDecision, error) { + approved++ + Expect(request.Tool).To(Equal("invoice_delete")) + return api.PermissionDecision{Allow: true}, nil + }) + def := api.ToolDefinition{ + Name: "invoice_delete", Group: "billing", DefaultPermission: api.ToolModeOn, Handler: noop, + } + selected, err := resolveToolDefinitions([]api.ToolDefinition{def}, api.ToolPreferences{ + "billing": api.ToolModeOn, + "invoice_delete": api.ToolModeAsk, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(selected).To(HaveLen(1)) + + _, err = runCorrelatedTool(provider, selected[0], map[string]any{}, func(ai.Event) {}) + Expect(err).NotTo(HaveOccurred()) + Expect(approved).To(Equal(1)) + }) + + It("rejects invalid preferences instead of silently using defaults", func() { + _, err := resolveToolDefinitions([]api.ToolDefinition{{Name: "search", Handler: noop}}, api.ToolPreferences{"search": "sometimes"}) + Expect(err).To(MatchError(ContainSubstring(`invalid tool preference "sometimes" for "search"`))) + }) + + It("rejects an invalid tool default even when a preference disables the tool", func() { + _, err := resolveToolDefinitions([]api.ToolDefinition{ + {Name: "search", DefaultPermission: "sometimes", Handler: noop}, + }, api.ToolPreferences{"search": api.ToolModeOff}) + Expect(err).To(MatchError(ContainSubstring(`tool "search" has invalid default permission "sometimes"`))) + }) + + It("caps Anthropic strict tools deterministically without dropping tools", func() { + defs := make([]api.ToolDefinition, 0, anthropicMaxStrictTools+3) + for i := 0; i < anthropicMaxStrictTools+2; i++ { + defs = append(defs, api.ToolDefinition{ + Name: fmt.Sprintf("readonly_%02d", i), Strict: boolPointer(true), ReadOnlyHint: boolPointer(true), Handler: noop, + }) + } + defs = append(defs, api.ToolDefinition{ + Name: "destructive", Strict: boolPointer(true), DestructiveHint: boolPointer(true), Handler: noop, + }) + + shaped := anthropicStrictToolDefinitions(defs) + Expect(shaped).To(HaveLen(len(defs))) + strictNames := make([]string, 0, anthropicMaxStrictTools) + for _, def := range shaped { + Expect(def.Strict).NotTo(BeNil()) + if *def.Strict { + strictNames = append(strictNames, def.Name) + } + } + Expect(strictNames).To(HaveLen(anthropicMaxStrictTools)) + Expect(strictNames).To(ContainElement("destructive")) + Expect(strictNames).NotTo(ContainElements("readonly_19", "readonly_20", "readonly_21")) + Expect(*defs[0].Strict).To(BeTrue(), "shaping must not mutate Config.Tools") + }) + + It("writes explicit Anthropic strict metadata onto every Genkit tool", func() { + provider := newToolProvider(nil) + strict := api.ToolDefinition{Name: "strict", Strict: boolPointer(true), Handler: noop} + loose := api.ToolDefinition{Name: "loose", Strict: boolPointer(false), Handler: noop} + + Expect(genkitStrictMetadata(provider.genkitTool(strict, nil, nil))).To(BeTrue()) + Expect(genkitStrictMetadata(provider.genkitTool(loose, nil, nil))).To(BeFalse()) + }) +}) + +func genkitStrictMetadata(ref gkai.ToolRef) bool { + tool, ok := ref.(gkai.Tool) + Expect(ok).To(BeTrue()) + strict, ok := tool.Definition().Metadata["strict"].(bool) + Expect(ok).To(BeTrue()) + return strict +} + +func boolPointer(value bool) *bool { return &value } diff --git a/pkg/ai/provider/genkit/tools.go b/pkg/ai/provider/genkit/tools.go new file mode 100644 index 0000000..4eb5f84 --- /dev/null +++ b/pkg/ai/provider/genkit/tools.go @@ -0,0 +1,289 @@ +package genkit + +import ( + "context" + "encoding/json" + "fmt" + "sort" + + "github.com/flanksource/captain/pkg/ai" + captools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" +) + +// maxToolTurns caps the model↔tool iteration so a misbehaving loop terminates. +const ( + maxToolTurns = 16 + anthropicMaxStrictTools = 20 +) + +// SupportsCallerTools reports that the genkit provider can expose and execute +// api.Config.Tools. It satisfies api.ToolCapableProvider. +func (p *Provider) SupportsCallerTools() bool { return true } + +// toolOptions builds the genkit generate options for the caller-supplied tools +// on the provider config. Each tool runs in-process via its handler; emit, when +// non-nil, receives EventToolUse / EventPermission / EventToolResult as the +// model calls them. Returns nil when there are no caller tools, so a run with +// none is byte-for-byte unchanged. +func (p *Provider) toolOptions(preferences api.ToolPreferences, emit func(ai.Event)) ([]gkai.GenerateOption, error) { + definitions, err := resolveToolDefinitions(p.cfg.Tools, preferences) + if err != nil { + return nil, err + } + if len(definitions) == 0 { + return nil, nil + } + if p.backend == ai.BackendAnthropic { + if count := explicitStrictToolCount(definitions); count > anthropicMaxStrictTools { + log.Warnf("Anthropic supports at most %d strict tools; keeping %d of %d explicit opt-ins strict and sending %d as non-strict", anthropicMaxStrictTools, anthropicMaxStrictTools, count, count-anthropicMaxStrictTools) + } + definitions = anthropicStrictToolDefinitions(definitions) + } + refs := make([]gkai.ToolRef, 0, len(definitions)) + for i := range definitions { + refs = append(refs, p.genkitTool(definitions[i], emit, p.toolCorrelation)) + } + return []gkai.GenerateOption{gkai.WithTools(refs...), gkai.WithMaxTurns(maxToolTurns)}, nil +} + +func resolveToolDefinitions(definitions []api.ToolDefinition, preferences api.ToolPreferences) ([]api.ToolDefinition, error) { + if err := preferences.Validate(); err != nil { + return nil, err + } + selected := make([]api.ToolDefinition, 0, len(definitions)) + for _, definition := range definitions { + if definition.Name == "" { + return nil, fmt.Errorf("genkit tool name cannot be empty") + } + if definition.Handler == nil { + return nil, fmt.Errorf("genkit tool %q has no handler", definition.Name) + } + mode, err := effectiveToolMode(definition, preferences) + if err != nil { + return nil, err + } + if mode == api.ToolModeOff { + continue + } + definition.DefaultPermission = mode + selected = append(selected, definition) + } + return selected, nil +} + +func effectiveToolMode(definition api.ToolDefinition, preferences api.ToolPreferences) (api.ToolMode, error) { + defaultMode := api.ToolModeAuto + if definition.DefaultPermission != "" { + var ok bool + defaultMode, ok = api.NormalizeToolMode(definition.DefaultPermission) + if !ok { + return "", fmt.Errorf("genkit tool %q has invalid default permission %q", definition.Name, definition.DefaultPermission) + } + } + info := captools.ToolInfo{Name: definition.Name, Group: definition.Group} + if preferred, ok := captools.EffectivePreference(preferences, info); ok && preferred != api.ToolModeAuto { + return preferred, nil + } + return defaultMode, nil +} + +func anthropicStrictToolDefinitions(definitions []api.ToolDefinition) []api.ToolDefinition { + type candidate struct { + index int + def api.ToolDefinition + } + candidates := make([]candidate, 0, len(definitions)) + for i, definition := range definitions { + if definition.Strict != nil && *definition.Strict { + candidates = append(candidates, candidate{index: i, def: definition}) + } + } + sort.SliceStable(candidates, func(i, j int) bool { + if left, right := strictRiskRank(candidates[i].def), strictRiskRank(candidates[j].def); left != right { + return left < right + } + if candidates[i].def.Name != candidates[j].def.Name { + return candidates[i].def.Name < candidates[j].def.Name + } + return candidates[i].index < candidates[j].index + }) + strict := make(map[int]bool, min(len(candidates), anthropicMaxStrictTools)) + for _, candidate := range candidates[:min(len(candidates), anthropicMaxStrictTools)] { + strict[candidate.index] = true + } + shaped := make([]api.ToolDefinition, len(definitions)) + for i, definition := range definitions { + value := strict[i] + definition.Strict = &value + shaped[i] = definition + } + return shaped +} + +func explicitStrictToolCount(definitions []api.ToolDefinition) int { + count := 0 + for _, definition := range definitions { + if definition.Strict != nil && *definition.Strict { + count++ + } + } + return count +} + +func strictRiskRank(definition api.ToolDefinition) int { + switch { + case definition.DestructiveHint != nil && *definition.DestructiveHint: + return 0 + case definition.IdempotentHint != nil && !*definition.IdempotentHint: + return 1 + case definition.ReadOnlyHint != nil && !*definition.ReadOnlyHint: + return 2 + case definition.ReadOnlyHint == nil: + return 3 + default: + return 4 + } +} + +// genkitTool wraps one api.ToolDefinition as an ephemeral genkit tool. Ephemeral +// (NewTool, passed via WithTools) rather than genkit.DefineTool: +// the genkit instance is cached and shared per (backend, apiKey), so registering +// tools globally would leak them across unrelated requests. +func (p *Provider) genkitTool(def api.ToolDefinition, emit func(ai.Event), correlation *toolEventCorrelation) gkai.ToolRef { + handler := func(tc *gkai.ToolContext, input any) (any, error) { + return p.runTool(tc.Context, def, input, emit, correlation) + } + options := []gkai.ToolOption{gkai.WithInputSchema(def.InputSchema)} + if def.Strict != nil { + options = append(options, gkai.WithStrictSchema(*def.Strict)) + } + return gkai.NewTool(def.Name, def.Description, handler, options...) +} + +// runTool gates a caller tool through Config.CanUseTool (when it needs approval), +// runs its handler, and publishes the use/permission/result events. On a denied +// or errored call it feeds a message back to the model as the tool result so the +// model can react rather than the whole generation failing. +func (p *Provider) runTool( + ctx context.Context, + def api.ToolDefinition, + input any, + emit func(ai.Event), + correlation *toolEventCorrelation, +) (any, error) { + args := toArgsMap(input) + callID := "" + if emit != nil { + if correlation == nil { + return nil, fmt.Errorf("genkit tool %q execution has no correlation state", def.Name) + } + var err error + callID, err = correlation.begin(def.Name, args) + if err != nil { + return nil, err + } + emit(ai.Event{Kind: ai.EventToolUse, Tool: def.Name, Input: args, ToolCallID: callID, Model: p.cfg.Model.Name}) + } + + if def.NeedsApproval() && !gkai.IsToolResumed(ctx) { + if emit != nil { + emit(ai.Event{Kind: ai.EventPermission, Tool: def.Name, Input: args, ToolCallID: callID, Model: p.cfg.Model.Name}) + } + if p.cfg.CanUseTool == nil { + return nil, gkai.NewToolInterruptError(map[string]any{"approvalRequired": true}) + } + decision, err := p.cfg.CanUseTool(ctx, api.PermissionRequest{ + Tool: def.Name, + Input: args, + ToolUseID: callID, + SessionID: p.cfg.SessionID, + }) + if err != nil { + return p.toolDenied(def.Name, callID, err.Error(), emit, correlation) + } + if !decision.Allow { + msg := decision.Message + if msg == "" { + msg = "tool call denied" + } + return p.toolDenied(def.Name, callID, msg, emit, correlation) + } + if decision.UpdatedInput != nil { + args = decision.UpdatedInput + } + } + + out, err := def.Handler(ctx, args) + if err != nil { + if emitErr := p.emitToolResult(ai.Event{Kind: ai.EventToolResult, Tool: def.Name, ToolCallID: callID, Success: false, Text: err.Error(), Model: p.cfg.Model.Name}, emit, correlation); emitErr != nil { + return nil, emitErr + } + // Feed the error back as the tool result rather than aborting the turn. + return map[string]any{"error": err.Error()}, nil + } + + if err := p.emitToolResult(ai.Event{Kind: ai.EventToolResult, Tool: def.Name, ToolCallID: callID, Success: true, Text: toolOutputText(out), Model: p.cfg.Model.Name}, emit, correlation); err != nil { + return nil, err + } + return out, nil +} + +// toolDenied emits a failed EventToolResult and returns the deny reason to the +// model as the tool output (not an error, so the model can adapt). +func (p *Provider) toolDenied(name, callID, reason string, emit func(ai.Event), correlation *toolEventCorrelation) (any, error) { + if err := p.emitToolResult(ai.Event{Kind: ai.EventToolResult, Tool: name, ToolCallID: callID, Success: false, Text: reason, Model: p.cfg.Model.Name}, emit, correlation); err != nil { + return nil, err + } + return map[string]any{"denied": true, "reason": reason}, nil +} + +func (p *Provider) emitToolResult(event ai.Event, emit func(ai.Event), correlation *toolEventCorrelation) error { + if emit == nil { + return nil + } + if err := correlation.finish(event.ToolCallID); err != nil { + return err + } + emit(event) + return nil +} + +// toArgsMap normalizes genkit's decoded tool input into a JSON object map. +func toArgsMap(input any) map[string]any { + if input == nil { + return map[string]any{} + } + if m, ok := input.(map[string]any); ok { + return m + } + data, err := json.Marshal(input) + if err != nil { + return map[string]any{} + } + var m map[string]any + if err := json.Unmarshal(data, &m); err != nil { + return map[string]any{} + } + return m +} + +// toolOutputText renders a handler's return value as the EventToolResult text. +func toolOutputText(out any) string { + if out == nil { + return "" + } + if s, ok := out.(string); ok { + return s + } + data, err := json.Marshal(out) + if err != nil { + return fmt.Sprintf("%v", out) + } + return string(data) +} + +// static assertion: the genkit provider advertises caller-tool support. +var _ api.ToolCapableProvider = (*Provider)(nil) diff --git a/pkg/ai/provider/genkit/tools_test.go b/pkg/ai/provider/genkit/tools_test.go new file mode 100644 index 0000000..87c87cf --- /dev/null +++ b/pkg/ai/provider/genkit/tools_test.go @@ -0,0 +1,160 @@ +package genkit + +import ( + "context" + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" +) + +func newToolProvider(canUse api.PermissionFunc) *Provider { + return &Provider{cfg: ai.Config{Model: api.Model{Name: "test-model"}, CanUseTool: canUse}} +} + +func collectEvents() (func(ai.Event), *[]ai.Event) { + var events []ai.Event + return func(ev ai.Event) { events = append(events, ev) }, &events +} + +func TestRunToolAutoRunsAndEmitsCorrelatedEvents(t *testing.T) { + p := newToolProvider(nil) + emit, events := collectEvents() + + var gotInput map[string]any + def := api.ToolDefinition{ + Name: "echo", + DefaultPermission: api.ToolModeOn, + Handler: func(_ context.Context, in map[string]any) (any, error) { + gotInput = in + return map[string]any{"ok": true}, nil + }, + } + + out, err := runCorrelatedTool(p, def, map[string]any{"x": 1}, emit) + if err != nil { + t.Fatalf("runTool: %v", err) + } + if gotInput["x"] != 1 { + t.Errorf("handler input = %v, want {x:1}", gotInput) + } + if m, ok := out.(map[string]any); !ok || m["ok"] != true { + t.Errorf("out = %v, want {ok:true}", out) + } + + if len(*events) != 2 { + t.Fatalf("events = %d, want 2 (use, result); got %+v", len(*events), *events) + } + use, result := (*events)[0], (*events)[1] + if use.Kind != ai.EventToolUse || result.Kind != ai.EventToolResult { + t.Fatalf("kinds = %q,%q, want tool_use,tool_result", use.Kind, result.Kind) + } + if use.ToolCallID == "" || use.ToolCallID != result.ToolCallID { + t.Errorf("call ids not correlated: use=%q result=%q", use.ToolCallID, result.ToolCallID) + } + if !result.Success { + t.Error("result.Success = false, want true") + } +} + +func TestRunToolApprovalDeniedSkipsHandler(t *testing.T) { + ran := false + deny := func(_ context.Context, _ api.PermissionRequest) (api.PermissionDecision, error) { + return api.PermissionDecision{Allow: false, Message: "nope"}, nil + } + p := newToolProvider(deny) + emit, events := collectEvents() + + def := api.ToolDefinition{ + Name: "danger", + DefaultPermission: api.ToolModeAsk, + Handler: func(context.Context, map[string]any) (any, error) { ran = true; return "ran", nil }, + } + + out, err := runCorrelatedTool(p, def, map[string]any{}, emit) + if err != nil { + t.Fatalf("runTool: %v", err) + } + if ran { + t.Error("handler ran despite denial") + } + if m, ok := out.(map[string]any); !ok || m["denied"] != true { + t.Errorf("out = %v, want denied", out) + } + kinds := eventKinds(*events) + if len(kinds) != 3 || kinds[0] != ai.EventToolUse || kinds[1] != ai.EventPermission || kinds[2] != ai.EventToolResult { + t.Fatalf("event kinds = %v, want [tool_use permission tool_result]", kinds) + } + if (*events)[2].Success { + t.Error("denied result Success = true, want false") + } +} + +func TestRunToolApprovalAllowsAndSubstitutesInput(t *testing.T) { + allow := func(_ context.Context, _ api.PermissionRequest) (api.PermissionDecision, error) { + return api.PermissionDecision{Allow: true, UpdatedInput: map[string]any{"amount": 5}}, nil + } + p := newToolProvider(allow) + emit, _ := collectEvents() + + var seen map[string]any + def := api.ToolDefinition{ + Name: "pay", + DefaultPermission: api.ToolModeAsk, + Handler: func(_ context.Context, in map[string]any) (any, error) { seen = in; return "done", nil }, + } + + if _, err := runCorrelatedTool(p, def, map[string]any{"amount": 1}, emit); err != nil { + t.Fatalf("runTool: %v", err) + } + if seen["amount"] != 5 { + t.Errorf("handler input = %v, want the UpdatedInput amount=5", seen) + } +} + +func TestRunToolHandlerErrorFedBack(t *testing.T) { + p := newToolProvider(nil) + emit, events := collectEvents() + def := api.ToolDefinition{ + Name: "boom", + DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return nil, context.Canceled }, + } + out, err := runCorrelatedTool(p, def, map[string]any{}, emit) + if err != nil { + t.Fatalf("runTool should not surface handler error, got %v", err) + } + if m, ok := out.(map[string]any); !ok || m["error"] == nil { + t.Errorf("out = %v, want {error:...}", out) + } + last := (*events)[len(*events)-1] + if last.Kind != ai.EventToolResult || last.Success { + t.Errorf("last event = %+v, want failed tool_result", last) + } +} + +func TestToolOptionsEmptyWhenNoTools(t *testing.T) { + p := newToolProvider(nil) + if opts, err := p.toolOptions(nil, nil); err != nil || opts != nil { + t.Errorf("toolOptions with no tools = %v, want nil", opts) + } + if !p.SupportsCallerTools() { + t.Error("SupportsCallerTools = false, want true") + } +} + +func eventKinds(events []ai.Event) []ai.EventKind { + kinds := make([]ai.EventKind, len(events)) + for i, e := range events { + kinds[i] = e.Kind + } + return kinds +} + +func runCorrelatedTool(p *Provider, def api.ToolDefinition, input any, emit func(ai.Event)) (any, error) { + correlation := newToolEventCorrelation() + correlation.observeRequest(&gkai.ToolRequest{Name: def.Name, Ref: "provider-call-1", Input: input}) + return p.runTool(context.Background(), def, input, emit, correlation) +} diff --git a/pkg/ai/provider/init.go b/pkg/ai/provider/init.go index e2583dd..1214d33 100644 --- a/pkg/ai/provider/init.go +++ b/pkg/ai/provider/init.go @@ -16,13 +16,11 @@ func init() { // with a custom base URL (see pluginFor). ai.RegisterProvider(ai.BackendDeepSeek, func(cfg ai.Config) (ai.Provider, error) { return genkit.New(cfg) }) - // Claude Agent SDK as a supervised TS process over JSON-RPC replaces the - // one-shot `claude -p` path; claude-code-* models route here too. ai.RegisterProvider(ai.BackendClaudeAgent, func(cfg ai.Config) (ai.Provider, error) { return claudeagent.New(cfg) }) - ai.RegisterProvider(ai.BackendClaudeCLI, func(cfg ai.Config) (ai.Provider, error) { return claudeagent.New(cfg) }) + ai.RegisterProvider(ai.BackendClaudeCLI, func(cfg ai.Config) (ai.Provider, error) { return NewClaudeCLI(cfg.Model.Name), nil }) - // Codex via `codex app-server` (JSON-RPC) replaces `codex exec --json`. - ai.RegisterProvider(ai.BackendCodexCLI, func(cfg ai.Config) (ai.Provider, error) { return NewCodexAppServer(cfg.Model.Name) }) + ai.RegisterProvider(ai.BackendCodexCLI, func(cfg ai.Config) (ai.Provider, error) { return NewCodexCLI(cfg.Model.Name), nil }) + ai.RegisterProvider(ai.BackendCodexAgent, func(cfg ai.Config) (ai.Provider, error) { return NewCodexAppServer(cfg.Model.Name) }) // cmux drives an interactive claude/codex TUI inside a tmux/cmux surface, // tailing the session JSONL; the same provider serves both agents (it reads @@ -30,6 +28,5 @@ func init() { ai.RegisterProvider(ai.BackendClaudeCmux, func(cfg ai.Config) (ai.Provider, error) { return cmux.New(cfg) }) ai.RegisterProvider(ai.BackendCodexCmux, func(cfg ai.Config) (ai.Provider, error) { return cmux.New(cfg) }) - // Gemini CLI is unchanged. ai.RegisterProvider(ai.BackendGeminiCLI, func(cfg ai.Config) (ai.Provider, error) { return NewGeminiCLI(cfg.Model.Name), nil }) } diff --git a/pkg/ai/provider/monitor_hooks.go b/pkg/ai/provider/monitor_hooks.go new file mode 100644 index 0000000..f70fcde --- /dev/null +++ b/pkg/ai/provider/monitor_hooks.go @@ -0,0 +1,95 @@ +package provider + +import ( + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + + "github.com/flanksource/captain/pkg/claude" +) + +// monitorHookTimeout is the settings.json timeout for the injected hooks; the +// receiver itself delivers-or-drops within 1s, this is only the safety bound. +const monitorHookTimeout = 10 + +// monitorHookEvents are the Claude Code lifecycle events captain subscribes to +// for session monitoring (matching `captain hook monitor install`). +var monitorHookEvents = []claude.HookEventType{ + claude.HookEventSessionStart, + claude.HookEventUserPromptSubmit, + claude.HookEventStop, + claude.HookEventSubagentStop, + claude.HookEventSessionEnd, +} + +// captainBinary resolves the binary the injected hooks invoke. This package is +// also consumed as a library from other binaries (gavel, tests), where +// os.Executable is not captain; fall back to PATH, and report false when no +// captain exists — injection is then skipped and the user-level hook install +// plus the monitor's recon cover those sessions. +func captainBinary() (string, bool) { + path, err := os.Executable() + if err == nil && strings.HasPrefix(filepath.Base(path), "captain") { + return path, true + } + if found, err := exec.LookPath("captain"); err == nil { + return found, true + } + return "", false +} + +// claudeMonitorSettings is the --settings document injected into +// captain-launched claude CLI sessions so they report lifecycle events even +// without the user-level hook install. +func claudeMonitorSettings(binary string) ([]byte, error) { + command := fmt.Sprintf("%s hook monitor notify --provider claude", binary) + hooks := claude.HooksConfig{Hooks: map[claude.HookEventType][]claude.HookMatcher{}} + for _, event := range monitorHookEvents { + hooks.Hooks[event] = []claude.HookMatcher{{ + Hooks: []claude.Hook{{Type: claude.HookTypeCommand, Command: command, Timeout: monitorHookTimeout}}, + }} + } + data, err := json.Marshal(hooks) + if err != nil { + return nil, fmt.Errorf("encode monitor hook settings: %w", err) + } + return data, nil +} + +// writeClaudeMonitorSettings writes the --settings document to a temp file and +// returns its path with a cleanup, mirroring writeTempSchema. +func writeClaudeMonitorSettings(binary string) (string, func(), error) { + data, err := claudeMonitorSettings(binary) + if err != nil { + return "", nil, err + } + f, err := os.CreateTemp("", "captain-claude-hooks-*.json") + if err != nil { + return "", nil, fmt.Errorf("claude-cli: create monitor hooks settings file: %w", err) + } + path := f.Name() + if _, err := f.Write(data); err != nil { + _ = f.Close() + _ = os.Remove(path) + return "", nil, fmt.Errorf("claude-cli: write monitor hooks settings: %w", err) + } + if err := f.Close(); err != nil { + _ = os.Remove(path) + return "", nil, fmt.Errorf("claude-cli: close monitor hooks settings: %w", err) + } + return path, func() { _ = os.Remove(path) }, nil +} + +// codexNotifyOverride renders the -c override installing captain as codex's +// notify program for one invocation. A JSON string array is valid TOML. +func codexNotifyOverride(binary string) (string, error) { + argv := []string{binary, "hook", "monitor", "notify", "--provider", "codex"} + value, err := json.Marshal(argv) + if err != nil { + return "", fmt.Errorf("encode codex notify override: %w", err) + } + return "notify=" + string(value), nil +} diff --git a/pkg/ai/provider/monitor_hooks_test.go b/pkg/ai/provider/monitor_hooks_test.go new file mode 100644 index 0000000..087635d --- /dev/null +++ b/pkg/ai/provider/monitor_hooks_test.go @@ -0,0 +1,152 @@ +package provider + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" +) + +// fakeCaptainOnPath makes captainBinary() resolve deterministically in tests: +// the test binary is not captain-named, so PATH lookup finds this fake. +func fakeCaptainOnPath(t *testing.T) string { + t.Helper() + dir := t.TempDir() + path := filepath.Join(dir, "captain") + if err := os.WriteFile(path, []byte("#!/bin/sh\n"), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("PATH", dir) + t.Setenv(api.MonitorHooksEnv, "") + return path +} + +func monitorHookEventNames() []string { + names := make([]string, 0, len(monitorHookEvents)) + for _, event := range monitorHookEvents { + names = append(names, string(event)) + } + return names +} + +func TestClaudeCLIMonitorHooksInjection(t *testing.T) { + binary := fakeCaptainOnPath(t) + req := ai.Request{Prompt: api.Prompt{User: "hi"}} + + args, cleanup, err := buildClaudeCLIArgs("claude-sonnet-5", req) + if err != nil { + t.Fatalf("buildClaudeCLIArgs: %v", err) + } + settingsPath := flagValue(t, args, "--settings") + + data, err := os.ReadFile(settingsPath) + if err != nil { + t.Fatalf("read injected settings: %v", err) + } + var config claude.HooksConfig + if err := json.Unmarshal(data, &config); err != nil { + t.Fatalf("parse injected settings: %v", err) + } + for _, event := range monitorHookEventNames() { + matchers := config.Hooks[claude.HookEventType(event)] + if len(matchers) != 1 || len(matchers[0].Hooks) != 1 { + t.Fatalf("event %s: expected exactly one hook, got %+v", event, matchers) + } + command := matchers[0].Hooks[0].Command + if command != binary+" hook monitor notify --provider claude" { + t.Fatalf("event %s: unexpected command %q", event, command) + } + } + + cleanup() + if _, err := os.Stat(settingsPath); !os.IsNotExist(err) { + t.Fatalf("cleanup must remove the temp settings file") + } + + t.Run("bare request suppresses injection", func(t *testing.T) { + bare := req + bare.Memory.Bare = true + args, cleanup, err := buildClaudeCLIArgs("claude-sonnet-5", bare) + if err != nil { + t.Fatal(err) + } + defer cleanup() + assertNoFlag(t, args, "--settings") + }) + + t.Run("CAPTAIN_MONITOR_HOOKS=off suppresses injection", func(t *testing.T) { + t.Setenv(api.MonitorHooksEnv, "off") + args, cleanup, err := buildClaudeCLIArgs("claude-sonnet-5", req) + if err != nil { + t.Fatal(err) + } + defer cleanup() + assertNoFlag(t, args, "--settings") + }) +} + +func TestCodexCLIMonitorHooksInjection(t *testing.T) { + binary := fakeCaptainOnPath(t) + req := ai.Request{Prompt: api.Prompt{User: "hi"}} + + args, cleanup, err := buildCodexCLIArgs("gpt-5", req) + if err != nil { + t.Fatalf("buildCodexCLIArgs: %v", err) + } + defer cleanup() + + notify := codexOverrideValue(t, args, "notify=") + var argv []string + if err := json.Unmarshal([]byte(strings.TrimPrefix(notify, "notify=")), &argv); err != nil { + t.Fatalf("notify override is not a JSON string array: %v", err) + } + want := []string{binary, "hook", "monitor", "notify", "--provider", "codex"} + if len(argv) != len(want) { + t.Fatalf("notify argv = %v, want %v", argv, want) + } + for i := range want { + if argv[i] != want[i] { + t.Fatalf("notify argv = %v, want %v", argv, want) + } + } + + t.Run("CAPTAIN_MONITOR_HOOKS=off suppresses injection", func(t *testing.T) { + t.Setenv(api.MonitorHooksEnv, "off") + args, cleanup, err := buildCodexCLIArgs("gpt-5", req) + if err != nil { + t.Fatal(err) + } + defer cleanup() + for _, arg := range args { + if strings.HasPrefix(arg, "notify=") { + t.Fatalf("notify override must be suppressed, got args %v", args) + } + } + }) +} + +// codexOverrideValue returns the -c override value with the given prefix. +func codexOverrideValue(t *testing.T, args []string, prefix string) string { + t.Helper() + for i := 0; i < len(args)-1; i++ { + if args[i] == "-c" && strings.HasPrefix(args[i+1], prefix) { + return args[i+1] + } + } + t.Fatalf("no -c %s… override in args %v", prefix, args) + return "" +} + +func assertNoFlag(t *testing.T, args []string, flag string) { + t.Helper() + for _, arg := range args { + if arg == flag { + t.Fatalf("args must not contain %s: %v", flag, args) + } + } +} diff --git a/pkg/ai/provider_test.go b/pkg/ai/provider_test.go index f3487be..b50fcc7 100644 --- a/pkg/ai/provider_test.go +++ b/pkg/ai/provider_test.go @@ -68,6 +68,7 @@ func TestAuthEnvVars(t *testing.T) { BackendClaudeCmux: nil, BackendOpenAI: {"OPENAI_API_KEY"}, BackendCodexCLI: {"OPENAI_API_KEY"}, + BackendCodexAgent: {"OPENAI_API_KEY"}, BackendCodexCmux: nil, BackendDeepSeek: {"DEEPSEEK_API_KEY"}, BackendGemini: {"GEMINI_API_KEY", "GOOGLE_API_KEY"}, diff --git a/pkg/ai/runtime_selection_ginkgo_test.go b/pkg/ai/runtime_selection_ginkgo_test.go new file mode 100644 index 0000000..3700d3c --- /dev/null +++ b/pkg/ai/runtime_selection_ginkgo_test.go @@ -0,0 +1,59 @@ +package ai + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/api" +) + +var _ = Describe("runtime model selection", func() { + It("infers a concrete backend for a bare model", func() { + model, err := ResolveModelSelectors(api.Model{Name: "gemini-3.5-flash"}) + + Expect(err).NotTo(HaveOccurred()) + Expect(model.Name).To(Equal("gemini-3.5-flash")) + Expect(model.Backend).To(Equal(api.BackendGemini)) + }) + + It("infers each bare multi-model runtime independently", func() { + models, err := ResolveRuntimeSelectors( + []string{"gemini-3.5-flash,gpt-5.5"}, + api.Model{Name: "opus", Backend: api.BackendClaudeAgent}, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(models).To(HaveLen(2)) + Expect(models[0].Name).To(Equal("gemini-3.5-flash")) + Expect(models[0].Backend).To(Equal(api.BackendGemini)) + Expect(models[1].Name).To(Equal("gpt-5.5")) + Expect(models[1].Backend).To(Equal(api.BackendOpenAI)) + }) + + It("rejects a known model forced through another family", func() { + _, err := ResolveModelSelectors(api.Model{ + Name: "gemini-3.5-flash", + Backend: api.BackendClaudeAgent, + }) + + Expect(err).To(MatchError(ContainSubstring("gemini"))) + Expect(err).To(MatchError(ContainSubstring("claude-agent"))) + }) + + It("allows an unknown provider model when its backend is explicit", func() { + model, err := ResolveModelSelectors(api.Model{ + Name: "provider-future-model", + Backend: api.BackendGemini, + }) + + Expect(err).NotTo(HaveOccurred()) + Expect(model.Name).To(Equal("provider-future-model")) + Expect(model.Backend).To(Equal(api.BackendGemini)) + }) + + It("requires a backend for an unknown bare model", func() { + _, err := ResolveModelSelectors(api.Model{Name: "provider-future-model"}) + + Expect(err).To(MatchError(ContainSubstring("pass an explicit backend"))) + }) +}) diff --git a/pkg/ai/runtime_selector.go b/pkg/ai/runtime_selector.go new file mode 100644 index 0000000..a0ec8af --- /dev/null +++ b/pkg/ai/runtime_selector.go @@ -0,0 +1,41 @@ +package ai + +import ( + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" +) + +// Model selectors are parsed by pkg/api/registry — the same parser that decodes +// specs and prompt frontmatter. This file used to hold a second implementation of +// the grammar (resolveSelectorPart, selectorBackend, selectorModelFamily, +// wildcardBackends, splitSelectorEffort) over its own family table, which is why +// `--model agent:sol` and `model: agent:sol` in frontmatter disagreed. + +// ContainsRuntimeSelector reports whether s contains a backend/mode selector such +// as "cmux:gpt-5.6-sol" or "*:sonnet-5". Comma-separated fallback lists are +// scanned item by item. +func ContainsRuntimeSelector(s string) bool { + return registry.ContainsSelector(s) +} + +// ResolveModelSelectors turns every recognizable model name into a concrete +// model/backend pair while preserving fallback semantics. Unknown model names +// require an explicit backend. +func ResolveModelSelectors(model api.Model) (api.Model, error) { + return registry.ResolveModel(model) +} + +// ResolveRuntimeSelectors expands --multi-models values into concrete runtime +// model/backend pairs. Values may be repeated and/or comma-separated. +func ResolveRuntimeSelectors(values []string, base api.Model) ([]api.Model, error) { + return registry.ResolveMulti(values, base) +} + +// NormalizeModelForBackend maps a catalog/live/provider model id onto the exact +// runtime id accepted by backend. Provider prefixes and old backend-specific +// aliases are input compatibility only; the returned value is never a family +// alias such as "opus" or a synthetic id such as "claude-agent-opus". +func NormalizeModelForBackend(backend api.Backend, model string) string { + resolved, _ := ResolveExactModelForBackend(backend, model) + return resolved +} diff --git a/pkg/ai/runtime_selector_test.go b/pkg/ai/runtime_selector_test.go new file mode 100644 index 0000000..b4ce75d --- /dev/null +++ b/pkg/ai/runtime_selector_test.go @@ -0,0 +1,347 @@ +package ai + +import ( + "reflect" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" +) + +func TestResolveModelSelectors(t *testing.T) { + cases := []struct { + name string + in api.Model + wantName string + wantBackend api.Backend + }{ + { + name: "api claude shorthand", + in: api.Model{Name: "api:sonnet-5"}, + wantName: "claude-sonnet-5", + wantBackend: api.BackendAnthropic, + }, + { + name: "cmux codex shorthand", + in: api.Model{Name: "cmux:gpt-5.5"}, + wantName: "gpt-5.5", + wantBackend: api.BackendCodexCmux, + }, + { + name: "agent claude shorthand", + in: api.Model{Name: "agent:opus"}, + wantName: "claude-opus-5", + wantBackend: api.BackendClaudeAgent, + }, + { + name: "agent codex app server", + in: api.Model{Name: "agent:gpt-5.5"}, + wantName: "gpt-5.5", + wantBackend: api.BackendCodexAgent, + }, + { + name: "exact backend", + in: api.Model{Name: "claude-cmux:opus"}, + wantName: "claude-opus-5", + wantBackend: api.BackendClaudeCmux, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, err := ResolveModelSelectors(tc.in) + if err != nil { + t.Fatalf("ResolveModelSelectors: %v", err) + } + if got.Name != tc.wantName || got.Backend != tc.wantBackend { + t.Fatalf("got %s/%s, want %s/%s", got.Backend, got.Name, tc.wantBackend, tc.wantName) + } + }) + } +} + +func TestNormalizeModelForBackend(t *testing.T) { + cases := []struct { + backend api.Backend + model string + want string + }{ + {api.BackendClaudeAgent, "opus-4-8", "claude-opus-4-8"}, + {api.BackendClaudeCmux, "claude-opus-4-8", "claude-opus-4-8"}, + {api.BackendClaudeCLI, "fable-5", "claude-fable-5"}, + {api.BackendAnthropic, "opus-4-8", "claude-opus-4-8"}, + {api.BackendClaudeAgent, "claude-agent-opus", "claude-opus-5"}, + {api.BackendClaudeAgent, "sonnet", "claude-sonnet-5"}, + {api.BackendClaudeAgent, "sonnet-4", "claude-sonnet-4-6"}, + {api.BackendClaudeAgent, "claude-sonnet-4-5", "claude-sonnet-4-6"}, + {api.BackendClaudeAgent, "haiku", "claude-haiku-4-5"}, + {api.BackendClaudeAgent, "haiku-4", "claude-haiku-4-5"}, + {api.BackendClaudeAgent, "haiku-4.5", "claude-haiku-4-5"}, + {api.BackendCodexCmux, "openai/gpt-5.5", "gpt-5.5"}, + {api.BackendCodexAgent, "gpt-5.4-mini", "gpt-5.4-mini"}, + } + for _, tc := range cases { + if got := NormalizeModelForBackend(tc.backend, tc.model); got != tc.want { + t.Fatalf("NormalizeModelForBackend(%s, %q) = %q, want %q", tc.backend, tc.model, got, tc.want) + } + } +} + +func TestParseModelIdentity(t *testing.T) { + cases := []struct { + model string + want ModelIdentity + }{ + {"opus-4-8", ModelIdentity{Provider: registry.Anthropic.Name, Family: "opus", Version: "4.8"}}, + {"claude-opus-4-8-3003", ModelIdentity{Provider: registry.Anthropic.Name, Family: "opus", Version: "4.8-3003"}}, + {"claude-agent-sonnet", ModelIdentity{Provider: registry.Anthropic.Name, Family: "sonnet"}}, + {"openai/gpt-5.5", ModelIdentity{Provider: registry.OpenAI.Name, Family: "gpt", Version: "5.5"}}, + {"codex-gpt-5.4", ModelIdentity{Provider: registry.OpenAI.Name, Family: "gpt", Version: "5.4"}}, + {"gpt-5.4-mini", ModelIdentity{Provider: registry.OpenAI.Name, Family: "gpt", Version: "5.4-mini"}}, + } + for _, tc := range cases { + got, ok := ParseModelIdentity(registry.Anthropic.Name, tc.model) + if !ok { + t.Fatalf("ParseModelIdentity(%q) did not parse", tc.model) + } + if got != tc.want { + t.Fatalf("ParseModelIdentity(%q) = %+v, want %+v", tc.model, got, tc.want) + } + } +} + +func TestResolveModelSelectors_FallbackSelectors(t *testing.T) { + got, err := ResolveModelSelectors(api.Model{Name: "api:sonnet-5,cmux:opus"}) + if err != nil { + t.Fatalf("ResolveModelSelectors: %v", err) + } + if got.Name != "claude-sonnet-5" || got.Backend != api.BackendAnthropic { + t.Fatalf("primary = %s/%s", got.Backend, got.Name) + } + if len(got.Fallbacks) != 1 { + t.Fatalf("fallback count = %d, want 1", len(got.Fallbacks)) + } + if fb := got.Fallbacks[0]; fb.Name != "claude-opus-5" || fb.Backend != api.BackendClaudeCmux { + t.Fatalf("fallback = %s/%s, want %s/%s", fb.Backend, fb.Name, api.BackendClaudeCmux, "claude-opus-5") + } +} + +func TestResolveRuntimeSelectors(t *testing.T) { + got, err := ResolveRuntimeSelectors( + []string{"cli:sonnet-5,cmux:opus", "*:gpt-5.5"}, + api.Model{Name: "claude-sonnet-5", Backend: api.BackendAnthropic}, + ) + if err != nil { + t.Fatalf("ResolveRuntimeSelectors: %v", err) + } + var labels []string + for _, model := range got { + labels = append(labels, string(model.Backend)+":"+model.Name) + } + want := []string{ + "claude-cli:claude-sonnet-5", + "claude-cmux:claude-opus-5", + "openai:gpt-5.5", + "codex-agent:gpt-5.5", + "codex-cli:gpt-5.5", + "codex-cmux:gpt-5.5", + } + if !reflect.DeepEqual(labels, want) { + t.Fatalf("labels = %#v, want %#v", labels, want) + } +} + +func TestResolveRuntimeSelectors_RequiresModelForPrefixOnly(t *testing.T) { + got, err := ResolveRuntimeSelectors([]string{"cmux"}, api.Model{Name: "gpt-5.5"}) + if err != nil { + t.Fatalf("ResolveRuntimeSelectors: %v", err) + } + if len(got) != 1 || got[0].Backend != api.BackendCodexCmux || got[0].Name != "gpt-5.5" { + t.Fatalf("got %#v, want codex-cmux:gpt-5.5", got) + } + if _, err := ResolveRuntimeSelectors([]string{"cmux"}, api.Model{}); err == nil { + t.Fatal("expected missing base model error") + } +} + +func TestResolveRuntimeSelectors_PreservesExplicitUncatalogedVariant(t *testing.T) { + got, err := ResolveRuntimeSelectors([]string{"agent:gpt-5.4-mini"}, api.Model{}) + if err != nil { + t.Fatalf("ResolveRuntimeSelectors: %v", err) + } + if len(got) != 1 { + t.Fatalf("got %d models, want 1: %#v", len(got), got) + } + if got[0].Backend != api.BackendCodexAgent || got[0].Name != "gpt-5.4-mini" { + t.Fatalf("got %s/%s, want %s/gpt-5.4-mini", got[0].Backend, got[0].Name, api.BackendCodexAgent) + } +} + +func TestResolveModelSelectors_WildcardOnlyForMultiModels(t *testing.T) { + if _, err := ResolveModelSelectors(api.Model{Name: "*:sonnet-5"}); err == nil { + t.Fatal("expected wildcard selector to be rejected for single model") + } +} + +func TestResolveModelSelectors_UnknownPrefixFails(t *testing.T) { + if _, err := ResolveModelSelectors(api.Model{Name: "nope:sonnet-5"}); err == nil { + t.Fatal("expected unknown selector prefix error") + } +} + +func TestResolveModelSelectors_EffortQualifiedAlias(t *testing.T) { + got, err := ResolveModelSelectors(api.Model{ + Name: "agent:sol:high", + Effort: api.EffortLow, + }) + if err != nil { + t.Fatalf("ResolveModelSelectors: %v", err) + } + if got.Backend != api.BackendCodexAgent || got.Name != "gpt-5.6-sol" || got.Effort != api.EffortHigh { + t.Fatalf("got %s/%s/%s, want codex-agent/gpt-5.6-sol/high", got.Backend, got.Name, got.Effort) + } +} + +func TestResolveRuntimeSelectors_PerSelectorEffortAndDedup(t *testing.T) { + got, err := ResolveRuntimeSelectors( + []string{"agent:sol:high,agent:sol:xhigh,cmux:terra:max"}, + api.Model{Effort: api.EffortLow}, + ) + if err != nil { + t.Fatalf("ResolveRuntimeSelectors: %v", err) + } + // Compare the identity triple: this pins per-selector effort and dedup. The + // derived capability fields are covered by TestResolvedModelCarriesCapabilities. + type runtime struct { + name string + backend api.Backend + effort api.Effort + } + want := []runtime{ + {"gpt-5.6-sol", api.BackendCodexAgent, api.EffortHigh}, + {"gpt-5.6-sol", api.BackendCodexAgent, api.EffortXHigh}, + {"gpt-5.6-terra", api.BackendCodexCmux, api.EffortMax}, + } + gotRuntimes := make([]runtime, 0, len(got)) + for _, m := range got { + gotRuntimes = append(gotRuntimes, runtime{m.Name, m.Backend, m.Effort}) + } + if !reflect.DeepEqual(gotRuntimes, want) { + t.Fatalf("got %+v, want %+v", gotRuntimes, want) + } +} + +// TestResolvedModelCarriesCapabilities pins that resolution answers what the +// chosen adapter can do, not just which model it is. The values are asserted +// against the known adapters: only the claude agent SDK steers, both agent SDKs +// interrupt, and the claude CLI carries no attachments. +func TestResolvedModelCarriesCapabilities(t *testing.T) { + cases := []struct { + selector string + wantMode registry.RuntimeMode + wantResume bool + wantIntr bool + wantSteer bool + wantMedia []string + }{ + {"agent:sonnet", registry.ModeAgent, true, true, true, []string{"image/png", "image/jpeg", "image/gif", "image/webp", "application/pdf"}}, + {"cli:sonnet", registry.ModeCLI, true, false, false, []string{}}, + {"api:sonnet", registry.ModeAPI, false, false, false, []string{"image/*"}}, + {"agent:sol", registry.ModeAgent, true, true, false, []string{"image/*"}}, + } + for _, tc := range cases { + t.Run(tc.selector, func(t *testing.T) { + got, err := ResolveModelSelectors(api.Model{Name: tc.selector}) + if err != nil { + t.Fatalf("ResolveModelSelectors(%q): %v", tc.selector, err) + } + if got.Mode != tc.wantMode { + t.Errorf("Mode = %q, want %q", got.Mode, tc.wantMode) + } + if !got.Streaming { + t.Error("Streaming = false; every adapter streams") + } + if got.Resume != tc.wantResume || got.Interrupt != tc.wantIntr || got.Steer != tc.wantSteer { + t.Errorf("resume/interrupt/steer = %v/%v/%v, want %v/%v/%v", + got.Resume, got.Interrupt, got.Steer, tc.wantResume, tc.wantIntr, tc.wantSteer) + } + if !reflect.DeepEqual(got.MediaTypes, tc.wantMedia) { + t.Errorf("MediaTypes = %v, want %v", got.MediaTypes, tc.wantMedia) + } + if got.Provider == nil { + t.Fatal("Provider is nil; a resolved model must know its family") + } + }) + } +} + +// TestModeContradictingBackendFailsLoud: Backend is exactly (provider, mode), so +// a spec naming both must not name two different runtimes. +func TestModeContradictingBackendFailsLoud(t *testing.T) { + _, err := ResolveModelSelectors(api.Model{Name: "sonnet", Backend: api.BackendAnthropic, Mode: registry.ModeAgent}) + if err == nil { + t.Fatal("mode agent + backend anthropic should fail loud, not silently pick one") + } + if !strings.Contains(err.Error(), "contradicts") { + t.Errorf("err = %v, want it to name the contradiction", err) + } +} + +// TestModeSelectsBackend: {model, mode} is the object form of the compact +// "mode:model" string and must resolve identically. +func TestModeSelectsBackend(t *testing.T) { + byMode, err := ResolveModelSelectors(api.Model{Name: "sonnet", Mode: registry.ModeAgent}) + if err != nil { + t.Fatalf("mode form: %v", err) + } + byPrefix, err := ResolveModelSelectors(api.Model{Name: "agent:sonnet"}) + if err != nil { + t.Fatalf("compact form: %v", err) + } + if byMode.Backend != byPrefix.Backend || byMode.Name != byPrefix.Name { + t.Errorf("mode form = %s/%s, compact form = %s/%s; the two spellings must agree", + byMode.Backend, byMode.Name, byPrefix.Backend, byPrefix.Name) + } +} + +func TestResolveRuntimeSelectors_WildcardRespectsAvailability(t *testing.T) { + got, err := ResolveRuntimeSelectors([]string{"*:sol:high"}, api.Model{}) + if err != nil { + t.Fatalf("ResolveRuntimeSelectors: %v", err) + } + var backends []api.Backend + for _, model := range got { + backends = append(backends, model.Backend) + if model.Name != "gpt-5.6-sol" || model.Effort != api.EffortHigh { + t.Fatalf("unexpected model: %+v", model) + } + } + want := []api.Backend{api.BackendOpenAI, api.BackendCodexAgent, api.BackendCodexCLI, api.BackendCodexCmux} + if !reflect.DeepEqual(backends, want) { + t.Fatalf("backends = %v, want %v", backends, want) + } +} + +func TestResolveModelSelectors_EffortErrors(t *testing.T) { + if _, err := ResolveModelSelectors(api.Model{Name: "agent:sol:extreme"}); err == nil { + t.Fatal("expected invalid effort suffix error") + } +} + +func TestResolveModelSelectors_OpenAI56VariantsAvailableViaAPI(t *testing.T) { + for alias, want := range map[string]string{ + "luna": "gpt-5.6-luna", + "sol": "gpt-5.6-sol", + "terra": "gpt-5.6-terra", + } { + t.Run(alias, func(t *testing.T) { + got, err := ResolveModelSelectors(api.Model{Name: "api:" + alias}) + if err != nil { + t.Fatalf("ResolveModelSelectors(api:%s): %v", alias, err) + } + if got.Backend != api.BackendOpenAI || got.Name != want { + t.Fatalf("got %s/%s, want openai/%s", got.Backend, got.Name, want) + } + }) + } +} diff --git a/pkg/ai/schema.go b/pkg/ai/schema.go index f18a2c9..aab200c 100644 --- a/pkg/ai/schema.go +++ b/pkg/ai/schema.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "reflect" + "sort" "strings" "github.com/flanksource/captain/pkg/api" @@ -150,6 +151,298 @@ func SchemaJSONFor(p api.Prompt) (json.RawMessage, error) { return json.Marshal(schema) } +// SchemaJSONForBackend resolves the schema a provider should send to its +// backend. Native structured-output backends receive a transformed subset of +// the caller's original JSON Schema when their API imposes stricter schema +// rules; every other backend receives the original schema from SchemaJSONFor. +func SchemaJSONForBackend(backend api.Backend, p api.Prompt) (json.RawMessage, error) { + schema, err := SchemaJSONFor(p) + if err != nil || len(schema) == 0 { + return schema, err + } + switch { + case UsesAnthropicSchemaSubset(backend): + return AnthropicCompatibleSchema(schema) + case UsesOpenAISchemaSubset(backend): + return OpenAICompatibleSchema(schema) + default: + return schema, nil + } +} + +// UsesAnthropicSchemaSubset reports whether a backend sends JSON Schema to +// Claude's native structured-output machinery and therefore needs Captain to +// apply Anthropic's supported-subset transformation before dispatch. +func UsesAnthropicSchemaSubset(backend api.Backend) bool { + switch backend { + case api.BackendAnthropic, api.BackendClaudeAgent, api.BackendClaudeCLI: + return true + default: + return false + } +} + +// UsesOpenAISchemaSubset reports whether a backend sends JSON Schema to +// OpenAI's native strict structured-output machinery. Prompt-only Codex cmux +// sessions are deliberately excluded because they do not submit a response +// format to OpenAI. +func UsesOpenAISchemaSubset(backend api.Backend) bool { + switch backend { + case api.BackendOpenAI, api.BackendCodexCLI, api.BackendCodexAgent: + return true + default: + return false + } +} + +// OpenAICompatibleSchema returns a provider-facing copy of schema that obeys +// OpenAI strict structured-output object rules: every declared property is +// required and every object is closed. Property types and constraints are left +// unchanged so the original schema remains the source of truth for local +// validation. Open-ended map schemas cannot be represented by this subset and +// fail loudly instead of being silently narrowed to an empty object. +func OpenAICompatibleSchema(schema json.RawMessage) (json.RawMessage, error) { + var v any + if err := json.Unmarshal(schema, &v); err != nil { + return nil, fmt.Errorf("openai schema transform: invalid JSON schema: %w", err) + } + if err := normalizeOpenAISchema(v, "$"); err != nil { + return nil, err + } + out, err := json.Marshal(v) + if err != nil { + return nil, fmt.Errorf("openai schema transform: marshal schema: %w", err) + } + return out, nil +} + +func normalizeOpenAISchema(v any, path string) error { + node, ok := v.(map[string]any) + if !ok { + return nil + } + refOnly, err := normalizeOpenAIRef(node, path) + if err != nil { + return err + } + if refOnly { + return nil + } + + if isObjectSchema(node) { + if additional, exists := node["additionalProperties"]; exists { + allowed, isBool := additional.(bool) + if !isBool || allowed { + return fmt.Errorf("openai schema transform: %s is open-ended; additionalProperties must be false", path) + } + } + + var required []string + if rawProperties, exists := node["properties"]; exists { + properties, ok := rawProperties.(map[string]any) + if !ok { + return fmt.Errorf("openai schema transform: %s.properties must be an object", path) + } + required = make([]string, 0, len(properties)) + for name := range properties { + required = append(required, name) + } + sort.Strings(required) + } + if required == nil { + required = []string{} + } + node["required"] = required + node["additionalProperties"] = false + } + + for _, key := range []string{"properties", "$defs", "definitions", "patternProperties", "dependentSchemas"} { + children, _ := node[key].(map[string]any) + for name, child := range children { + if err := normalizeOpenAISchema(child, path+"."+key+"."+name); err != nil { + return err + } + } + } + for _, key := range []string{"items", "contains", "propertyNames", "not", "if", "then", "else"} { + if child, ok := node[key].(map[string]any); ok { + if err := normalizeOpenAISchema(child, path+"."+key); err != nil { + return err + } + } + } + for _, key := range []string{"allOf", "anyOf", "oneOf", "prefixItems"} { + children, _ := node[key].([]any) + for i, child := range children { + if err := normalizeOpenAISchema(child, fmt.Sprintf("%s.%s[%d]", path, key, i)); err != nil { + return err + } + } + } + return nil +} + +var openAIRefAnnotationKeywords = map[string]bool{ + "$id": true, + "$comment": true, + "default": true, + "deprecated": true, + "description": true, + "examples": true, + "readOnly": true, + "title": true, + "writeOnly": true, +} + +// normalizeOpenAIRef makes referenced nodes conform to OpenAI's strict schema +// subset. A nested $ref must stand alone; invopop/jsonschema commonly attaches +// field annotations such as description alongside it, which OpenAI rejects. +// Root schema dialect metadata and definitions are retained so the reference +// remains resolvable. $id is stripped even at the root because OpenAI rejects +// it as a $ref sibling. Unexpected semantic siblings fail loudly rather than +// being silently discarded. +func normalizeOpenAIRef(node map[string]any, path string) (bool, error) { + rawRef, hasRef := node["$ref"] + if !hasRef { + return false, nil + } + ref, ok := rawRef.(string) + if !ok || strings.TrimSpace(ref) == "" { + return false, fmt.Errorf("openai schema transform: %s.$ref must be a non-empty string", path) + } + + for key := range node { + if key == "$ref" { + continue + } + if openAIRefAnnotationKeywords[key] { + delete(node, key) + continue + } + if path == "$" && (key == "$schema" || key == "$defs" || key == "definitions") { + continue + } + return false, fmt.Errorf("openai schema transform: %s.$ref has unsupported sibling %q", path, key) + } + + return path != "$", nil +} + +// AnthropicCompatibleSchema returns a copy of schema with constraints that +// Claude structured outputs do not accept removed from the provider-facing +// payload. The original schema must still be used for local validation. +func AnthropicCompatibleSchema(schema json.RawMessage) (json.RawMessage, error) { + var v any + if err := json.Unmarshal(schema, &v); err != nil { + return nil, fmt.Errorf("anthropic schema transform: invalid JSON schema: %w", err) + } + sanitizeAnthropicSchema(v) + out, err := json.Marshal(v) + if err != nil { + return nil, fmt.Errorf("anthropic schema transform: marshal schema: %w", err) + } + return out, nil +} + +var anthropicUnsupportedConstraints = map[string]string{ + "minimum": "minimum", + "maximum": "maximum", + "exclusiveMinimum": "exclusiveMinimum", + "exclusiveMaximum": "exclusiveMaximum", + "multipleOf": "multipleOf", + "minLength": "minLength", + "maxLength": "maxLength", + "pattern": "pattern", + "minItems": "minItems", + "maxItems": "maxItems", + "uniqueItems": "uniqueItems", + "contains": "contains", + "minContains": "minContains", + "maxContains": "maxContains", + "minProperties": "minProperties", + "maxProperties": "maxProperties", + "dependentRequired": "dependentRequired", + "dependentSchemas": "dependentSchemas", + "propertyNames": "propertyNames", + "unevaluatedItems": "unevaluatedItems", + "unevaluatedProperties": "unevaluatedProperties", +} + +var anthropicSupportedFormats = map[string]bool{ + "date": true, + "date-time": true, + "email": true, + "hostname": true, + "ipv4": true, + "ipv6": true, + "uri": true, + "uuid": true, +} + +func sanitizeAnthropicSchema(v any) { + switch node := v.(type) { + case map[string]any: + var hints []string + for key, label := range anthropicUnsupportedConstraints { + if value, ok := node[key]; ok { + hints = append(hints, label+"="+jsonValue(value)) + delete(node, key) + } + } + if value, ok := node["format"].(string); ok && !anthropicSupportedFormats[value] { + hints = append(hints, "format="+value) + delete(node, "format") + } + + for _, value := range node { + sanitizeAnthropicSchema(value) + } + if isObjectSchema(node) { + if _, ok := node["additionalProperties"]; !ok { + node["additionalProperties"] = false + } + } + if len(hints) > 0 { + node["description"] = appendConstraintDescription(node["description"], hints) + } + case []any: + for _, value := range node { + sanitizeAnthropicSchema(value) + } + } +} + +func isObjectSchema(node map[string]any) bool { + if t, ok := node["type"].(string); ok && t == "object" { + return true + } + if ts, ok := node["type"].([]any); ok { + for _, t := range ts { + if s, ok := t.(string); ok && s == "object" { + return true + } + } + } + _, hasProperties := node["properties"] + return hasProperties +} + +func appendConstraintDescription(existing any, hints []string) string { + prefix := strings.TrimSpace(fmt.Sprint(existing)) + if prefix == "" || prefix == "" { + return "Constraints preserved for local validation: " + strings.Join(hints, ", ") + "." + } + return strings.TrimRight(prefix, ".") + ". Constraints preserved for local validation: " + strings.Join(hints, ", ") + "." +} + +func jsonValue(v any) string { + data, err := json.Marshal(v) + if err != nil { + return fmt.Sprint(v) + } + return string(data) +} + // BindStructuredOutput unmarshals a provider's validated structured-output JSON // into the caller's Go target. Missing or malformed output fails loudly with // ErrSchemaValidation rather than leaving the target zero-valued. diff --git a/pkg/ai/schema_prompt.go b/pkg/ai/schema_prompt.go new file mode 100644 index 0000000..bd3063e --- /dev/null +++ b/pkg/ai/schema_prompt.go @@ -0,0 +1,47 @@ +package ai + +import ( + "encoding/json" + "fmt" + "strings" +) + +// WithSchemaPrompt makes a structured-output request runnable on a backend that +// cannot enforce a schema natively (cmux, gemini-cli) by appending a JSON-only +// schema instruction to the user prompt and clearing the native schema fields so +// the run is a plain text turn. The returned schema is what +// ValidatedStructuredData checks the reply against once the run finishes. +func WithSchemaPrompt(req Request) (Request, json.RawMessage, error) { + schema, err := SchemaJSONFor(req.Prompt) + if err != nil { + return req, nil, err + } + if len(schema) > 0 { + req.Prompt.User = strings.TrimRight(req.Prompt.User, "\n") + "\n\n" + SchemaInstruction(string(schema)) + } + req.Prompt.Schema = nil + req.Prompt.SchemaJSON = nil + return req, schema, nil +} + +// ValidatedStructuredData extracts the JSON object a WithSchemaPrompt run asked +// for and validates it against the schema. A native terminal outcome (plan exit, +// ask) means the run ended on control flow rather than the requested answer, so +// there is nothing to extract. +func ValidatedStructuredData(schema json.RawMessage, text string, outcome *TerminalOutcome) (json.RawMessage, error) { + if len(schema) == 0 || outcome != nil { + return nil, nil + } + object, ok := ExtractJSONObject(text) + if !ok { + return nil, fmt.Errorf("%w: response carried no JSON object", ErrSchemaValidation) + } + violations, err := ValidateStructuredJSON(schema, object) + if err != nil { + return nil, err + } + if violations != "" { + return nil, fmt.Errorf("%w: %s", ErrSchemaValidation, violations) + } + return json.RawMessage(object), nil +} diff --git a/pkg/ai/schema_prompt_test.go b/pkg/ai/schema_prompt_test.go new file mode 100644 index 0000000..4156d4f --- /dev/null +++ b/pkg/ai/schema_prompt_test.go @@ -0,0 +1,70 @@ +package ai + +import ( + "encoding/json" + "errors" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/api" +) + +func TestWithSchemaPrompt(t *testing.T) { + // A raw JSON schema is appended to the prompt and the native fields cleared, + // so the run is a plain text turn that still asks for JSON. + req := Request{Prompt: api.Prompt{ + User: "review the diff", + SchemaJSON: []byte(`{"type":"object","required":["pass"]}`), + }} + got, schema, err := WithSchemaPrompt(req) + if err != nil { + t.Fatalf("WithSchemaPrompt: %v", err) + } + if got.Prompt.Schema != nil || got.Prompt.SchemaJSON != nil { + t.Errorf("native schema fields must be cleared, got Schema=%v SchemaJSON=%s", got.Prompt.Schema, got.Prompt.SchemaJSON) + } + if string(schema) != string(req.Prompt.SchemaJSON) { + t.Errorf("preserved schema = %s, want %s", schema, req.Prompt.SchemaJSON) + } + if !strings.Contains(got.Prompt.User, "review the diff") { + t.Errorf("original prompt lost: %q", got.Prompt.User) + } + if !strings.Contains(got.Prompt.User, `"required":["pass"]`) { + t.Errorf("schema not appended to prompt: %q", got.Prompt.User) + } + + // A text-mode request is returned unchanged. + plain := Request{Prompt: api.Prompt{User: "hi"}} + got, schema, err = WithSchemaPrompt(plain) + if err != nil { + t.Fatalf("WithSchemaPrompt(text): %v", err) + } + if got.Prompt.User != "hi" { + t.Errorf("text prompt altered: %q", got.Prompt.User) + } + if len(schema) != 0 { + t.Errorf("text prompt preserved unexpected schema: %s", schema) + } +} + +func TestValidatedStructuredData(t *testing.T) { + schema := json.RawMessage(`{"type":"object","required":["pass"],"properties":{"pass":{"type":"boolean"}}}`) + + got, err := ValidatedStructuredData(schema, "Result:\n```json\n{\"pass\":true}\n```", nil) + if err != nil { + t.Fatalf("ValidatedStructuredData(valid) error = %v", err) + } + if string(got) != `{"pass":true}` { + t.Fatalf("ValidatedStructuredData(valid) = %s", got) + } + + if _, err := ValidatedStructuredData(schema, `{"pass":"yes"}`, nil); !errors.Is(err, ErrSchemaValidation) { + t.Fatalf("ValidatedStructuredData(invalid) error = %v, want ErrSchemaValidation", err) + } + + outcome := &TerminalOutcome{Kind: TerminalOutcomePlan, Plan: &TerminalPlan{Content: "1. Inspect"}} + got, err = ValidatedStructuredData(schema, "not a schema envelope", outcome) + if err != nil || got != nil { + t.Fatalf("native terminal outcome must bypass schema extraction, got (%s, %v)", got, err) + } +} diff --git a/pkg/ai/schema_test.go b/pkg/ai/schema_test.go index d06ac8e..1283569 100644 --- a/pkg/ai/schema_test.go +++ b/pkg/ai/schema_test.go @@ -3,6 +3,8 @@ package ai import ( "encoding/json" "errors" + "reflect" + "strings" "testing" "github.com/flanksource/captain/pkg/api" @@ -46,6 +48,299 @@ func TestSchemaJSONFor(t *testing.T) { } } +func TestSchemaJSONForBackend_AnthropicSanitizesUnsupportedConstraints(t *testing.T) { + raw := json.RawMessage(`{ + "type":"object", + "properties":{ + "title":{"type":"string","minLength":3,"maxLength":40,"format":"slug"}, + "items":{"type":"array","minItems":1,"maxItems":2,"items":{"type":"object","properties":{"name":{"type":"string","pattern":"^[a-z]+$"}}}} + }, + "required":["title"] + }`) + got, err := SchemaJSONForBackend(api.BackendAnthropic, api.Prompt{SchemaJSON: raw}) + if err != nil { + t.Fatalf("SchemaJSONForBackend: %v", err) + } + if string(raw) == string(got) { + t.Fatal("anthropic schema should be transformed, got original") + } + if !json.Valid(raw) || !containsJSONKey(t, raw, "maxItems") { + t.Fatalf("original schema should remain unchanged: %s", raw) + } + + var schema map[string]any + if err := json.Unmarshal(got, &schema); err != nil { + t.Fatalf("transformed schema invalid JSON: %v", err) + } + if schema["additionalProperties"] != false { + t.Fatalf("root additionalProperties = %v, want false", schema["additionalProperties"]) + } + title := schema["properties"].(map[string]any)["title"].(map[string]any) + for _, key := range []string{"minLength", "maxLength", "format"} { + if _, ok := title[key]; ok { + t.Fatalf("title still has unsupported %s: %#v", key, title) + } + } + if desc := title["description"].(string); !strings.Contains(desc, "minLength=3") || !strings.Contains(desc, "format=slug") { + t.Fatalf("title description = %q, want removed constraints", desc) + } + items := schema["properties"].(map[string]any)["items"].(map[string]any) + if _, ok := items["maxItems"]; ok { + t.Fatalf("items still has maxItems: %#v", items) + } + nested := items["items"].(map[string]any) + if nested["additionalProperties"] != false { + t.Fatalf("nested additionalProperties = %v, want false", nested["additionalProperties"]) + } + name := nested["properties"].(map[string]any)["name"].(map[string]any) + if _, ok := name["pattern"]; ok { + t.Fatalf("name still has pattern: %#v", name) + } +} + +func TestSchemaJSONForBackend_OpenAIRequiresAllCommitMessageProperties(t *testing.T) { + original := json.RawMessage(`{ + "type":"object", + "additionalProperties":false, + "required":["type","subject"], + "properties":{ + "type":{"type":"string"}, + "scope":{"type":"string","description":"Optional scope"}, + "subject":{"type":"string"}, + "body":{"type":"string","description":"Optional body explaining why and impact"} + } + }`) + prompt := api.Prompt{SchemaJSON: original} + got, err := SchemaJSONForBackend(api.BackendCodexAgent, prompt) + if err != nil { + t.Fatalf("SchemaJSONForBackend: %v", err) + } + + originalObject := decodeSchemaObject(t, original) + if required := schemaRequired(t, originalObject); !reflect.DeepEqual(required, []string{"type", "subject"}) { + t.Fatalf("original required = %v, want [type subject]", required) + } + if originalObject["additionalProperties"] != false { + t.Fatal("provider transform mutated the original commit-message schema") + } + + strictObject := decodeSchemaObject(t, got) + if required := schemaRequired(t, strictObject); !reflect.DeepEqual(required, []string{"body", "scope", "subject", "type"}) { + t.Fatalf("openai required = %v, want every property in sorted order", required) + } + if strictObject["additionalProperties"] != false { + t.Fatalf("additionalProperties = %v, want false", strictObject["additionalProperties"]) + } + body := strictObject["properties"].(map[string]any)["body"].(map[string]any) + if body["type"] != "string" { + t.Fatalf("optional body type = %v, want string (not nullable)", body["type"]) + } +} + +func TestOpenAICompatibleSchema_RecursesWithoutMutatingInput(t *testing.T) { + raw := json.RawMessage(`{ + "type":"object", + "properties":{ + "result":{"type":"object","properties":{"z":{"type":"string"},"a":{"type":"string"}}}, + "rows":{"type":"array","items":{"type":"object","properties":{"value":{"type":"integer"}}}} + }, + "$defs":{"detail":{"type":"object","properties":{"note":{"type":"string"}}}} + }`) + original := append(json.RawMessage(nil), raw...) + + got, err := OpenAICompatibleSchema(raw) + if err != nil { + t.Fatalf("OpenAICompatibleSchema: %v", err) + } + if string(raw) != string(original) { + t.Fatal("OpenAICompatibleSchema mutated its input") + } + + root := decodeSchemaObject(t, got) + if required := schemaRequired(t, root); !reflect.DeepEqual(required, []string{"result", "rows"}) { + t.Fatalf("root required = %v", required) + } + properties := root["properties"].(map[string]any) + result := properties["result"].(map[string]any) + if required := schemaRequired(t, result); !reflect.DeepEqual(required, []string{"a", "z"}) { + t.Fatalf("nested required = %v", required) + } + row := properties["rows"].(map[string]any)["items"].(map[string]any) + if required := schemaRequired(t, row); !reflect.DeepEqual(required, []string{"value"}) { + t.Fatalf("array item required = %v", required) + } + detail := root["$defs"].(map[string]any)["detail"].(map[string]any) + if required := schemaRequired(t, detail); !reflect.DeepEqual(required, []string{"note"}) { + t.Fatalf("$defs required = %v", required) + } + for name, object := range map[string]map[string]any{"root": root, "result": result, "row": row, "detail": detail} { + if object["additionalProperties"] != false { + t.Errorf("%s additionalProperties = %v, want false", name, object["additionalProperties"]) + } + } +} + +func TestOpenAICompatibleSchema_RemovesAnnotationsFromReferences(t *testing.T) { + raw := json.RawMessage(`{ + "$schema":"https://json-schema.org/draft/2020-12/schema", + "$id":"https://example.com/plan-envelope", + "$ref":"#/$defs/PlanEnvelope", + "$defs":{ + "PlanEnvelope":{ + "type":"object", + "properties":{ + "summary":{"type":"string","description":"What the session did"}, + "plan":{"$ref":"#/$defs/PlanResult","description":"The plan this session produced"} + } + }, + "PlanResult":{ + "type":"object", + "properties":{ + "status":{"type":"string"}, + "path":{"type":"string"} + } + } + } + }`) + original := append(json.RawMessage(nil), raw...) + + got, err := OpenAICompatibleSchema(raw) + if err != nil { + t.Fatalf("OpenAICompatibleSchema: %v", err) + } + if string(raw) != string(original) { + t.Fatal("OpenAICompatibleSchema mutated its input") + } + + root := decodeSchemaObject(t, got) + if root["$ref"] != "#/$defs/PlanEnvelope" || root["$schema"] == nil || root["$defs"] == nil { + t.Fatalf("root reference metadata was not preserved: %#v", root) + } + if _, exists := root["$id"]; exists { + t.Fatalf("root reference retained unsupported $id sibling: %#v", root) + } + defs := root["$defs"].(map[string]any) + planEnvelope := defs["PlanEnvelope"].(map[string]any) + plan := planEnvelope["properties"].(map[string]any)["plan"].(map[string]any) + if !reflect.DeepEqual(plan, map[string]any{"$ref": "#/$defs/PlanResult"}) { + t.Fatalf("plan reference = %#v, want a standalone $ref", plan) + } + summary := planEnvelope["properties"].(map[string]any)["summary"].(map[string]any) + if summary["description"] != "What the session did" { + t.Fatalf("ordinary property description was removed: %#v", summary) + } + for name, object := range map[string]map[string]any{ + "PlanEnvelope": planEnvelope, + "PlanResult": defs["PlanResult"].(map[string]any), + } { + if object["additionalProperties"] != false { + t.Errorf("%s additionalProperties = %v, want false", name, object["additionalProperties"]) + } + } +} + +func TestOpenAICompatibleSchema_RejectsSemanticReferenceSiblings(t *testing.T) { + raw := json.RawMessage(`{ + "type":"object", + "properties":{"value":{"$ref":"#/$defs/Value","minLength":1}}, + "$defs":{"Value":{"type":"string"}} + }`) + _, err := OpenAICompatibleSchema(raw) + if err == nil || !strings.Contains(err.Error(), `$.properties.value.$ref has unsupported sibling "minLength"`) { + t.Fatalf("error = %v, want path-aware unsupported sibling error", err) + } +} + +func TestOpenAICompatibleSchema_RejectsOpenEndedObjects(t *testing.T) { + type labels struct { + Values map[string]string `json:"values"` + } + _, err := SchemaJSONForBackend(api.BackendOpenAI, api.Prompt{Schema: &labels{}}) + if err == nil || !strings.Contains(err.Error(), "additionalProperties must be false") { + t.Fatalf("error = %v, want open-ended object rejection", err) + } +} + +func TestUsesOpenAISchemaSubset(t *testing.T) { + for _, backend := range []api.Backend{api.BackendOpenAI, api.BackendCodexCLI, api.BackendCodexAgent} { + if !UsesOpenAISchemaSubset(backend) { + t.Errorf("UsesOpenAISchemaSubset(%s) = false, want true", backend) + } + } + for _, backend := range []api.Backend{api.BackendAnthropic, api.BackendGemini, api.BackendCodexCmux} { + if UsesOpenAISchemaSubset(backend) { + t.Errorf("UsesOpenAISchemaSubset(%s) = true, want false", backend) + } + } +} + +func TestSchemaJSONForBackend_NativeTransformNotUsedForGemini(t *testing.T) { + raw := json.RawMessage(`{"type":"array","maxItems":2}`) + got, err := SchemaJSONForBackend(api.BackendGemini, api.Prompt{SchemaJSON: raw}) + if err != nil { + t.Fatalf("SchemaJSONForBackend: %v", err) + } + if string(got) != string(raw) { + t.Fatalf("gemini schema changed: got %s want %s", got, raw) + } +} + +func decodeSchemaObject(t *testing.T, raw json.RawMessage) map[string]any { + t.Helper() + var object map[string]any + if err := json.Unmarshal(raw, &object); err != nil { + t.Fatalf("decode schema: %v", err) + } + return object +} + +func schemaRequired(t *testing.T, object map[string]any) []string { + t.Helper() + raw, ok := object["required"].([]any) + if !ok { + t.Fatalf("required = %#v, want array", object["required"]) + } + out := make([]string, len(raw)) + for i, value := range raw { + var ok bool + out[i], ok = value.(string) + if !ok { + t.Fatalf("required[%d] = %#v, want string", i, value) + } + } + return out +} + +func containsJSONKey(t *testing.T, raw json.RawMessage, key string) bool { + t.Helper() + var v any + if err := json.Unmarshal(raw, &v); err != nil { + t.Fatalf("unmarshal: %v", err) + } + var walk func(any) bool + walk = func(node any) bool { + switch n := node.(type) { + case map[string]any: + if _, ok := n[key]; ok { + return true + } + for _, child := range n { + if walk(child) { + return true + } + } + case []any: + for _, child := range n { + if walk(child) { + return true + } + } + } + return false + } + return walk(v) +} + func TestGenerateJSONSchema(t *testing.T) { schema, err := GenerateJSONSchema(testStruct{}) if err != nil { diff --git a/pkg/ai/schema_validation.go b/pkg/ai/schema_validation.go new file mode 100644 index 0000000..d482e02 --- /dev/null +++ b/pkg/ai/schema_validation.go @@ -0,0 +1,32 @@ +package ai + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/xeipuuv/gojsonschema" +) + +// ValidateStructuredJSON returns joined schema violations, or a hard error when +// the schema or JSON document cannot be evaluated. +func ValidateStructuredJSON(schema json.RawMessage, document string) (string, error) { + if len(schema) == 0 { + return "", fmt.Errorf("%w: schema is required", ErrSchemaValidation) + } + if strings.TrimSpace(document) == "" { + return "response carried no JSON to validate", nil + } + result, err := gojsonschema.Validate(gojsonschema.NewBytesLoader(schema), gojsonschema.NewStringLoader(document)) + if err != nil { + return "", fmt.Errorf("%w: validation could not run: %v", ErrSchemaValidation, err) + } + if result.Valid() { + return "", nil + } + messages := make([]string, 0, len(result.Errors())) + for _, validationErr := range result.Errors() { + messages = append(messages, validationErr.String()) + } + return strings.Join(messages, "; "), nil +} diff --git a/pkg/ai/schema_validation_test.go b/pkg/ai/schema_validation_test.go new file mode 100644 index 0000000..1e26f41 --- /dev/null +++ b/pkg/ai/schema_validation_test.go @@ -0,0 +1,25 @@ +package ai + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateStructuredJSON(t *testing.T) { + schema := json.RawMessage(`{"type":"object","required":["pass"],"properties":{"pass":{"type":"boolean"}}}`) + + violations, err := ValidateStructuredJSON(schema, `{"pass":true}`) + require.NoError(t, err) + assert.Empty(t, violations) + + violations, err = ValidateStructuredJSON(schema, `{"pass":"yes"}`) + require.NoError(t, err) + assert.Contains(t, violations, "boolean") + + violations, err = ValidateStructuredJSON(schema, "") + require.NoError(t, err) + assert.Equal(t, "response carried no JSON to validate", violations) +} diff --git a/pkg/ai/session/session.go b/pkg/ai/session/session.go deleted file mode 100644 index 220213c..0000000 --- a/pkg/ai/session/session.go +++ /dev/null @@ -1,41 +0,0 @@ -package session - -import ( - "sync" - - "github.com/flanksource/captain/pkg/ai" -) - -type Session struct { - ID string - ProjectName string - Costs ai.Costs - mu sync.RWMutex -} - -func New(id, projectName string) *Session { - return &Session{ - ID: id, - ProjectName: projectName, - } -} - -func (s *Session) AddCost(cost ai.Cost) { - s.mu.Lock() - defer s.mu.Unlock() - s.Costs = append(s.Costs, cost) -} - -func (s *Session) TotalCost() float64 { - s.mu.RLock() - defer s.mu.RUnlock() - return s.Costs.Sum().Total() -} - -func (s *Session) GetCosts() ai.Costs { - s.mu.RLock() - defer s.mu.RUnlock() - result := make(ai.Costs, len(s.Costs)) - copy(result, s.Costs) - return result -} diff --git a/pkg/ai/terminal_outcome.go b/pkg/ai/terminal_outcome.go new file mode 100644 index 0000000..1319cf0 --- /dev/null +++ b/pkg/ai/terminal_outcome.go @@ -0,0 +1,187 @@ +package ai + +import ( + "fmt" + "strings" +) + +// TerminalOutcomeFromEvent normalizes supported native terminal tool events. +func TerminalOutcomeFromEvent(event Event) (*TerminalOutcome, error) { + if event.Kind != EventToolUse { + return nil, nil + } + switch event.Tool { + case "ExitPlanMode": + return terminalPlanFromInput(event.Input) + case "AskUserQuestion": + return terminalQuestionsFromInput(event.Input) + default: + return nil, nil + } +} + +// PlanTerminalPermission is the shared plan-mode permission policy: in a +// plan-only run the ExitPlanMode call is the turn's terminal signal, not a +// permission to grant — allowing it would leave plan mode and start +// implementation inside the same agent turn. Transports consult this before +// brokering a can_use_tool round-trip; when handled they apply the decision +// themselves (the SDK answers the deny, cmux dismisses its plan surface). +// AskUserQuestion is deliberately not handled: its round-trip is the +// interactive ask flow. +func PlanTerminalPermission(planMode bool, req PermissionRequest) (PermissionDecision, bool) { + if !planMode || req.Tool != "ExitPlanMode" { + return PermissionDecision{}, false + } + return PermissionDecision{ + Allow: false, + Message: "Plan captured for human review. Do not implement it in this session — end the turn and return your final response.", + }, true +} + +func terminalPlanFromInput(input map[string]any) (*TerminalOutcome, error) { + content, err := requiredString(input, "plan") + if err != nil { + return nil, fmt.Errorf("ExitPlanMode: %w", err) + } + path, err := optionalString(input, "planFilePath") + if err != nil { + return nil, fmt.Errorf("ExitPlanMode: %w", err) + } + outcome := &TerminalOutcome{ + Kind: TerminalOutcomePlan, + Plan: &TerminalPlan{Content: content, Path: path}, + } + return outcome, outcome.Validate() +} + +func terminalQuestionsFromInput(input map[string]any) (*TerminalOutcome, error) { + raw, ok := input["questions"] + if !ok { + question, err := terminalQuestionFromMap(input) + if err != nil { + return nil, fmt.Errorf("AskUserQuestion: %w", err) + } + outcome := &TerminalOutcome{Kind: TerminalOutcomeQuestions, Questions: []TerminalQuestion{question}} + return outcome, outcome.Validate() + } + items, ok := raw.([]any) + if !ok { + return nil, fmt.Errorf("AskUserQuestion: questions must be an array, got %T", raw) + } + questions := make([]TerminalQuestion, 0, len(items)) + for i, item := range items { + values, ok := item.(map[string]any) + if !ok { + return nil, fmt.Errorf("AskUserQuestion: question %d must be an object, got %T", i+1, item) + } + question, err := terminalQuestionFromMap(values) + if err != nil { + return nil, fmt.Errorf("AskUserQuestion: question %d: %w", i+1, err) + } + questions = append(questions, question) + } + outcome := &TerminalOutcome{Kind: TerminalOutcomeQuestions, Questions: questions} + return outcome, outcome.Validate() +} + +func terminalQuestionFromMap(values map[string]any) (TerminalQuestion, error) { + text, err := firstRequiredString(values, "question", "prompt", "text") + if err != nil { + return TerminalQuestion{}, err + } + context, err := firstOptionalString(values, "context", "header") + if err != nil { + return TerminalQuestion{}, err + } + options, err := terminalQuestionOptions(values["options"]) + if err != nil { + return TerminalQuestion{}, err + } + return TerminalQuestion{Text: text, Context: context, Options: options}, nil +} + +func terminalQuestionOptions(raw any) ([]string, error) { + if raw == nil { + return nil, nil + } + items, ok := raw.([]any) + if !ok { + return nil, fmt.Errorf("options must be an array, got %T", raw) + } + options := make([]string, 0, len(items)) + for i, item := range items { + var option string + var err error + switch value := item.(type) { + case string: + option = strings.TrimSpace(value) + case map[string]any: + option, err = firstRequiredString(value, "label", "value") + default: + err = fmt.Errorf("must be a string or object, got %T", item) + } + if err != nil || option == "" { + if err == nil { + err = fmt.Errorf("must not be empty") + } + return nil, fmt.Errorf("option %d %w", i+1, err) + } + options = append(options, option) + } + return options, nil +} + +func requiredString(values map[string]any, key string) (string, error) { + value, ok := values[key] + if !ok { + return "", fmt.Errorf("%s is required", key) + } + text, ok := value.(string) + if !ok { + return "", fmt.Errorf("%s must be a string, got %T", key, value) + } + if strings.TrimSpace(text) == "" { + return "", fmt.Errorf("%s is required", key) + } + return text, nil +} + +func optionalString(values map[string]any, key string) (string, error) { + value, ok := values[key] + if !ok || value == nil { + return "", nil + } + text, ok := value.(string) + if !ok { + return "", fmt.Errorf("%s must be a string, got %T", key, value) + } + return text, nil +} + +func firstRequiredString(values map[string]any, keys ...string) (string, error) { + text, err := firstOptionalString(values, keys...) + if err != nil { + return "", err + } + if text == "" { + return "", fmt.Errorf("one of %s is required", strings.Join(keys, ", ")) + } + return text, nil +} + +func firstOptionalString(values map[string]any, keys ...string) (string, error) { + for _, key := range keys { + value, ok := values[key] + if !ok || value == nil { + continue + } + text, ok := value.(string) + if !ok { + return "", fmt.Errorf("%s must be a string, got %T", key, value) + } + if strings.TrimSpace(text) != "" { + return text, nil + } + } + return "", nil +} diff --git a/pkg/ai/terminal_outcome_test.go b/pkg/ai/terminal_outcome_test.go new file mode 100644 index 0000000..a56f1c5 --- /dev/null +++ b/pkg/ai/terminal_outcome_test.go @@ -0,0 +1,100 @@ +package ai + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTerminalOutcomeFromEventPlan(t *testing.T) { + outcome, err := TerminalOutcomeFromEvent(Event{ + Kind: EventToolUse, + Tool: "ExitPlanMode", + Input: map[string]any{ + "plan": "1. Inspect the seam\n2. Implement it", + "planFilePath": "/repo/.claude/plans/example.md", + }, + }) + + require.NoError(t, err) + require.NotNil(t, outcome) + assert.Equal(t, TerminalOutcomePlan, outcome.Kind) + require.NotNil(t, outcome.Plan) + assert.Equal(t, "1. Inspect the seam\n2. Implement it", outcome.Plan.Content) + assert.Equal(t, "/repo/.claude/plans/example.md", outcome.Plan.Path) +} + +func TestTerminalOutcomeFromEventQuestions(t *testing.T) { + outcome, err := TerminalOutcomeFromEvent(Event{ + Kind: EventToolUse, + Tool: "AskUserQuestion", + Input: map[string]any{"questions": []any{ + map[string]any{ + "question": "Which database?", + "header": "Storage", + "options": []any{ + map[string]any{"label": "PostgreSQL", "description": "Production database"}, + "SQLite", + }, + }, + }}, + }) + + require.NoError(t, err) + require.NotNil(t, outcome) + assert.Equal(t, TerminalOutcomeQuestions, outcome.Kind) + require.Equal(t, []TerminalQuestion{{ + Text: "Which database?", + Context: "Storage", + Options: []string{"PostgreSQL", "SQLite"}, + }}, outcome.Questions) +} + +func TestTerminalOutcomeFromEventFailsOnMalformedNativeOutcome(t *testing.T) { + tests := []struct { + name string + event Event + want string + }{ + { + name: "missing plan", + event: Event{Kind: EventToolUse, Tool: "ExitPlanMode", Input: map[string]any{"planFilePath": "/repo/plan.md"}}, + want: "plan is required", + }, + { + name: "invalid option", + event: Event{Kind: EventToolUse, Tool: "AskUserQuestion", Input: map[string]any{ + "questions": []any{map[string]any{"question": "Continue?", "options": []any{42}}}, + }}, + want: "option 1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + outcome, err := TerminalOutcomeFromEvent(tt.event) + assert.Nil(t, outcome) + require.ErrorContains(t, err, tt.want) + }) + } +} + +func TestTerminalOutcomeFromEventIgnoresOrdinaryTools(t *testing.T) { + outcome, err := TerminalOutcomeFromEvent(Event{Kind: EventToolUse, Tool: "Read", Input: map[string]any{"file_path": "README.md"}}) + require.NoError(t, err) + assert.Nil(t, outcome) +} + +func TestPlanTerminalPermission(t *testing.T) { + decision, handled := PlanTerminalPermission(true, PermissionRequest{Tool: "ExitPlanMode", Input: map[string]any{"plan": "1. do it"}}) + require.True(t, handled, "ExitPlanMode in plan mode is the terminal signal, not a brokered approval") + assert.False(t, decision.Allow) + assert.NotEmpty(t, decision.Message) + + _, handled = PlanTerminalPermission(false, PermissionRequest{Tool: "ExitPlanMode"}) + assert.False(t, handled, "outside plan mode ExitPlanMode brokers normally") + + _, handled = PlanTerminalPermission(true, PermissionRequest{Tool: "AskUserQuestion"}) + assert.False(t, handled, "AskUserQuestion keeps its interactive broker round-trip") +} diff --git a/pkg/ai/tools/catalog.go b/pkg/ai/tools/catalog.go new file mode 100644 index 0000000..2ebd558 --- /dev/null +++ b/pkg/ai/tools/catalog.go @@ -0,0 +1,141 @@ +package tools + +import "strings" + +// ToolCatalog is the GET /api/chat/tools payload. +type ToolCatalog struct { + Tools []ToolCatalogEntry `json:"tools"` +} + +// ToolCatalogEntry is the frontend-facing DTO for one tool. Unlike ToolInfo it +// keeps method/path/operationName as typed fields, since the tool-preferences UI +// renders them. +type ToolCatalogEntry struct { + Name string `json:"name"` + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` + Source string `json:"source"` + Server string `json:"server,omitempty"` + Group string `json:"group,omitempty"` + Parent string `json:"parent,omitempty"` + Icon string `json:"icon,omitempty"` + PreferenceKey string `json:"preferenceKey,omitempty"` + DefaultPermission ToolMode `json:"defaultPermission,omitempty"` + Strict *bool `json:"strict,omitempty"` + Method string `json:"method,omitempty"` + Path string `json:"path,omitempty"` + OperationName string `json:"operationName,omitempty"` + InputSchema map[string]any `json:"inputSchema"` + OutputSchema map[string]any `json:"outputSchema,omitempty"` +} + +// CustomCatalogEntry builds the catalog DTO for an app-owned custom tool. The +// method/path/operationName it surfaces come from the definition's Annotations +// (clicky/method, clicky/path, clicky/operation) when present. +func CustomCatalogEntry(def ToolDefinition, name string, schema map[string]any) ToolCatalogEntry { + info := ToolInfo{ + Name: name, + Group: def.Group, + Parent: def.Parent, + Icon: def.Icon, + DefaultPermission: DefaultPermissionMode(def.DefaultPermission), + Strict: def.Strict, + Annotations: def.Annotations, + } + return ToolCatalogEntry{ + Name: name, + Title: def.Name, + Description: def.Description, + Source: "custom", + Group: def.Group, + Parent: def.Parent, + Icon: def.Icon, + PreferenceKey: PreferenceKey(info), + DefaultPermission: DefaultPermissionMode(def.DefaultPermission), + Strict: def.Strict, + Method: def.Annotations["clicky/method"], + Path: def.Annotations["clicky/path"], + OperationName: def.Name, + InputSchema: ObjectSchema(schema), + } +} + +// PreferenceKey returns the key a tool is governed by in the preferences UI: its +// group when grouped, otherwise its own name. +func PreferenceKey(info ToolInfo) string { + if info.Group != "" { + return info.Group + } + return info.Name +} + +// ObjectSchema defaults a nil/typeless schema to an empty JSON object schema. +func ObjectSchema(schema map[string]any) map[string]any { + if schema == nil { + return map[string]any{"type": "object", "properties": map[string]any{}} + } + if _, ok := schema["type"]; !ok { + schema["type"] = "object" + } + if schema["type"] == "object" { + if _, ok := schema["properties"]; !ok { + schema["properties"] = map[string]any{} + } + } + return schema +} + +// StringMetadata returns the first non-empty string value among keys. +func StringMetadata(meta map[string]any, keys ...string) (string, bool) { + for _, key := range keys { + if v, ok := meta[key].(string); ok && v != "" { + return v, true + } + } + return "", false +} + +// ApplyToolMetadata overlays MCP tool metadata (group/parent/icon/permission/ +// strict, incl. the nested com.flanksource.clicky/tool block) onto an entry. +func ApplyToolMetadata(entry *ToolCatalogEntry, meta map[string]any) { + if entry == nil || len(meta) == 0 { + return + } + if group, ok := StringMetadata(meta, "group", "toolGroup"); ok && entry.Group == "" { + entry.Group = group + entry.PreferenceKey = group + } + if parent, ok := StringMetadata(meta, "parent", "toolParent"); ok && entry.Parent == "" { + entry.Parent = parent + } + if icon, ok := StringMetadata(meta, "icon", "toolIcon"); ok && entry.Icon == "" { + entry.Icon = icon + } + if permission, ok := StringMetadata(meta, "defaultPermission", "permission", "defaultMode"); ok { + entry.DefaultPermission = DefaultPermissionMode(ToolMode(permission)) + } + if strict, ok := BoolMetadata(meta, "strict"); ok && entry.Strict == nil { + entry.Strict = &strict + } + if nested, ok := meta["com.flanksource.clicky/tool"].(map[string]any); ok { + ApplyToolMetadata(entry, nested) + } +} + +// BoolMetadata returns the first bool-valued (or "true"/"false" string) key. +func BoolMetadata(meta map[string]any, keys ...string) (bool, bool) { + for _, key := range keys { + switch v := meta[key].(type) { + case bool: + return v, true + case string: + switch strings.ToLower(strings.TrimSpace(v)) { + case "true": + return true, true + case "false": + return false, true + } + } + } + return false, false +} diff --git a/pkg/ai/tools/preferences_ginkgo_test.go b/pkg/ai/tools/preferences_ginkgo_test.go new file mode 100644 index 0000000..97850e1 --- /dev/null +++ b/pkg/ai/tools/preferences_ginkgo_test.go @@ -0,0 +1,36 @@ +package tools_test + +import ( + "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Tool preference resolution", func() { + It("prefers an exact tool entry over its group", func() { + mode, ok := tools.EffectivePreference(api.ToolPreferences{ + "billing": api.ToolModeOff, + "invoice_delete": api.ToolModeAsk, + }, tools.ToolInfo{Name: "invoice_delete", Group: "billing"}) + + Expect(ok).To(BeTrue()) + Expect(mode).To(Equal(api.ToolModeAsk)) + }) + + It("normalizes only the canonical modes", func() { + on, ok := tools.NormalizedPreference(api.ToolPreferences{"search": api.ToolModeOn}, "search") + Expect(ok).To(BeTrue()) + Expect(on).To(Equal(api.ToolModeOn)) + + off, ok := tools.NormalizedPreference(api.ToolPreferences{"search": api.ToolModeOff}, "search") + Expect(ok).To(BeTrue()) + Expect(off).To(Equal(api.ToolModeOff)) + + _, ok = tools.NormalizedPreference(api.ToolPreferences{"search": "enabled"}, "search") + Expect(ok).To(BeFalse()) + _, ok = tools.NormalizedPreference(api.ToolPreferences{"search": "disabled"}, "search") + Expect(ok).To(BeFalse()) + }) +}) diff --git a/pkg/ai/tools/tools.go b/pkg/ai/tools/tools.go new file mode 100644 index 0000000..530ca8b --- /dev/null +++ b/pkg/ai/tools/tools.go @@ -0,0 +1,257 @@ +// Package tools is captain's home for the chat tool registry's data model and +// approval policy: the tool definition/info types, the per-request tool mode and +// preferences, the tool catalog DTO, and the approval-decision logic. It is +// genkit- and clicky-free — the genkit binding (registering these as model tools) +// and the clicky-RPC→tool mapping live in the consumer (clicky/aichat), which +// imports this package. Clicky-specific metadata (the originating verb/method/ +// path) rides in the opaque Annotations map rather than as typed fields. +package tools + +import ( + "context" + "sort" + + "github.com/flanksource/captain/pkg/api" +) + +// ToolMode controls how a tool is exposed for one request. +type ToolMode = api.ToolMode + +const ( + ToolModeOn = api.ToolModeOn + ToolModeAsk = api.ToolModeAsk + ToolModeOff = api.ToolModeOff + ToolModeAuto = api.ToolModeAuto +) + +// NormalizeToolMode canonicalizes a mode string. The bool is false for an +// unrecognized value. +func NormalizeToolMode(mode ToolMode) (ToolMode, bool) { + return api.NormalizeToolMode(mode) +} + +// DefaultPermissionMode resolves a mode to its canonical value, defaulting an +// unset/unknown value to Auto (defer to the approval policy). +func DefaultPermissionMode(mode ToolMode) ToolMode { + if normalized, ok := NormalizeToolMode(mode); ok { + return normalized + } + return ToolModeAuto +} + +// ApprovalDecisionForMode maps a resolved mode to an approve/auto decision. The +// second bool is false only for Auto, which defers to the policy. +func ApprovalDecisionForMode(mode ToolMode) (require bool, handled bool) { + switch DefaultPermissionMode(mode) { + case ToolModeOn: + return false, true + case ToolModeAsk: + return true, true + case ToolModeOff: + return false, true + case ToolModeAuto: + return false, false + default: + return false, false + } +} + +// ToolPreferences carries the clicky-ui tool preference payload. The UI sends +// "on", "ask", "off", or "auto". +type ToolPreferences = api.ToolPreferences + +// ToolInfo is the concrete tool being considered for approval and preference +// resolution. Clicky-RPC specifics (verb/method/path/operation) live in +// Annotations, not as typed fields. +type ToolInfo struct { + Name string + // Group is the tool-group this tool belongs to. When non-empty the + // preferences UI presents the group as one entry governing every member. + Group string + Parent string + Icon string + DefaultPermission ToolMode + Strict *bool + ReadOnlyHint *bool + DestructiveHint *bool + IdempotentHint *bool + // Annotations carries opaque caller metadata (e.g. clicky/verb, clicky/method, + // clicky/path, clicky/operation) for policies that want the raw values. + Annotations map[string]string +} + +// Annotation returns the named annotation (empty when absent). +func (i ToolInfo) Annotation(key string) string { + if i.Annotations == nil { + return "" + } + return i.Annotations[key] +} + +// ToolDefinition describes an app-owned tool registered alongside clicky RPC and +// MCP tools. Handlers should return JSON-serializable values. +type ToolDefinition struct { + Name string + Description string + InputSchema map[string]any + Parent string + Icon string + DefaultPermission ToolMode + Strict *bool + ReadOnlyHint *bool + DestructiveHint *bool + IdempotentHint *bool + // Group places this custom tool in a tool-group so the preferences UI presents + // it under the group rather than individually. + Group string + // Annotations carries opaque caller metadata (see ToolInfo.Annotations). + Annotations map[string]string + Handler func(context.Context, any) (any, error) +} + +// ApprovalPolicy reports whether a tool call must be approved before it runs. +type ApprovalPolicy func(toolName string, input any) bool + +// ToolApprovalPolicy is the metadata-aware approval hook; it takes precedence +// over ApprovalPolicy when both are configured. +type ToolApprovalPolicy func(tool ToolInfo, input any) bool + +// ApprovalPredicate is the internal gate type shared by the resolvers. +type ApprovalPredicate func(tool ToolInfo, input any) bool + +// ResolveApprovalPolicy picks the effective gate: an explicit tool policy wins, +// then a name-based policy, then an exact-name list, else nil (auto-approve). +func ResolveApprovalPolicy(toolPolicy ToolApprovalPolicy, policy ApprovalPolicy, names []string) ApprovalPredicate { + if toolPolicy != nil { + return ApprovalPredicate(toolPolicy) + } + if policy != nil { + return func(tool ToolInfo, input any) bool { + return policy(tool.Name, input) + } + } + return RequireApprovalFor(names) +} + +// RequireApprovalFor builds a predicate requiring approval for exactly the named +// tools. An empty list yields nil (auto-approve everything). +func RequireApprovalFor(names []string) ApprovalPredicate { + if len(names) == 0 { + return nil + } + set := make(map[string]bool, len(names)) + for _, n := range names { + set[n] = true + } + return func(tool ToolInfo, _ any) bool { + return set[tool.Name] + } +} + +type toolRuntimeConfig struct { + preferences ToolPreferences + defaultApproval ApprovalPredicate +} + +type toolRuntimeContextKey struct{} + +// WithRuntime stores the per-request tool preferences + default approval gate on +// the context so a tool handler can resolve its approval decision. +func WithRuntime(ctx context.Context, prefs ToolPreferences, defaultApproval ApprovalPredicate) context.Context { + return context.WithValue(ctx, toolRuntimeContextKey{}, toolRuntimeConfig{preferences: prefs, defaultApproval: defaultApproval}) +} + +func runtimeConfig(ctx context.Context) (toolRuntimeConfig, bool) { + cfg, ok := ctx.Value(toolRuntimeContextKey{}).(toolRuntimeConfig) + return cfg, ok +} + +// ShouldRequireApproval resolves whether a tool call must be approved, honoring +// (in order) the per-request preference, the tool's default permission, the +// runtime's default approval gate, and finally the supplied fallback. +func ShouldRequireApproval(ctx context.Context, fallback ApprovalPredicate, tool ToolInfo, input any) bool { + if ctx != nil { + if cfg, ok := runtimeConfig(ctx); ok { + if mode, ok := EffectivePreference(cfg.preferences, tool); ok { + if decision, handled := ApprovalDecisionForMode(mode); handled { + return decision + } + } + if decision, handled := ApprovalDecisionForMode(DefaultPermissionMode(tool.DefaultPermission)); handled { + return decision + } + if cfg.defaultApproval != nil { + return cfg.defaultApproval(tool, input) + } + } + } + if decision, handled := ApprovalDecisionForMode(DefaultPermissionMode(tool.DefaultPermission)); handled { + return decision + } + if fallback == nil { + return false + } + return fallback(tool, input) +} + +// EffectivePreference resolves the ToolMode for a tool: an exact tool-name +// preference wins, else the tool's group preference; ungrouped tools resolve by +// their own name. +func EffectivePreference(prefs ToolPreferences, info ToolInfo) (ToolMode, bool) { + if mode, ok := NormalizedPreference(prefs, info.Name); ok { + return mode, true + } + if info.Group != "" { + return NormalizedPreference(prefs, info.Group) + } + return "", false +} + +// NormalizedPreference looks up and normalizes a preference by key. +func NormalizedPreference(prefs ToolPreferences, name string) (ToolMode, bool) { + if len(prefs) == 0 { + return "", false + } + mode, ok := prefs[name] + if !ok { + return "", false + } + return NormalizeToolMode(mode) +} + +// ToolEntry is one row in the tool-preferences UI: a single ungrouped tool, or a +// collapsed group listing its member names. +type ToolEntry struct { + Key string `json:"key"` + Group string `json:"group,omitempty"` + Tools []string `json:"tools"` + Mode ToolMode `json:"mode,omitempty"` +} + +// ListToolEntries collapses grouped tools into one entry per group and leaves +// ungrouped tools individual, sorted by Key. prefs (may be nil) annotates Mode. +func ListToolEntries(infos []ToolInfo, prefs ToolPreferences) []ToolEntry { + groups := map[string][]string{} + var entries []ToolEntry + for _, info := range infos { + if g := info.Group; g != "" { + groups[g] = append(groups[g], info.Name) + continue + } + entry := ToolEntry{Key: info.Name, Tools: []string{info.Name}} + if mode, ok := NormalizedPreference(prefs, info.Name); ok { + entry.Mode = mode + } + entries = append(entries, entry) + } + for group, members := range groups { + sort.Strings(members) + entry := ToolEntry{Key: group, Group: group, Tools: members} + if mode, ok := NormalizedPreference(prefs, group); ok { + entry.Mode = mode + } + entries = append(entries, entry) + } + sort.Slice(entries, func(i, j int) bool { return entries[i].Key < entries[j].Key }) + return entries +} diff --git a/pkg/ai/tools/tools_suite_test.go b/pkg/ai/tools/tools_suite_test.go new file mode 100644 index 0000000..42a1dde --- /dev/null +++ b/pkg/ai/tools/tools_suite_test.go @@ -0,0 +1,13 @@ +package tools_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestTools(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "AI Tools Suite") +} diff --git a/pkg/ai/types.go b/pkg/ai/types.go index 67ab6a9..c9e9977 100644 --- a/pkg/ai/types.go +++ b/pkg/ai/types.go @@ -18,13 +18,22 @@ type Request = api.Spec // contract). They are re-exported here as aliases so existing call sites and // clicky/aichat's captainai.* keep compiling unchanged. type ( - PermissionFunc = api.PermissionFunc - PermissionRequest = api.PermissionRequest - PermissionDecision = api.PermissionDecision - Response = api.Response - EventKind = api.EventKind - Event = api.Event - Config = api.Config + PermissionFunc = api.PermissionFunc + PermissionRequest = api.PermissionRequest + PermissionDecision = api.PermissionDecision + Response = api.Response + TerminalOutcome = api.TerminalOutcome + TerminalOutcomeKind = api.TerminalOutcomeKind + TerminalPlan = api.TerminalPlan + TerminalQuestion = api.TerminalQuestion + EventKind = api.EventKind + Event = api.Event + Config = api.Config +) + +const ( + TerminalOutcomePlan = api.TerminalOutcomePlan + TerminalOutcomeQuestions = api.TerminalOutcomeQuestions ) const ( @@ -41,6 +50,14 @@ const ( // Usage is an alias for the canonical api.Usage (per-call token breakdown). type Usage = api.Usage +// NetInputTokens / NetOutputTokens re-export the disjoint-bucket normalizers so +// provider packages (which import ai, not api) can enforce the Usage invariant +// at their parse boundary. +var ( + NetInputTokens = api.NetInputTokens + NetOutputTokens = api.NetOutputTokens +) + // Cost and Costs are aliases for the canonical api types (token + money // accounting). The methods (Total/Add/Sum/ByModel) live on the api types. type Cost = api.Cost diff --git a/pkg/aichat/aichat_suite_test.go b/pkg/aichat/aichat_suite_test.go new file mode 100644 index 0000000..f0547fe --- /dev/null +++ b/pkg/aichat/aichat_suite_test.go @@ -0,0 +1,13 @@ +package aichat_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAIChat(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "AI Chat Transport Suite") +} diff --git a/pkg/aichat/events.go b/pkg/aichat/events.go new file mode 100644 index 0000000..b87cff6 --- /dev/null +++ b/pkg/aichat/events.go @@ -0,0 +1,362 @@ +package aichat + +import ( + "encoding/json" + "fmt" + "reflect" + "sort" + + "github.com/flanksource/captain/pkg/api" +) + +type toolState struct { + name string + input map[string]any + approvalRequested bool +} + +type eventStream struct { + writer *SSEWriter + blockType string + blockID string + nextBlock int + tools map[string]toolState + metadata *MessageMetadata + sessionID string + model string + terminal bool + finished bool +} + +// WriteEventStream translates a Captain event channel into one complete UI Message Stream. +// Translation errors are written as an error part before the stream is closed and are +// also returned so malformed provider output cannot pass silently. +func WriteEventStream(writer *SSEWriter, events <-chan api.Event) error { + stream := &eventStream{writer: writer, tools: map[string]toolState{}} + if err := stream.start(); err != nil { + return err + } + for event := range events { + if err := stream.event(event); err != nil { + return stream.fail(err) + } + } + if err := stream.finish(); err != nil { + return stream.fail(err) + } + return nil +} + +func (s *eventStream) start() error { + if err := s.writer.WritePart(Part{Type: "start"}); err != nil { + return err + } + return s.writer.WritePart(Part{Type: "start-step"}) +} + +func (s *eventStream) event(event api.Event) error { + if s.terminal { + return fmt.Errorf("event %q received after terminal event", event.Kind) + } + if event.SessionID != "" { + s.sessionID = event.SessionID + } + if event.Model != "" { + s.model = event.Model + } + switch event.Kind { + case api.EventText: + return s.delta("text", event.Text) + case api.EventThinking: + return s.delta("reasoning", event.Text) + case api.EventToolUse: + return s.toolUse(event) + case api.EventPermission: + return s.permission(event) + case api.EventToolResult: + return s.toolResult(event) + case api.EventSystem: + return nil + case api.EventResult: + return s.result(event) + case api.EventError: + return s.providerError(event) + default: + return fmt.Errorf("unsupported Captain event kind %q", event.Kind) + } +} + +func (s *eventStream) delta(kind, delta string) error { + if delta == "" { + return nil + } + if s.blockType != kind { + if err := s.closeBlock(); err != nil { + return err + } + s.blockType = kind + s.blockID = fmt.Sprintf("%s-%d", kind, s.nextBlock) + s.nextBlock++ + if err := s.writer.WritePart(Part{Type: kind + "-start", ID: s.blockID}); err != nil { + return err + } + } + return s.writer.WritePart(Part{Type: kind + "-delta", ID: s.blockID, Delta: delta}) +} + +func (s *eventStream) toolUse(event api.Event) error { + if err := s.closeBlock(); err != nil { + return err + } + if event.Tool == "" { + return fmt.Errorf("tool use has no tool name") + } + if event.ToolCallID == "" { + return fmt.Errorf("tool use %q has no tool call id", event.Tool) + } + if _, exists := s.tools[event.ToolCallID]; exists { + return fmt.Errorf("duplicate tool call id %q", event.ToolCallID) + } + input := event.Input + if input == nil { + input = map[string]any{} + } + if err := s.writer.WritePart(Part{ + Type: "tool-input-available", ToolCallID: event.ToolCallID, + ToolName: event.Tool, Input: input, Dynamic: true, + }); err != nil { + return err + } + s.tools[event.ToolCallID] = toolState{name: event.Tool, input: event.Input} + return nil +} + +func (s *eventStream) permission(event api.Event) error { + if err := s.closeBlock(); err != nil { + return err + } + state, err := s.correlatedTool("permission", event) + if err != nil { + return err + } + if state.approvalRequested { + return fmt.Errorf("duplicate permission for tool call %q", event.ToolCallID) + } + state.approvalRequested = true + s.tools[event.ToolCallID] = state + return s.writer.WritePart(Part{ + Type: "tool-approval-request", ApprovalID: event.ToolCallID, ToolCallID: event.ToolCallID, + }) +} + +func (s *eventStream) toolResult(event api.Event) error { + if err := s.closeBlock(); err != nil { + return err + } + if _, err := s.correlatedTool("result", event); err != nil { + return err + } + output := map[string]any{"output": event.Text} + if !event.Success { + output["isError"] = true + } + if err := s.writer.WritePart(Part{Type: "tool-output-available", ToolCallID: event.ToolCallID, Output: output}); err != nil { + return err + } + delete(s.tools, event.ToolCallID) + return nil +} + +func (s *eventStream) correlatedTool(kind string, event api.Event) (toolState, error) { + if event.ToolCallID == "" { + return toolState{}, fmt.Errorf("%s for tool %q has no tool call id", kind, event.Tool) + } + state, exists := s.tools[event.ToolCallID] + if !exists { + return toolState{}, fmt.Errorf("%s for tool call %q has no matching tool use", kind, event.ToolCallID) + } + if event.Tool != "" && event.Tool != state.name { + return toolState{}, fmt.Errorf("%s for tool call %q names %q, want %q", kind, event.ToolCallID, event.Tool, state.name) + } + return state, nil +} + +func (s *eventStream) result(event api.Event) error { + if err := s.closeBlock(); err != nil { + return err + } + if err := s.unresolvedToolError(event.ToolApproval != nil); err != nil { + return err + } + if event.ToolApproval != nil { + if err := event.ToolApproval.Validate(); err != nil { + return fmt.Errorf("validate tool approval state: %w", err) + } + if err := s.validateApprovalCorrelation(event.ToolApproval); err != nil { + return err + } + if err := s.writer.WritePart(Part{Type: "data-tool-approval", Data: event.ToolApproval}); err != nil { + return err + } + } else { + data := any(map[string]any{"success": event.Success}) + if len(event.StructuredData) > 0 { + if err := json.Unmarshal(event.StructuredData, &data); err != nil { + return fmt.Errorf("decode structured result: %w", err) + } + } + if err := s.writer.WritePart(Part{Type: "data-result", Data: data}); err != nil { + return err + } + } + success := event.Success + s.metadata = &MessageMetadata{ + ProviderSessionID: s.sessionID, + Model: s.model, + Cost: event.CostUSD, + Success: &success, + } + if event.Usage != nil { + s.metadata.Usage = usageMetadata(*event.Usage) + s.metadata.ContextTokens = event.Usage.InputTokens + } + s.terminal = true + return nil +} + +func (s *eventStream) validateApprovalCorrelation(approval *api.ToolApprovalState) error { + seen := make(map[string]bool) + pending := approval.Pending() + sort.Slice(pending, func(i, j int) bool { return pending[i].ToolCallID < pending[j].ToolCallID }) + for _, request := range pending { + state, ok := s.tools[request.ToolCallID] + if !ok || !state.approvalRequested { + return fmt.Errorf("approval state call %q has no matching streamed approval request", request.ToolCallID) + } + if request.Tool != state.name { + return fmt.Errorf("approval state call %q names %q, want %q", request.ToolCallID, request.Tool, state.name) + } + if !approvalInputMatches(request.Input, state.input) { + return fmt.Errorf("approval state call %q input does not match the streamed tool request", request.ToolCallID) + } + seen[request.ToolCallID] = true + } + streamed := make([]string, 0, len(s.tools)) + for id, state := range s.tools { + if state.approvalRequested && !seen[id] { + streamed = append(streamed, id) + } + } + if len(streamed) > 0 { + sort.Strings(streamed) + return fmt.Errorf("streamed approval request %q is absent from the approval state", streamed[0]) + } + return nil +} + +func approvalInputMatches(raw json.RawMessage, input map[string]any) bool { + if len(raw) == 0 { + return len(input) == 0 + } + var stateInput any + if json.Unmarshal(raw, &stateInput) != nil { + return false + } + streamedRaw, err := json.Marshal(input) + if err != nil { + return false + } + var streamedInput any + if json.Unmarshal(streamedRaw, &streamedInput) != nil { + return false + } + return reflect.DeepEqual(stateInput, streamedInput) +} + +func usageMetadata(usage api.Usage) *UsageMetadata { + return &UsageMetadata{ + InputTokens: usage.InputTokens, OutputTokens: usage.OutputTokens, + ReasoningTokens: usage.ReasoningTokens, CacheReadTokens: usage.CacheReadTokens, + CacheWriteTokens: usage.CacheWriteTokens, TotalTokens: usage.TotalTokens(), + } +} + +func (s *eventStream) providerError(event api.Event) error { + if err := s.closeBlock(); err != nil { + return err + } + if event.Error == "" { + return fmt.Errorf("captain error event has no error message") + } + s.tools = map[string]toolState{} + s.terminal = true + return s.writer.WritePart(Part{Type: "error", ErrorText: event.Error}) +} + +func (s *eventStream) closeBlock() error { + if s.blockType == "" { + return nil + } + err := s.writer.WritePart(Part{Type: s.blockType + "-end", ID: s.blockID}) + s.blockType = "" + s.blockID = "" + return err +} + +func (s *eventStream) unresolvedToolError(allowApproval bool) error { + ids := make([]string, 0, len(s.tools)) + for id, state := range s.tools { + if allowApproval && state.approvalRequested { + continue + } + ids = append(ids, id) + } + if len(ids) == 0 { + return nil + } + sort.Strings(ids) + if !allowApproval { + return fmt.Errorf("tool call %q ended without a result", ids[0]) + } + return fmt.Errorf("tool call %q ended without a result or approval request", ids[0]) +} + +func (s *eventStream) finish() error { + if s.finished { + return fmt.Errorf("AI SDK event stream is already finished") + } + if err := s.closeBlock(); err != nil { + return err + } + if !s.terminal { + if err := s.unresolvedToolError(true); err != nil { + return err + } + } + if err := s.writer.WritePart(Part{Type: "finish-step"}); err != nil { + return err + } + if err := s.writer.WritePart(Part{Type: "finish", MessageMetadata: s.metadata}); err != nil { + return err + } + s.finished = true + return s.writer.Done() +} + +func (s *eventStream) fail(cause error) error { + if s.finished { + return cause + } + if err := s.closeBlock(); err != nil { + return err + } + if err := s.writer.WritePart(Part{Type: "error", ErrorText: cause.Error()}); err != nil { + return err + } + s.terminal = true + s.tools = map[string]toolState{} + if err := s.finish(); err != nil { + return err + } + return cause +} diff --git a/pkg/aichat/mcp_provider.go b/pkg/aichat/mcp_provider.go new file mode 100644 index 0000000..8864122 --- /dev/null +++ b/pkg/aichat/mcp_provider.go @@ -0,0 +1,302 @@ +package aichat + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "maps" + "sync" + + mcpclient "github.com/mark3labs/mcp-go/client" + "github.com/mark3labs/mcp-go/client/transport" + "github.com/mark3labs/mcp-go/mcp" + + aitools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" +) + +// Identity sent in the MCP initialize handshake. Servers log it, and some gate +// their capabilities on the client they are talking to. +const ( + mcpClientName = "captain" + mcpClientVersion = "1.0.0" +) + +// MCPServer configures one external MCP server. Exactly one transport must be +// set: Command for stdio, or URL for SSE/streamable HTTP. +type MCPServer struct { + Name string + Command string + Args []string + Env []string + URL string + Headers map[string]string + StreamableHTTP bool +} + +type MCPToolProviderOptions struct { + Servers []MCPServer +} + +// mcpClient is the slice of an MCP session this provider needs: list the +// server's tools, invoke one, and shut the transport down. +type mcpClient interface { + ListTools(context.Context) ([]mcp.Tool, error) + CallTool(ctx context.Context, name string, arguments map[string]any) (any, error) + Disconnect() error +} + +type mcpClientFactory func(context.Context, MCPServer) (mcpClient, error) + +// MCPToolProvider owns explicit MCP client sessions and their projected Captain +// tool definitions. Call Close when the mounted chat service shuts down. +type MCPToolProvider struct { + options MCPToolProviderOptions + clientFactory mcpClientFactory + + mu sync.Mutex + clients map[string]mcpClient + tools *ToolSet +} + +func NewMCPToolProvider(options MCPToolProviderOptions) *MCPToolProvider { + options.Servers = append([]MCPServer(nil), options.Servers...) + return &MCPToolProvider{ + options: options, + clientFactory: dialMCPServer, + clients: map[string]mcpClient{}, + } +} + +func (p *MCPToolProvider) ToolSet(ctx context.Context) (ToolSet, error) { + p.mu.Lock() + defer p.mu.Unlock() + if p.tools != nil { + return cloneToolSet(*p.tools), nil + } + if err := p.validate(); err != nil { + return ToolSet{}, err + } + + set, err := p.load(ctx) + if err != nil { + return ToolSet{}, errors.Join(err, p.closeClientsLocked()) + } + p.tools = &set + return cloneToolSet(set), nil +} + +func (p *MCPToolProvider) Close() error { + p.mu.Lock() + defer p.mu.Unlock() + p.tools = nil + return p.closeClientsLocked() +} + +func (p *MCPToolProvider) validate() error { + seen := make(map[string]bool, len(p.options.Servers)) + for _, server := range p.options.Servers { + if server.Name == "" { + return fmt.Errorf("MCP server name is required") + } + if seen[server.Name] { + return fmt.Errorf("duplicate MCP server %q", server.Name) + } + seen[server.Name] = true + if (server.Command == "") == (server.URL == "") { + return fmt.Errorf("MCP server %q must configure exactly one of Command or URL", server.Name) + } + } + return nil +} + +func (p *MCPToolProvider) load(ctx context.Context) (ToolSet, error) { + set := ToolSet{} + seen := make(map[string]string) + for _, server := range p.options.Servers { + client, err := p.client(ctx, server) + if err != nil { + return ToolSet{}, fmt.Errorf("connect MCP server %q: %w", server.Name, err) + } + tools, err := client.ListTools(ctx) + if err != nil { + return ToolSet{}, fmt.Errorf("discover MCP tools from server %q: %w", server.Name, err) + } + for _, tool := range tools { + name := namespacedToolName(server.Name, tool.Name) + if prior, exists := seen[name]; exists { + return ToolSet{}, fmt.Errorf("duplicate MCP tool definition %q from servers %q and %q", name, prior, server.Name) + } + seen[name] = server.Name + definition, catalog, err := projectMCPTool(server.Name, tool, client) + if err != nil { + return ToolSet{}, err + } + set.Definitions = append(set.Definitions, definition) + set.Catalog = append(set.Catalog, catalog) + } + } + return set, nil +} + +func (p *MCPToolProvider) client(ctx context.Context, server MCPServer) (mcpClient, error) { + if client := p.clients[server.Name]; client != nil { + return client, nil + } + client, err := p.clientFactory(ctx, server) + if err != nil { + return nil, err + } + p.clients[server.Name] = client + return client, nil +} + +func (p *MCPToolProvider) closeClientsLocked() error { + var errs []error + for _, server := range p.options.Servers { + client := p.clients[server.Name] + if client == nil { + continue + } + if err := client.Disconnect(); err != nil { + errs = append(errs, fmt.Errorf("disconnect MCP server %q: %w", server.Name, err)) + } + delete(p.clients, server.Name) + } + return errors.Join(errs...) +} + +// dialMCPServer opens and initializes one MCP session. The handshake happens +// here rather than lazily on first use: a server that cannot initialize has no +// tools to publish, and failing at connect time attributes the failure to the +// server that caused it. +func dialMCPServer(ctx context.Context, server MCPServer) (mcpClient, error) { + channel, err := mcpTransport(server) + if err != nil { + return nil, err + } + client := mcpclient.NewClient(channel) + if err := client.Start(ctx); err != nil { + return nil, fmt.Errorf("start transport: %w", err) + } + request := mcp.InitializeRequest{} + request.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION + request.Params.ClientInfo = mcp.Implementation{Name: mcpClientName, Version: mcpClientVersion} + if _, err := client.Initialize(ctx, request); err != nil { + return nil, errors.Join(fmt.Errorf("initialize session: %w", err), client.Close()) + } + return &mcpSession{client: client}, nil +} + +func mcpTransport(server MCPServer) (transport.Interface, error) { + switch { + case server.Command != "": + return transport.NewStdio(server.Command, append([]string(nil), server.Env...), server.Args...), nil + case server.StreamableHTTP: + return transport.NewStreamableHTTP(server.URL, transport.WithHTTPHeaders(maps.Clone(server.Headers))) + default: + return transport.NewSSE(server.URL, transport.WithHeaders(maps.Clone(server.Headers))) + } +} + +// mcpSession adapts an mcp-go client to the narrow surface this provider uses. +type mcpSession struct { + client *mcpclient.Client +} + +func (s *mcpSession) ListTools(ctx context.Context) ([]mcp.Tool, error) { + result, err := s.client.ListTools(ctx, mcp.ListToolsRequest{}) + if err != nil { + return nil, err + } + return result.Tools, nil +} + +func (s *mcpSession) CallTool(ctx context.Context, name string, arguments map[string]any) (any, error) { + request := mcp.CallToolRequest{} + request.Params.Name = name + request.Params.Arguments = arguments + result, err := s.client.CallTool(ctx, request) + if err != nil { + return nil, err + } + return result, nil +} + +func (s *mcpSession) Disconnect() error { return s.client.Close() } + +// namespacedToolName scopes a server's tool to that server, so two servers +// exposing "search" stay distinguishable. The name is also the catalog's +// preference key, so it is part of the stored contract, not cosmetic. +func namespacedToolName(server, tool string) string { + return fmt.Sprintf("%s_%s", server, tool) +} + +func projectMCPTool(server string, tool mcp.Tool, client mcpClient) (api.ToolDefinition, aitools.ToolCatalogEntry, error) { + name := namespacedToolName(server, tool.Name) + inputSchema, err := argumentsSchema(mcp.ToolArgumentsSchema(tool.InputSchema)) + if err != nil { + return api.ToolDefinition{}, aitools.ToolCatalogEntry{}, fmt.Errorf("read input schema of MCP tool %q from server %q: %w", name, server, err) + } + outputSchema, err := argumentsSchema(mcp.ToolArgumentsSchema(tool.OutputSchema)) + if err != nil { + return api.ToolDefinition{}, aitools.ToolCatalogEntry{}, fmt.Errorf("read output schema of MCP tool %q from server %q: %w", name, server, err) + } + title := tool.Title + if title == "" { + title = name + } + catalog := aitools.ToolCatalogEntry{ + Name: name, Title: title, Description: tool.Description, + Source: "mcp", Server: server, PreferenceKey: name, + DefaultPermission: api.ToolModeAuto, InputSchema: aitools.ObjectSchema(inputSchema), + OutputSchema: outputSchema, + } + // _meta is where an MCP server publishes the grouping, icon and permission + // hints the catalog understands. + if tool.Meta != nil { + aitools.ApplyToolMetadata(&catalog, tool.Meta.AdditionalFields) + } + required := append([]string(nil), tool.InputSchema.Required...) + definition := api.ToolDefinition{ + Name: name, Description: tool.Description, + InputSchema: maps.Clone(catalog.InputSchema), Group: catalog.Group, + Parent: catalog.Parent, Icon: catalog.Icon, Strict: catalog.Strict, + DefaultPermission: api.ToolMode(catalog.DefaultPermission), + Handler: func(ctx context.Context, input map[string]any) (any, error) { + for _, field := range required { + if _, present := input[field]; !present { + return nil, fmt.Errorf("MCP tool %q requires field %q", name, field) + } + } + return client.CallTool(ctx, tool.Name, input) + }, + } + return definition, catalog, nil +} + +// argumentsSchema renders an mcp-go schema as the plain JSON Schema map the +// catalog stores, and nil when the server declared no schema at all — the +// zero struct marshals to `{"type":"","properties":null}`, which is not one. +func argumentsSchema(schema mcp.ToolArgumentsSchema) (map[string]any, error) { + if schema.Type == "" && schema.Properties == nil && schema.Defs == nil { + return nil, nil + } + encoded, err := json.Marshal(schema) + if err != nil { + return nil, err + } + var decoded map[string]any + if err := json.Unmarshal(encoded, &decoded); err != nil { + return nil, err + } + return decoded, nil +} + +func cloneToolSet(set ToolSet) ToolSet { + return ToolSet{ + Definitions: append([]api.ToolDefinition(nil), set.Definitions...), + Catalog: append([]aitools.ToolCatalogEntry(nil), set.Catalog...), + } +} diff --git a/pkg/aichat/mcp_provider_ginkgo_test.go b/pkg/aichat/mcp_provider_ginkgo_test.go new file mode 100644 index 0000000..1fcf459 --- /dev/null +++ b/pkg/aichat/mcp_provider_ginkgo_test.go @@ -0,0 +1,194 @@ +package aichat + +import ( + "context" + "errors" + + "github.com/mark3labs/mcp-go/mcp" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type fakeMCPClient struct { + tools []mcp.Tool + discoverErr error + discoveries int + disconnects int + disconnect error + calls []map[string]any +} + +func (f *fakeMCPClient) ListTools(context.Context) ([]mcp.Tool, error) { + f.discoveries++ + return f.tools, f.discoverErr +} + +func (f *fakeMCPClient) CallTool(_ context.Context, name string, arguments map[string]any) (any, error) { + f.calls = append(f.calls, arguments) + return map[string]any{"tool": name, "arguments": arguments}, nil +} + +func (f *fakeMCPClient) Disconnect() error { + f.disconnects++ + return f.disconnect +} + +var _ = Describe("Captain MCP tool provider", func() { + It("caches explicit clients and projects executable, server-scoped tools", func(ctx SpecContext) { + client := &fakeMCPClient{tools: []mcp.Tool{newFakeMCPTool("lookup")}} + provider := NewMCPToolProvider(MCPToolProviderOptions{ + Servers: []MCPServer{{Name: "weather", URL: "https://example.com/mcp", StreamableHTTP: true}}, + }) + DeferCleanup(provider.Close) + created := 0 + provider.clientFactory = func(_ context.Context, server MCPServer) (mcpClient, error) { + created++ + Expect(server.Name).To(Equal("weather")) + Expect(server.URL).To(Equal("https://example.com/mcp")) + Expect(server.StreamableHTTP).To(BeTrue()) + return client, nil + } + + first, err := provider.ToolSet(ctx) + Expect(err).NotTo(HaveOccurred()) + second, err := provider.ToolSet(ctx) + Expect(err).NotTo(HaveOccurred()) + Expect(created).To(Equal(1)) + Expect(client.discoveries).To(Equal(1)) + Expect(second.Definitions).To(HaveLen(1)) + Expect(second.Definitions[0].Name).To(Equal(first.Definitions[0].Name)) + Expect(second.Catalog).To(Equal(first.Catalog)) + Expect(first.Definitions).To(HaveLen(1)) + // The server name scopes the tool, and the same string is the stored + // preference key, so it may not drift. + Expect(first.Definitions[0].Name).To(Equal("weather_lookup")) + Expect(first.Catalog).To(HaveLen(1)) + Expect(first.Catalog[0].Source).To(Equal("mcp")) + Expect(first.Catalog[0].Server).To(Equal("weather")) + Expect(first.Catalog[0].PreferenceKey).To(Equal("weather_lookup")) + Expect(first.Catalog[0].InputSchema).To(HaveKeyWithValue("type", "object")) + + // The server is called with its own unscoped name, not the scoped one. + result, err := first.Definitions[0].Handler(ctx, map[string]any{"city": "Cape Town"}) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(Equal(map[string]any{ + "tool": "lookup", + "arguments": map[string]any{"city": "Cape Town"}, + })) + }) + + It("rejects a call that omits a field the server declared required", func(ctx SpecContext) { + tool := newFakeMCPTool("lookup") + tool.InputSchema.Required = []string{"city"} + provider := NewMCPToolProvider(MCPToolProviderOptions{ + Servers: []MCPServer{{Name: "weather", Command: "weather-server"}}, + }) + DeferCleanup(provider.Close) + client := &fakeMCPClient{tools: []mcp.Tool{tool}} + provider.clientFactory = func(context.Context, MCPServer) (mcpClient, error) { return client, nil } + + set, err := provider.ToolSet(ctx) + Expect(err).NotTo(HaveOccurred()) + + _, err = set.Definitions[0].Handler(ctx, map[string]any{"country": "ZA"}) + Expect(err).To(MatchError(`MCP tool "weather_lookup" requires field "city"`)) + Expect(client.calls).To(BeEmpty(), "a call missing a required field must not reach the server") + }) + + It("fails the whole load and disconnects clients when scoped tool names collide", func(ctx SpecContext) { + // Scoping joins server and tool with an underscore, so a server whose + // name is a prefix of another's can still collide. That ambiguity is + // the only way two distinct servers can produce one name. + clients := []*fakeMCPClient{ + {tools: []mcp.Tool{newFakeMCPTool("b_c")}}, + {tools: []mcp.Tool{newFakeMCPTool("c")}}, + } + provider := NewMCPToolProvider(MCPToolProviderOptions{Servers: []MCPServer{ + {Name: "a", Command: "first-server"}, + {Name: "a_b", Command: "second-server"}, + }}) + created := 0 + provider.clientFactory = func(context.Context, MCPServer) (mcpClient, error) { + client := clients[created] + created++ + return client, nil + } + + _, err := provider.ToolSet(ctx) + Expect(err).To(MatchError(ContainSubstring(`duplicate MCP tool definition "a_b_c" from servers "a" and "a_b"`))) + Expect(clients[0].disconnects).To(Equal(1)) + Expect(clients[1].disconnects).To(Equal(1)) + }) + + It("reports the failing server and closes earlier clients", func(ctx SpecContext) { + first := &fakeMCPClient{tools: []mcp.Tool{newFakeMCPTool("first_tool")}} + provider := NewMCPToolProvider(MCPToolProviderOptions{Servers: []MCPServer{ + {Name: "first", Command: "first-server"}, + {Name: "broken", Command: "broken-server"}, + }}) + provider.clientFactory = func(_ context.Context, server MCPServer) (mcpClient, error) { + if server.Name == "broken" { + return nil, errors.New("connection refused") + } + return first, nil + } + + _, err := provider.ToolSet(ctx) + Expect(err).To(MatchError(ContainSubstring(`connect MCP server "broken": connection refused`))) + Expect(first.disconnects).To(Equal(1)) + }) + + It("reports discovery failures and deterministically closes every client", func(ctx SpecContext) { + clients := map[string]*fakeMCPClient{ + "first": {tools: []mcp.Tool{newFakeMCPTool("first_tool")}}, + "broken": {discoverErr: errors.New("list failed")}, + } + provider := NewMCPToolProvider(MCPToolProviderOptions{Servers: []MCPServer{ + {Name: "first", Command: "first-server"}, + {Name: "broken", Command: "broken-server"}, + }}) + provider.clientFactory = func(_ context.Context, server MCPServer) (mcpClient, error) { + return clients[server.Name], nil + } + + _, err := provider.ToolSet(ctx) + Expect(err).To(MatchError(ContainSubstring(`discover MCP tools from server "broken": list failed`))) + Expect(clients["first"].disconnects).To(Equal(1)) + Expect(clients["broken"].disconnects).To(Equal(1)) + + Expect(provider.Close()).To(Succeed()) + Expect(clients["first"].disconnects).To(Equal(1)) + Expect(clients["broken"].disconnects).To(Equal(1)) + }) + + It("returns every disconnect error with its server name", func(ctx SpecContext) { + clients := map[string]*fakeMCPClient{ + "first": {disconnect: errors.New("first close")}, + "second": {disconnect: errors.New("second close")}, + } + provider := NewMCPToolProvider(MCPToolProviderOptions{Servers: []MCPServer{ + {Name: "first", Command: "first-server"}, + {Name: "second", Command: "second-server"}, + }}) + provider.clientFactory = func(_ context.Context, server MCPServer) (mcpClient, error) { + return clients[server.Name], nil + } + _, err := provider.ToolSet(ctx) + Expect(err).NotTo(HaveOccurred()) + + err = provider.Close() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(`disconnect MCP server "first": first close`)) + Expect(err.Error()).To(ContainSubstring(`disconnect MCP server "second": second close`)) + Expect(clients["first"].disconnects).To(Equal(1)) + Expect(clients["second"].disconnects).To(Equal(1)) + }) +}) + +func newFakeMCPTool(name string) mcp.Tool { + return mcp.Tool{ + Name: name, + Description: "Run " + name, + InputSchema: mcp.ToolInputSchema{Type: "object", Properties: map[string]any{}}, + } +} diff --git a/pkg/aichat/messages.go b/pkg/aichat/messages.go new file mode 100644 index 0000000..fb9c1a5 --- /dev/null +++ b/pkg/aichat/messages.go @@ -0,0 +1,202 @@ +package aichat + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/api" +) + +// AttachmentInput is an untrusted browser file reference to be resolved into a +// prepared Captain attachment before provider execution. +type AttachmentInput struct { + ID string + URL string + Filename string + MediaType string +} + +// AttachmentResolver prepares browser references for provider consumption. +type AttachmentResolver interface { + Resolve(context.Context, []AttachmentInput) ([]api.AttachmentRef, error) +} + +type partLocation struct{ message, part int } + +func (s *Service) resolveAttachments(ctx context.Context, messages []UIMessage) (map[partLocation]api.AttachmentRef, error) { + inputs := make([]AttachmentInput, 0) + locations := make([]partLocation, 0) + for messageIndex, message := range messages { + for partIndex, part := range message.Parts { + if part.Type != "file" { + continue + } + if message.Role != string(api.RoleUser) { + return nil, fmt.Errorf("message %d: file parts require user role", messageIndex+1) + } + inputs = append(inputs, AttachmentInput{ + ID: part.AttachmentID, URL: part.URL, Filename: part.Filename, MediaType: part.MediaType, + }) + locations = append(locations, partLocation{message: messageIndex, part: partIndex}) + } + } + if len(inputs) == 0 { + return nil, nil + } + if s.options.Attachments == nil { + return nil, fmt.Errorf("chat attachments require an attachment resolver") + } + refs, err := s.options.Attachments.Resolve(ctx, inputs) + if err != nil { + return nil, fmt.Errorf("resolve chat attachments: %w", err) + } + if len(refs) != len(inputs) { + return nil, fmt.Errorf("attachment resolver returned %d results for %d inputs", len(refs), len(inputs)) + } + resolved := make(map[partLocation]api.AttachmentRef, len(refs)) + for i, ref := range refs { + if err := ref.Validate(); err != nil { + return nil, fmt.Errorf("resolved attachment %d: %w", i+1, err) + } + if !ref.IsPrepared() { + return nil, fmt.Errorf("resolved attachment %d is not prepared", i+1) + } + resolved[locations[i]] = ref + } + return resolved, nil +} + +func requestSpec(request ChatRequest, settings RuntimeSettings, attachments map[partLocation]api.AttachmentRef) (api.Spec, error) { + model := strings.TrimSpace(request.Model) + if model == "" { + model = strings.TrimSpace(settings.Spec.Name) + } + if model == "" { + return api.Spec{}, fmt.Errorf("chat model is required") + } + // Expand before merging: a compact selector ("agent:sol") carries its own + // backend, and merging it unexpanded would keep settings.Spec's backend and run + // a different runtime than the caller asked for. + override, err := api.Model{Name: model, Effort: request.ReasoningEffort, Temperature: request.Temperature}.Expand() + if err != nil { + return api.Spec{}, fmt.Errorf("invalid chat model %q: %w", model, err) + } + spec := settings.Spec.Merge(api.Spec{ + Model: override, + Budget: request.Budget, + ToolPreferences: request.ToolPreferences, + ToolApproval: request.ToolApproval, + Permissions: api.Permissions{Mode: request.PermissionMode}, + SessionID: request.ProviderSessionID, + }) + spec.Prompt.User = "" + spec.Prompt.System = "" + spec.Prompt.AppendSystem = "" + spec.Prompt.Attachments = nil + if request.ToolApproval == nil { + messages, err := canonicalMessages(request.Messages, attachments) + if err != nil { + return api.Spec{}, err + } + system, err := requestSystem(settings.System, request) + if err != nil { + return api.Spec{}, err + } + if system != "" { + messages = append([]api.Message{{Role: api.RoleSystem, Parts: []api.Part{{Type: api.PartText, Text: system}}}}, messages...) + } + spec.Messages = messages + } else { + spec.Messages = nil + } + if err := spec.Validate(); err != nil { + return api.Spec{}, err + } + return spec, nil +} + +func canonicalMessages(messages []UIMessage, attachments map[partLocation]api.AttachmentRef) ([]api.Message, error) { + out := make([]api.Message, 0, len(messages)) + for messageIndex, message := range messages { + role := api.MessageRole(message.Role) + parts := make([]api.Part, 0, len(message.Parts)) + results := make([]api.Part, 0) + for partIndex, part := range message.Parts { + mapped, result, err := canonicalPart(role, part, attachments[partLocation{message: messageIndex, part: partIndex}]) + if err != nil { + return nil, fmt.Errorf("message %d part %d: %w", messageIndex+1, partIndex+1, err) + } + if mapped != nil { + parts = append(parts, *mapped) + } + if result != nil { + results = append(results, *result) + } + } + if len(parts) == 0 { + return nil, fmt.Errorf("message %d (%s) has no provider content", messageIndex+1, role) + } + out = append(out, api.Message{Role: role, Parts: parts}) + if len(results) > 0 { + out = append(out, api.Message{Role: api.RoleTool, Parts: results}) + } + } + return out, nil +} + +func canonicalPart(role api.MessageRole, part UIPart, attachment api.AttachmentRef) (*api.Part, *api.Part, error) { + switch { + case part.Type == "text": + return &api.Part{Type: api.PartText, Text: part.Text}, nil, nil + case part.Type == "reasoning": + return &api.Part{Type: api.PartReasoning, Text: part.Text}, nil, nil + case part.Type == "file": + return &api.Part{Type: api.PartAttachment, Attachment: &attachment}, nil, nil + case part.IsTool(): + if role != api.RoleAssistant { + return nil, nil, fmt.Errorf("tool parts require assistant role") + } + request := &api.Part{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: part.ToolCallID, Name: part.EffectiveToolName(), Input: part.Input, + }} + if part.State != "output-available" && part.State != "output-error" { + return request, nil, nil + } + result := &api.Part{Type: api.PartToolResult, ToolResult: &api.ToolResult{ + ToolCallID: part.ToolCallID, Output: part.Output, + }} + if part.State == "output-error" { + result.ToolResult.Output = nil + result.ToolResult.Error = part.ErrorText + } + return request, result, nil + case part.Type == "step-start" || strings.HasPrefix(part.Type, "data-") || strings.HasPrefix(part.Type, "source-"): + return nil, nil, nil + default: + return nil, nil, fmt.Errorf("unsupported AI SDK part type %q", part.Type) + } +} + +func requestSystem(base string, request ChatRequest) (string, error) { + sections := make([]string, 0, 2) + if strings.TrimSpace(base) != "" { + sections = append(sections, strings.TrimSpace(base)) + } + contextSections := make([]string, 0, 2) + if strings.TrimSpace(request.Context) != "" { + contextSections = append(contextSections, strings.TrimSpace(request.Context)) + } + if len(request.ContextItems) > 0 { + payload, err := json.Marshal(request.ContextItems) + if err != nil { + return "", fmt.Errorf("marshal chat context items: %w", err) + } + contextSections = append(contextSections, "Structured context items JSON:\n"+string(payload)) + } + if len(contextSections) > 0 { + sections = append(sections, "Current UI context:\n"+strings.Join(contextSections, "\n")) + } + return strings.Join(sections, "\n\n"), nil +} diff --git a/pkg/aichat/persistence.go b/pkg/aichat/persistence.go new file mode 100644 index 0000000..ef0d8d4 --- /dev/null +++ b/pkg/aichat/persistence.go @@ -0,0 +1,228 @@ +package aichat + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/flanksource/captain/pkg/api" +) + +type assistantMessageBuilder struct { + message UIMessage + toolParts map[string]int + sessionID string + model string +} + +func newAssistantMessageBuilder(chatID string) *assistantMessageBuilder { + id := "" + if chatID != "" { + id = chatID + "-assistant" + } + return &assistantMessageBuilder{ + message: UIMessage{ID: id, Role: string(api.RoleAssistant), Parts: []UIPart{}}, + toolParts: map[string]int{}, + } +} + +func (s *Service) persistedEvents(ctx context.Context, request ChatRequest, source <-chan api.Event) <-chan api.Event { + if request.ThreadID == "" { + return source + } + out := make(chan api.Event) + go func() { + defer close(out) + builder := newAssistantMessageBuilder(request.ID) + persisted := false + for event := range source { + if err := builder.apply(event); err != nil { + sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error(), Model: event.Model}) + return + } + if !persisted && event.Kind != api.EventResult && event.Kind != api.EventError && event.SessionID != "" { + if err := s.persistEvent(ctx, request.ThreadID, event); err != nil { + sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error(), Model: event.Model}) + return + } + } + if !persisted && (event.Kind == api.EventResult || event.Kind == api.EventError) { + if err := s.persistCompletedTurn(ctx, request.ThreadID, builder, event); err != nil { + sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error(), Model: event.Model}) + return + } + persisted = true + } + if !sendEvent(ctx, out, event) { + return + } + } + if !persisted && len(builder.message.Parts) > 0 { + if err := s.options.Threads.AppendMessage(ctx, request.ThreadID, builder.message); err != nil { + sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: fmt.Sprintf("persist assistant message: %v", err)}) + } + } + }() + return out +} + +func sendEvent(ctx context.Context, target chan<- api.Event, event api.Event) bool { + select { + case target <- event: + return true + case <-ctx.Done(): + return false + } +} + +func (s *Service) persistCompletedTurn(ctx context.Context, threadID string, builder *assistantMessageBuilder, event api.Event) error { + if err := s.persistEvent(ctx, threadID, event); err != nil { + return err + } + if len(builder.message.Parts) == 0 { + return fmt.Errorf("completed assistant turn has no message parts") + } + if err := s.options.Threads.AppendMessage(ctx, threadID, builder.message); err != nil { + return fmt.Errorf("persist assistant message: %w", err) + } + return nil +} + +func (b *assistantMessageBuilder) apply(event api.Event) error { + if event.SessionID != "" { + b.sessionID = event.SessionID + } + if event.Model != "" { + b.model = event.Model + } + switch event.Kind { + case api.EventText: + b.appendText("text", event.Text) + case api.EventThinking: + b.appendText("reasoning", event.Text) + case api.EventToolUse: + return b.toolUse(event) + case api.EventPermission: + return b.permission(event) + case api.EventToolResult: + return b.toolResult(event) + case api.EventResult: + return b.result(event) + case api.EventError: + payload, err := json.Marshal(map[string]string{"error": event.Error}) + if err != nil { + return err + } + b.message.Parts = append(b.message.Parts, UIPart{Type: "data-error", Data: payload}) + case api.EventSystem: + } + return nil +} + +func (b *assistantMessageBuilder) appendText(partType, text string) { + if text == "" { + return + } + if len(b.message.Parts) > 0 && b.message.Parts[len(b.message.Parts)-1].Type == partType { + b.message.Parts[len(b.message.Parts)-1].Text += text + return + } + b.message.Parts = append(b.message.Parts, UIPart{Type: partType, Text: text}) +} + +func (b *assistantMessageBuilder) toolUse(event api.Event) error { + if event.ToolCallID == "" || event.Tool == "" { + return fmt.Errorf("persist tool use requires a tool name and call ID") + } + if _, exists := b.toolParts[event.ToolCallID]; exists { + return fmt.Errorf("persist duplicate tool call %q", event.ToolCallID) + } + input, err := json.Marshal(event.Input) + if err != nil { + return fmt.Errorf("marshal tool %q input: %w", event.Tool, err) + } + b.toolParts[event.ToolCallID] = len(b.message.Parts) + b.message.Parts = append(b.message.Parts, UIPart{ + Type: "dynamic-tool", ToolName: event.Tool, ToolCallID: event.ToolCallID, + State: "input-available", Input: input, + }) + return nil +} + +func (b *assistantMessageBuilder) permission(event api.Event) error { + part, err := b.toolPart(event.ToolCallID, event.Tool) + if err != nil { + return err + } + part.State = "approval-requested" + part.Approval = &Approval{ID: event.ToolCallID} + return nil +} + +func (b *assistantMessageBuilder) toolResult(event api.Event) error { + part, err := b.toolPart(event.ToolCallID, event.Tool) + if err != nil { + return err + } + if event.Success { + part.State = "output-available" + part.Output, err = jsonValue(event.Text) + if err != nil { + return fmt.Errorf("marshal tool %q output: %w", event.Tool, err) + } + return nil + } + part.State = "output-error" + part.ErrorText = event.Text + return nil +} + +func (b *assistantMessageBuilder) toolPart(callID, name string) (*UIPart, error) { + index, ok := b.toolParts[callID] + if !ok { + return nil, fmt.Errorf("persist tool event %q has no matching tool use", callID) + } + part := &b.message.Parts[index] + if name != "" && part.ToolName != name { + return nil, fmt.Errorf("persist tool event %q names %q, want %q", callID, name, part.ToolName) + } + return part, nil +} + +func (b *assistantMessageBuilder) result(event api.Event) error { + dataType := "data-result" + data := event.StructuredData + if event.ToolApproval != nil { + dataType = "data-tool-approval" + var err error + data, err = json.Marshal(event.ToolApproval) + if err != nil { + return fmt.Errorf("marshal tool approval state: %w", err) + } + } else if len(data) == 0 { + var err error + data, err = json.Marshal(map[string]bool{"success": event.Success}) + if err != nil { + return fmt.Errorf("marshal result state: %w", err) + } + } + b.message.Parts = append(b.message.Parts, UIPart{Type: dataType, Data: data}) + success := event.Success + b.message.Metadata = &MessageMetadata{ + ProviderSessionID: b.sessionID, Model: b.model, Cost: event.CostUSD, Success: &success, + } + if event.Usage != nil { + b.message.Metadata.Usage = usageMetadata(*event.Usage) + b.message.Metadata.ContextTokens = event.Usage.InputTokens + } + return nil +} + +func jsonValue(text string) (json.RawMessage, error) { + raw := json.RawMessage(text) + if json.Valid(raw) { + return raw, nil + } + payload, err := json.Marshal(text) + return payload, err +} diff --git a/pkg/aichat/provider_config.go b/pkg/aichat/provider_config.go new file mode 100644 index 0000000..0094108 --- /dev/null +++ b/pkg/aichat/provider_config.go @@ -0,0 +1,66 @@ +package aichat + +import ( + "context" + "fmt" + "reflect" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" +) + +// ProviderConfigRequest carries the canonically resolved model and the runtime +// config assembled by the chat service. +type ProviderConfigRequest struct { + Model api.Model + Config api.Config +} + +// ProviderConfigSource supplies request-scoped provider identities and +// credentials without owning provider resolution or construction. +type ProviderConfigSource interface { + ConfiguredProviders(context.Context) ([]api.Backend, error) + ProviderConfig(context.Context, ProviderConfigRequest) (api.Config, error) +} + +func (s *Service) annotateConfiguredModels(ctx context.Context, models ModelCatalogResponse) error { + if s.options.ProviderConfig == nil { + return nil + } + backends, err := s.options.ProviderConfig.ConfiguredProviders(ctx) + if err != nil { + return fmt.Errorf("load configured chat providers: %w", err) + } + configured := make(map[string]bool, len(backends)) + for _, backend := range backends { + if backend == "" { + return fmt.Errorf("configured chat provider backend is required") + } + configured[ai.BackendToProvider(backend)] = true + } + for i := range models { + models[i].Configured = models[i].Configured || configured[models[i].Provider] + } + return nil +} + +func (s *Service) resolveProvider(ctx context.Context, config api.Config) (api.StreamingProvider, error) { + if s.options.ProviderConfig != nil { + resolved, err := ai.ResolveModelSelectors(config.Model) + if err != nil { + return nil, fmt.Errorf("resolve chat model: %w", err) + } + config.Model = resolved + config, err = s.options.ProviderConfig.ProviderConfig(ctx, ProviderConfigRequest{ + Model: resolved, Config: config, + }) + if err != nil { + return nil, fmt.Errorf("load chat provider config for %s: %w", resolved.Backend, err) + } + if !reflect.DeepEqual(config.Model, resolved) { + return nil, fmt.Errorf("provider config source changed the resolved chat model from %q (%s) to %q (%s)", + resolved.Name, resolved.Backend, config.Model.Name, config.Model.Backend) + } + } + return s.resolver.Provider(ctx, config) +} diff --git a/pkg/aichat/resolver.go b/pkg/aichat/resolver.go new file mode 100644 index 0000000..9db2702 --- /dev/null +++ b/pkg/aichat/resolver.go @@ -0,0 +1,58 @@ +package aichat + +import ( + "context" + "fmt" + + "github.com/flanksource/captain/pkg/ai" + _ "github.com/flanksource/captain/pkg/ai/provider" + "github.com/flanksource/captain/pkg/api" +) + +// Resolver is the injectable boundary around Captain's model catalog and +// provider construction. The default implementation delegates to Captain's +// canonical resolver; tests may replace it with a fake provider. +type Resolver interface { + Models(context.Context) (ModelCatalogResponse, error) + Provider(context.Context, api.Config) (api.StreamingProvider, error) +} + +type captainResolver struct{} + +func (captainResolver) Models(_ context.Context) (ModelCatalogResponse, error) { + configured := make([]string, 0, 4) + for _, backend := range []api.Backend{ + api.BackendAnthropic, api.BackendOpenAI, api.BackendGemini, api.BackendDeepSeek, + } { + resolved, err := ai.ResolveAPIKey(backend) + if err != nil { + return nil, fmt.Errorf("resolve %s credentials: %w", backend, err) + } + if resolved.Token != "" { + configured = append(configured, ai.BackendToProvider(backend)) + } + } + return ai.LiveCatalogInfo(configured) +} + +func (captainResolver) Provider(_ context.Context, config api.Config) (api.StreamingProvider, error) { + provider, err := ai.NewProvider(config) + if err != nil { + return nil, err + } + streaming, ok := api.ProviderAs[api.StreamingProvider](provider) + if !ok { + if closeErr := closeProvider(provider); closeErr != nil { + return nil, fmt.Errorf("backend %q does not support streaming; close provider: %w", provider.GetBackend(), closeErr) + } + return nil, fmt.Errorf("backend %q does not support streaming", provider.GetBackend()) + } + return streaming, nil +} + +func closeProvider(provider api.Provider) error { + if closer, ok := api.ProviderAs[api.CloseableProvider](provider); ok { + return closer.Close() + } + return nil +} diff --git a/pkg/aichat/runtime_settings.go b/pkg/aichat/runtime_settings.go new file mode 100644 index 0000000..0e9fff1 --- /dev/null +++ b/pkg/aichat/runtime_settings.go @@ -0,0 +1,55 @@ +package aichat + +import ( + "encoding/json" + "fmt" + "net/http" +) + +type requestError struct { + status int + text string +} + +func (e requestError) Error() string { return e.text } + +func requestErrorStatus(err error) int { + if typed, ok := err.(requestError); ok { + return typed.status + } + return http.StatusBadRequest +} + +func enforceRuntimeSettings(request ChatRequest, settings RuntimeSettings) error { + if settings.MonthlyBudgetUSD > 0 && settings.CurrentMonthCostUSD >= settings.MonthlyBudgetUSD { + return requestError{status: http.StatusPaymentRequired, text: fmt.Sprintf( + "chat monthly cost budget exhausted: $%.4f used of $%.4f", + settings.CurrentMonthCostUSD, settings.MonthlyBudgetUSD, + )} + } + if settings.MonthlyTokenBudget > 0 && settings.CurrentMonthTokens >= settings.MonthlyTokenBudget { + return requestError{status: http.StatusPaymentRequired, text: fmt.Sprintf( + "chat monthly token budget exhausted: %d used of %d", + settings.CurrentMonthTokens, settings.MonthlyTokenBudget, + )} + } + if settings.MaxInputTokens <= 0 { + return nil + } + raw, err := json.Marshal(struct { + Messages []UIMessage `json:"messages,omitempty"` + Context string `json:"context,omitempty"` + ContextItems []ChatContextItem `json:"contextItems,omitempty"` + }{Messages: request.Messages, Context: request.Context, ContextItems: request.ContextItems}) + if err != nil { + return fmt.Errorf("estimate chat input tokens: %w", err) + } + estimated := (len(raw) + 3) / 4 + if estimated > settings.MaxInputTokens { + return requestError{status: http.StatusRequestEntityTooLarge, text: fmt.Sprintf( + "chat input is about %d tokens, exceeding the configured per-turn limit of %d", + estimated, settings.MaxInputTokens, + )} + } + return nil +} diff --git a/pkg/aichat/service.go b/pkg/aichat/service.go new file mode 100644 index 0000000..b7c08d3 --- /dev/null +++ b/pkg/aichat/service.go @@ -0,0 +1,239 @@ +package aichat + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + + aitools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/commons/logger" +) + +var serviceLog = logger.GetLogger("aichat") + +// RuntimeSettings are application-owned defaults and provider construction +// settings evaluated for each request. +// RuntimeSettings is the request-scoped application configuration for a chat. +// +// The default model lives in Spec.Model — there is deliberately no DefaultModel +// string beside it. A bare name next to a structured Spec is the lossy pattern: +// it cannot carry a backend/mode/effort, so whatever it named got re-inferred, and +// when both were set they could silently disagree. Spec.Model can say +// {Name: "sol", Mode: ModeAgent} and mean it. +type RuntimeSettings struct { + System string + Spec api.Spec + ProviderConfig api.Config + MaxInputTokens int + MonthlyTokenBudget int + CurrentMonthTokens int + MonthlyBudgetUSD float64 + CurrentMonthCostUSD float64 +} + +// RuntimeSettingsProvider supplies request-scoped application settings. +type RuntimeSettingsProvider interface { + RuntimeSettings(context.Context) (RuntimeSettings, error) +} + +type RuntimeSettingsProviderFunc func(context.Context) (RuntimeSettings, error) + +func (f RuntimeSettingsProviderFunc) RuntimeSettings(ctx context.Context) (RuntimeSettings, error) { + return f(ctx) +} + +// ServiceOptions injects every application-owned chat dependency. A nil +// Resolver uses Captain's canonical model/provider resolver. +type ServiceOptions struct { + Resolver Resolver + ProviderConfig ProviderConfigSource + Settings RuntimeSettingsProvider + Tools ToolProvider + MCP ToolProvider + Attachments AttachmentResolver + Threads ThreadStore +} + +// Service is Captain's AI SDK-compatible HTTP chat service. +type Service struct { + options ServiceOptions + resolver Resolver +} + +func NewService(options ServiceOptions) *Service { + resolver := options.Resolver + if resolver == nil { + resolver = captainResolver{} + } + return &Service{options: options, resolver: resolver} +} + +func (s *Service) Handler() http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("POST /api/chat", s.handleChat) + mux.HandleFunc("GET /api/chat/models", s.handleModels) + mux.HandleFunc("GET /api/chat/tools", s.handleTools) + s.registerThreadRoutes(mux) + return mux +} + +func (s *Service) handleModels(w http.ResponseWriter, request *http.Request) { + models, err := s.resolver.Models(request.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + if err := s.annotateConfiguredModels(request.Context(), models); err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + if err := writeJSON(w, http.StatusOK, models); err != nil { + serviceLog.Errorf("write chat models response: %v", err) + } +} + +func (s *Service) handleTools(w http.ResponseWriter, request *http.Request) { + set, err := s.loadTools(request.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := writeJSON(w, http.StatusOK, aitools.ToolCatalog{Tools: set.Catalog}); err != nil { + serviceLog.Errorf("write chat tools response: %v", err) + } +} + +func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { + var chat ChatRequest + if err := json.NewDecoder(request.Body).Decode(&chat); err != nil { + http.Error(w, fmt.Sprintf("invalid chat request: %v", err), http.StatusBadRequest) + return + } + settings, err := s.runtimeSettings(request.Context()) + if err != nil { + http.Error(w, fmt.Sprintf("load chat runtime settings: %v", err), http.StatusInternalServerError) + return + } + if err := enforceRuntimeSettings(chat, settings); err != nil { + http.Error(w, err.Error(), requestErrorStatus(err)) + return + } + if err := s.resolveThreadSession(request.Context(), &chat); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + attachments, err := s.resolveAttachments(request.Context(), chat.Messages) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + spec, err := requestSpec(chat, settings, attachments) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + set, err := s.loadTools(request.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := s.persistIncoming(request.Context(), chat); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + config := settings.ProviderConfig + config.Model = spec.Model + config.Budget = spec.Budget + config.SessionID = spec.SessionID + config.Tools = set.Definitions + provider, err := s.resolveProvider(request.Context(), config) + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + defer func() { + if closeErr := closeProvider(provider); closeErr != nil { + serviceLog.Errorf("close chat provider: %v", closeErr) + } + }() + if len(set.Definitions) > 0 { + capability, ok := api.ProviderAs[api.ToolCapableProvider](provider) + if !ok || !capability.SupportsCallerTools() { + http.Error(w, fmt.Sprintf("backend %q does not support caller tools", provider.GetBackend()), http.StatusBadRequest) + return + } + } + streamContext, cancel := context.WithCancel(request.Context()) + defer cancel() + events, err := provider.ExecuteStream(streamContext, spec) + if err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + return + } + writer, err := NewSSEWriter(w) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := WriteEventStream(writer, s.persistedEvents(streamContext, chat, events)); err != nil { + serviceLog.Errorf("stream chat response: %v", err) + } +} + +func (s *Service) runtimeSettings(ctx context.Context) (RuntimeSettings, error) { + if s.options.Settings == nil { + return RuntimeSettings{}, nil + } + return s.options.Settings.RuntimeSettings(ctx) +} + +func (s *Service) resolveThreadSession(ctx context.Context, request *ChatRequest) error { + if request.ThreadID == "" { + return nil + } + if s.options.Threads == nil { + return fmt.Errorf("thread persistence is not configured") + } + thread, err := s.options.Threads.Get(ctx, request.ThreadID) + if err != nil { + return err + } + if request.ProviderSessionID == "" { + request.ProviderSessionID = thread.ProviderSessionID + } + return nil +} + +func (s *Service) persistIncoming(ctx context.Context, request ChatRequest) error { + if request.ThreadID == "" || len(request.Messages) == 0 { + return nil + } + last := request.Messages[len(request.Messages)-1] + if strings.EqualFold(last.Role, string(api.RoleUser)) { + return s.options.Threads.AppendMessage(ctx, request.ThreadID, last) + } + return nil +} + +func (s *Service) persistEvent(ctx context.Context, threadID string, event api.Event) error { + if event.SessionID != "" { + if err := s.options.Threads.SetProviderSession(ctx, threadID, event.SessionID); err != nil { + return fmt.Errorf("persist provider session: %w", err) + } + } + if event.Kind != api.EventResult || event.Usage == nil { + return nil + } + _, err := s.options.Threads.AddUsage(ctx, threadID, TurnUsage{ + InputTokens: event.Usage.InputTokens, OutputTokens: event.Usage.OutputTokens, + ReasoningTokens: event.Usage.ReasoningTokens, CacheReadTokens: event.Usage.CacheReadTokens, + CacheWriteTokens: event.Usage.CacheWriteTokens, CostUSD: event.CostUSD, + }) + if err != nil { + return fmt.Errorf("persist thread usage: %w", err) + } + return nil +} diff --git a/pkg/aichat/service_ginkgo_test.go b/pkg/aichat/service_ginkgo_test.go new file mode 100644 index 0000000..299ad6f --- /dev/null +++ b/pkg/aichat/service_ginkgo_test.go @@ -0,0 +1,430 @@ +package aichat_test + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" +) + +type fakeResolver struct { + models aichat.ModelCatalogResponse + provider *fakeStreamingProvider + configs []api.Config +} + +type fakeProviderConfigSource struct { + backends []api.Backend + config api.Config + requests []aichat.ProviderConfigRequest + resolve func(aichat.ProviderConfigRequest) (api.Config, error) +} + +func (f *fakeProviderConfigSource) ConfiguredProviders(context.Context) ([]api.Backend, error) { + return append([]api.Backend(nil), f.backends...), nil +} + +func (f *fakeProviderConfigSource) ProviderConfig(_ context.Context, request aichat.ProviderConfigRequest) (api.Config, error) { + f.requests = append(f.requests, request) + if f.resolve != nil { + return f.resolve(request) + } + config := request.Config + config.APIKey = f.config.APIKey + config.APIURL = f.config.APIURL + return config, nil +} + +func (f *fakeResolver) Models(context.Context) (aichat.ModelCatalogResponse, error) { + return f.models, nil +} + +func (f *fakeResolver) Provider(_ context.Context, config api.Config) (api.StreamingProvider, error) { + f.configs = append(f.configs, config) + return f.provider, nil +} + +type fakeStreamingProvider struct { + events []api.Event + specs []api.Spec + execute func(context.Context, api.Spec) (<-chan api.Event, error) +} + +func (f *fakeStreamingProvider) Execute(context.Context, api.Spec) (*api.Response, error) { + return nil, fmt.Errorf("buffered execution is not used by chat") +} + +func (f *fakeStreamingProvider) ExecuteStream(ctx context.Context, spec api.Spec) (<-chan api.Event, error) { + f.specs = append(f.specs, spec) + if f.execute != nil { + return f.execute(ctx, spec) + } + events := make(chan api.Event, len(f.events)) + for _, event := range f.events { + events <- event + } + close(events) + return events, nil +} + +func (f *fakeStreamingProvider) GetModel() string { return "test-model" } +func (f *fakeStreamingProvider) GetBackend() api.Backend { return api.BackendOpenAI } +func (f *fakeStreamingProvider) SupportsCallerTools() bool { return true } + +type fakeAttachmentResolver struct{} + +func (fakeAttachmentResolver) Resolve(_ context.Context, inputs []aichat.AttachmentInput) ([]api.AttachmentRef, error) { + refs := make([]api.AttachmentRef, len(inputs)) + for i, input := range inputs { + refs[i] = api.AttachmentRef{ + ID: api.AttachmentIDPrefix + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + Filename: input.Filename, MediaType: input.MediaType, + }.WithPreparedContent(api.AttachmentContent{Bytes: []byte("image")}) + } + return refs, nil +} + +func requestJSON(method, path string, body any) *http.Request { + var payload bytes.Buffer + Expect(json.NewEncoder(&payload).Encode(body)).To(Succeed()) + return httptest.NewRequest(method, path, &payload) +} + +var _ = Describe("Captain aichat service", func() { + It("annotates the model catalog with request-scoped configured providers", func() { + resolver := &fakeResolver{models: aichat.ModelCatalogResponse{ + {ID: "anthropic/claude-sonnet", Provider: "anthropic", Label: "Claude"}, + {ID: "openai/gpt", Provider: "openai", Label: "GPT"}, + }} + source := &fakeProviderConfigSource{backends: []api.Backend{api.BackendOpenAI}} + service := aichat.NewService(aichat.ServiceOptions{Resolver: resolver, ProviderConfig: source}) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/api/chat/models", nil)) + Expect(response.Code).To(Equal(http.StatusOK)) + var models aichat.ModelCatalogResponse + Expect(json.Unmarshal(response.Body.Bytes(), &models)).To(Succeed()) + Expect(models).To(HaveLen(2)) + Expect(models[0].Configured).To(BeFalse()) + Expect(models[1].Configured).To(BeTrue()) + }) + + It("applies request-scoped credentials after canonical model selection", func() { + provider := &fakeStreamingProvider{events: []api.Event{ + {Kind: api.EventText, Text: "done"}, + {Kind: api.EventResult, Success: true, Model: "gpt-5.4"}, + }} + resolver := &fakeResolver{provider: provider} + source := &fakeProviderConfigSource{ + backends: []api.Backend{api.BackendOpenAI}, + config: api.Config{APIKey: "request-token", APIURL: "https://tenant-x.example/ai"}, + } + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, ProviderConfig: source, + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{Spec: api.Spec{Model: api.Model{Name: "api:gpt-5.4"}}}, nil + }), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "hello"}}}}, + })) + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(source.requests).To(HaveLen(1)) + Expect(source.requests[0].Model.Backend).To(Equal(api.BackendOpenAI)) + Expect(source.requests[0].Model.Name).To(Equal("gpt-5.4")) + Expect(resolver.configs).To(HaveLen(1)) + Expect(resolver.configs[0].APIKey).To(Equal("request-token")) + Expect(resolver.configs[0].APIURL).To(Equal("https://tenant-x.example/ai")) + }) + + It("rejects provider configuration that changes the canonical resolved model", func() { + resolver := &fakeResolver{provider: &fakeStreamingProvider{}} + source := &fakeProviderConfigSource{resolve: func(request aichat.ProviderConfigRequest) (api.Config, error) { + config := request.Config + config.Model = api.Model{Name: "claude-sonnet-4-6", Backend: api.BackendAnthropic} + return config, nil + }} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, ProviderConfig: source, + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{Spec: api.Spec{Model: api.Model{Name: "api:gpt-5.4"}}}, nil + }), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "hello"}}}}, + })) + Expect(response.Code).To(Equal(http.StatusServiceUnavailable)) + Expect(response.Body.String()).To(ContainSubstring("provider config source changed the resolved chat model")) + Expect(resolver.configs).To(BeEmpty()) + }) + + It("merges request budget overrides without erasing runtime defaults", func() { + provider := &fakeStreamingProvider{events: []api.Event{ + {Kind: api.EventText, Text: "done"}, + {Kind: api.EventResult, Success: true, Model: "test-model"}, + }} + resolver := &fakeResolver{provider: provider} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{Spec: api.Spec{ + Model: api.Model{Name: "openai/test-model"}, + Budget: api.Budget{Cost: 5, MaxTokens: 2_000, MaxTurns: 3}, + }}, nil + }), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "hello"}}}}, + Budget: api.Budget{MaxTokens: 1_000}, + })) + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(provider.specs).To(HaveLen(1)) + Expect(provider.specs[0].Budget).To(Equal(api.Budget{Cost: 5, MaxTokens: 1_000, MaxTurns: 3})) + }) + + It("rejects exhausted runtime budgets before provider construction", func() { + resolver := &fakeResolver{provider: &fakeStreamingProvider{}} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{ + Spec: api.Spec{Model: api.Model{Name: "openai/test-model"}}, + MonthlyBudgetUSD: 10, CurrentMonthCostUSD: 10, + }, nil + }), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "hello"}}}}, + })) + Expect(response.Code).To(Equal(http.StatusPaymentRequired)) + Expect(resolver.configs).To(BeEmpty()) + }) + + It("serves models and tools from injected Captain seams", func() { + resolver := &fakeResolver{models: aichat.ModelCatalogResponse{{ + ID: "openai/test-model", Provider: "openai", Label: "Test", Configured: true, + }}} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, + Tools: aichat.ToolProviderFunc(func(context.Context) (aichat.ToolSet, error) { + return aichat.ToolSet{Definitions: []api.ToolDefinition{{ + Name: "invoice_get", Group: "billing", Description: "Get invoice", + InputSchema: map[string]any{"type": "object"}, + Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}}, nil + }), + MCP: aichat.ToolProviderFunc(func(context.Context) (aichat.ToolSet, error) { + return aichat.ToolSet{Definitions: []api.ToolDefinition{{ + Name: "docs_search", Group: "docs", Description: "Search documentation", + Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}}, nil + }), + }) + + models := httptest.NewRecorder() + service.Handler().ServeHTTP(models, httptest.NewRequest(http.MethodGet, "/api/chat/models", nil)) + Expect(models.Code).To(Equal(http.StatusOK)) + Expect(models.Body.String()).To(MatchJSON(`[{"id":"openai/test-model","provider":"openai","label":"Test","reasoning":false,"temperature":false,"configured":true,"contextWindow":0,"inputMediaTypes":null}]`)) + + tools := httptest.NewRecorder() + service.Handler().ServeHTTP(tools, httptest.NewRequest(http.MethodGet, "/api/chat/tools", nil)) + Expect(tools.Code).To(Equal(http.StatusOK)) + var catalog aichat.ToolCatalogResponse + Expect(json.Unmarshal(tools.Body.Bytes(), &catalog)).To(Succeed()) + Expect(catalog.Tools).To(HaveLen(2)) + Expect(catalog.Tools[0].Name).To(Equal("invoice_get")) + Expect(catalog.Tools[0].Source).To(Equal("custom")) + Expect(catalog.Tools[1].Name).To(Equal("docs_search")) + Expect(catalog.Tools[1].Source).To(Equal("mcp")) + }) + + It("maps the HTTP request directly into api.Spec and streams provider events", func() { + provider := &fakeStreamingProvider{events: []api.Event{ + {Kind: api.EventText, Text: "done"}, + {Kind: api.EventResult, Success: true, Model: "test-model"}, + }} + resolver := &fakeResolver{provider: provider} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{ + Spec: api.Spec{Model: api.Model{Name: "openai/test-model"}}, + System: "Use application tools.", + ProviderConfig: api.Config{APIURL: "https://example.com/ai", ProjectName: "tenant-x"}, + }, nil + }), + Attachments: fakeAttachmentResolver{}, + Tools: aichat.ToolProviderFunc(func(context.Context) (aichat.ToolSet, error) { + return aichat.ToolSet{Definitions: []api.ToolDefinition{{ + Name: "invoice_get", Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}}, nil + }), + }) + request := aichat.ChatRequest{ + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{ + {Type: "text", Text: "inspect"}, + {Type: "file", URL: "https://example.com/image.png", Filename: "image.png", MediaType: "image/png"}, + }}}, + Context: "invoice editor", ToolPreferences: api.ToolPreferences{"billing": api.ToolModeAsk}, + ReasoningEffort: api.EffortHigh, PermissionMode: api.PermissionAcceptEdits, + } + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", request)) + + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(response.Header().Get("x-vercel-ai-ui-message-stream")).To(Equal("v1")) + Expect(response.Body.String()).To(ContainSubstring(`"delta":"done"`)) + Expect(provider.specs).To(HaveLen(1)) + spec := provider.specs[0] + Expect(spec.Model.Name).To(Equal("openai/test-model")) + Expect(spec.Model.Effort).To(Equal(api.EffortHigh)) + Expect(spec.ToolPreferences).To(Equal(api.ToolPreferences{"billing": api.ToolModeAsk})) + Expect(spec.Permissions.Mode).To(Equal(api.PermissionAcceptEdits)) + Expect(spec.Messages).To(HaveLen(2)) + Expect(spec.Messages[0]).To(Equal(api.Message{Role: api.RoleSystem, Parts: []api.Part{{ + Type: api.PartText, Text: "Use application tools.\n\nCurrent UI context:\ninvoice editor", + }}})) + Expect(spec.Messages[1].Parts).To(HaveLen(2)) + Expect(spec.Messages[1].Parts[1].Attachment).NotTo(BeNil()) + Expect(spec.Messages[1].Parts[1].Attachment.IsPrepared()).To(BeTrue()) + Expect(resolver.configs).To(HaveLen(1)) + Expect(resolver.configs[0].Tools).To(HaveLen(1)) + Expect(resolver.configs[0].APIURL).To(Equal("https://example.com/ai")) + Expect(resolver.configs[0].ProjectName).To(Equal("tenant-x")) + }) + + It("passes a durable approval resume without rebuilding conversation messages", func() { + state := api.ToolApprovalState{ + Messages: []api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "pay"}}}, + {Role: api.RoleAssistant, Parts: []api.Part{{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: "call-1", Name: "invoice_pay", Input: json.RawMessage(`{"id":"inv-1"}`), + }}}}, + }, + Calls: []api.ToolApprovalCall{{Request: api.ToolApprovalRequest{ + ToolCallID: "call-1", Tool: "invoice_pay", Input: json.RawMessage(`{"id":"inv-1"}`), + }}}, + } + resume := &api.ToolApprovalResume{State: state, Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-1", Tool: "invoice_pay", Action: api.ToolApprovalApprove, + }}} + provider := &fakeStreamingProvider{events: []api.Event{{Kind: api.EventResult, Success: true}}} + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}}) + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + Model: "openai/test-model", ToolApproval: resume, + })) + + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(provider.specs).To(HaveLen(1)) + Expect(provider.specs[0].ToolApproval).To(Equal(resume)) + Expect(provider.specs[0].Messages).To(BeNil()) + }) + + It("serves thread CRUD through the injected persistence store", func() { + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{}, Threads: aichat.NewMemoryThreadStore()}) + create := httptest.NewRecorder() + service.Handler().ServeHTTP(create, requestJSON(http.MethodPost, "/api/chat/threads", map[string]string{"title": "Review"})) + Expect(create.Code).To(Equal(http.StatusCreated)) + var thread aichat.Thread + Expect(json.Unmarshal(create.Body.Bytes(), &thread)).To(Succeed()) + + get := httptest.NewRecorder() + service.Handler().ServeHTTP(get, httptest.NewRequest(http.MethodGet, "/api/chat/threads/"+thread.ID, nil)) + Expect(get.Code).To(Equal(http.StatusOK)) + + remove := httptest.NewRecorder() + service.Handler().ServeHTTP(remove, httptest.NewRequest(http.MethodDelete, "/api/chat/threads/"+thread.ID, nil)) + Expect(remove.Code).To(Equal(http.StatusNoContent)) + }) + + It("persists the completed assistant event stream with tool results and metadata", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Review") + Expect(err).NotTo(HaveOccurred()) + usage := &api.Usage{InputTokens: 12, OutputTokens: 4} + provider := &fakeStreamingProvider{events: []api.Event{ + {Kind: api.EventText, Text: "checking"}, + {Kind: api.EventToolUse, Tool: "invoice_get", ToolCallID: "call-1", Input: map[string]any{"id": "inv-1"}}, + {Kind: api.EventToolResult, Tool: "invoice_get", ToolCallID: "call-1", Text: `{"status":"draft"}`, Success: true}, + {Kind: api.EventResult, Success: true, SessionID: "session-1", Model: "test-model", Usage: usage, CostUSD: 0.25}, + }} + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}, Threads: store}) + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: "chat-1", ThreadID: thread.ID, Model: "openai/test-model", + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "inspect"}}}}, + })) + Expect(response.Code).To(Equal(http.StatusOK)) + + stored, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(stored.Messages).To(HaveLen(2)) + assistant := stored.Messages[1] + Expect(assistant.Role).To(Equal("assistant")) + Expect(assistant.Parts).To(HaveLen(3)) + Expect(assistant.Parts[0].Type).To(Equal("text")) + Expect(assistant.Parts[0].Text).To(Equal("checking")) + Expect(assistant.Parts[1].Type).To(Equal("dynamic-tool")) + Expect(assistant.Parts[1].State).To(Equal("output-available")) + Expect(assistant.Parts[1].ToolCallID).To(Equal("call-1")) + Expect(assistant.Parts[1].Output).To(MatchJSON(`{"status":"draft"}`)) + Expect(assistant.Parts[2].Type).To(Equal("data-result")) + Expect(assistant.Metadata).NotTo(BeNil()) + Expect(assistant.Metadata.ProviderSessionID).To(Equal("session-1")) + Expect(stored.ProviderSessionID).To(Equal("session-1")) + Expect(stored.TotalInputTokens).To(Equal(12)) + Expect(stored.TotalCostUSD).To(Equal(0.25)) + }) + + It("cancels the provider and persistence forwarder when the SSE consumer stops", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Cancellation") + Expect(err).NotTo(HaveOccurred()) + exited := make(chan struct{}) + provider := &fakeStreamingProvider{} + provider.execute = func(ctx context.Context, _ api.Spec) (<-chan api.Event, error) { + events := make(chan api.Event) + go func() { + defer close(exited) + defer close(events) + for _, event := range []api.Event{ + {Kind: api.EventToolUse, Tool: "invoice_get", ToolCallID: "call-1"}, + {Kind: api.EventToolUse, Tool: "invoice_get", ToolCallID: "call-1"}, + {Kind: api.EventText, Text: "must not block"}, + } { + select { + case events <- event: + case <-ctx.Done(): + return + } + } + }() + return events, nil + } + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}, Threads: store}) + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ThreadID: thread.ID, Model: "openai/test-model", + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "inspect"}}}}, + })) + Eventually(exited).Should(BeClosed()) + Expect(response.Body.String()).To(ContainSubstring("persist duplicate tool call")) + }) +}) diff --git a/pkg/aichat/sse.go b/pkg/aichat/sse.go new file mode 100644 index 0000000..5f9c29f --- /dev/null +++ b/pkg/aichat/sse.go @@ -0,0 +1,102 @@ +package aichat + +import ( + "encoding/json" + "fmt" + "net/http" +) + +// Part is one AI SDK v6 UI Message Stream chunk. +type Part struct { + Type string `json:"type"` + + MessageID string `json:"messageId,omitempty"` + ID string `json:"id,omitempty"` + Delta string `json:"delta,omitempty"` + + ToolCallID string `json:"toolCallId,omitempty"` + ToolName string `json:"toolName,omitempty"` + Input any `json:"input,omitempty"` + Output any `json:"output,omitempty"` + Dynamic bool `json:"dynamic,omitempty"` + ApprovalID string `json:"approvalId,omitempty"` + + Data any `json:"data,omitempty"` + ErrorText string `json:"errorText,omitempty"` + MessageMetadata *MessageMetadata `json:"messageMetadata,omitempty"` +} + +// UsageMetadata is Captain usage normalized for frontend message metadata. +type UsageMetadata struct { + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` + ReasoningTokens int `json:"reasoningTokens"` + CacheReadTokens int `json:"cacheReadTokens"` + CacheWriteTokens int `json:"cacheWriteTokens"` + TotalTokens int `json:"totalTokens"` +} + +// MessageMetadata is attached to the assistant UIMessage by the finish part. +type MessageMetadata struct { + ProviderSessionID string `json:"providerSessionId,omitempty"` + Model string `json:"model,omitempty"` + Usage *UsageMetadata `json:"usage,omitempty"` + Cost float64 `json:"cost,omitempty"` + ContextTokens int `json:"contextTokens,omitempty"` + Success *bool `json:"success,omitempty"` +} + +// SSEWriter writes AI SDK v6 chunks using Server-Sent Events framing. +type SSEWriter struct { + w http.ResponseWriter + flusher http.Flusher + done bool +} + +// NewSSEWriter initializes a flushable HTTP response for the UI Message Stream protocol. +func NewSSEWriter(w http.ResponseWriter) (*SSEWriter, error) { + flusher, ok := w.(http.Flusher) + if !ok { + return nil, fmt.Errorf("response writer does not support flushing") + } + headers := w.Header() + headers.Set("Content-Type", "text/event-stream") + headers.Set("x-vercel-ai-ui-message-stream", "v1") + headers.Set("Cache-Control", "no-cache") + headers.Set("Connection", "keep-alive") + headers.Set("x-accel-buffering", "no") + w.WriteHeader(http.StatusOK) + return &SSEWriter{w: w, flusher: flusher}, nil +} + +// WritePart writes and flushes one JSON chunk. +func (w *SSEWriter) WritePart(part Part) error { + if w.done { + return fmt.Errorf("AI SDK stream is already done") + } + if part.Type == "" { + return fmt.Errorf("AI SDK stream part type is required") + } + payload, err := json.Marshal(part) + if err != nil { + return fmt.Errorf("marshal AI SDK stream part %q: %w", part.Type, err) + } + if _, err := fmt.Fprintf(w.w, "data: %s\n\n", payload); err != nil { + return fmt.Errorf("write AI SDK stream part %q: %w", part.Type, err) + } + w.flusher.Flush() + return nil +} + +// Done writes and flushes the literal stream terminator. +func (w *SSEWriter) Done() error { + if w.done { + return fmt.Errorf("AI SDK stream is already done") + } + w.done = true + if _, err := fmt.Fprint(w.w, "data: [DONE]\n\n"); err != nil { + return fmt.Errorf("write AI SDK stream terminator: %w", err) + } + w.flusher.Flush() + return nil +} diff --git a/pkg/aichat/stream_ginkgo_test.go b/pkg/aichat/stream_ginkgo_test.go new file mode 100644 index 0000000..b5819df --- /dev/null +++ b/pkg/aichat/stream_ginkgo_test.go @@ -0,0 +1,265 @@ +package aichat_test + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" +) + +type flushRecorder struct { + *httptest.ResponseRecorder + flushes int +} + +func (r *flushRecorder) Flush() { r.flushes++ } + +type nonFlushWriter struct { + header http.Header +} + +func (w *nonFlushWriter) Header() http.Header { return w.header } +func (w *nonFlushWriter) Write([]byte) (int, error) { return 0, nil } +func (w *nonFlushWriter) WriteHeader(statusCode int) {} + +func recordEvents(events ...api.Event) (*flushRecorder, error) { + recorder := &flushRecorder{ResponseRecorder: httptest.NewRecorder()} + writer, err := aichat.NewSSEWriter(recorder) + if err != nil { + return recorder, err + } + channel := make(chan api.Event, len(events)) + for _, event := range events { + channel <- event + } + close(channel) + return recorder, aichat.WriteEventStream(writer, channel) +} + +func decodedDataLines(body string) []map[string]any { + parts := []map[string]any{} + for line := range strings.SplitSeq(body, "\n") { + if !strings.HasPrefix(line, "data: ") || line == "data: [DONE]" { + continue + } + var part map[string]any + Expect(json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &part)).To(Succeed()) + parts = append(parts, part) + } + return parts +} + +func partTypes(parts []map[string]any) []string { + types := make([]string, len(parts)) + for i, part := range parts { + types[i], _ = part["type"].(string) + } + return types +} + +func pendingApprovalState(calls ...api.ToolApprovalRequest) *api.ToolApprovalState { + parts := make([]api.Part, len(calls)) + stateCalls := make([]api.ToolApprovalCall, len(calls)) + for i, call := range calls { + parts[i] = api.Part{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: call.ToolCallID, Name: call.Tool, Input: call.Input, + }} + stateCalls[i] = api.ToolApprovalCall{Request: call} + } + return &api.ToolApprovalState{ + Messages: []api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "update"}}}, + {Role: api.RoleAssistant, Parts: parts}, + }, + Calls: stateCalls, + } +} + +var _ = Describe("AI SDK v6 event stream", func() { + It("rejects a response writer that cannot stream", func() { + _, err := aichat.NewSSEWriter(&nonFlushWriter{header: http.Header{}}) + Expect(err).To(MatchError("response writer does not support flushing")) + }) + + It("sets the protocol headers and flushes every frame", func() { + recorder, err := recordEvents(api.Event{Kind: api.EventResult, Success: true}) + Expect(err).NotTo(HaveOccurred()) + Expect(recorder.Code).To(Equal(http.StatusOK)) + Expect(recorder.Header().Get("Content-Type")).To(Equal("text/event-stream")) + Expect(recorder.Header().Get("x-vercel-ai-ui-message-stream")).To(Equal("v1")) + Expect(recorder.Header().Get("Cache-Control")).To(Equal("no-cache")) + Expect(recorder.Header().Get("x-accel-buffering")).To(Equal("no")) + Expect(recorder.flushes).To(Equal(6)) + Expect(recorder.Body.String()).To(HaveSuffix("data: [DONE]\n\n")) + }) + + It("keeps every text, reasoning, tool, approval, result, and finish part ordered", func() { + usage := &api.Usage{InputTokens: 100, OutputTokens: 40, ReasoningTokens: 10, CacheReadTokens: 5} + recorder, err := recordEvents( + api.Event{Kind: api.EventSystem, SessionID: "session-1", Model: "claude-sonnet"}, + api.Event{Kind: api.EventThinking, Text: "check"}, + api.Event{Kind: api.EventThinking, Text: "ing"}, + api.Event{Kind: api.EventText, Text: "I will inspect."}, + api.Event{Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_get", Input: map[string]any{"id": "inv-1"}}, + api.Event{Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_get"}, + api.Event{Kind: api.EventToolResult, ToolCallID: "call-1", Tool: "invoice_get", Text: `{"status":"draft"}`, Success: true}, + api.Event{Kind: api.EventText, Text: "It is a draft."}, + api.Event{Kind: api.EventResult, SessionID: "session-1", Model: "claude-sonnet", Usage: usage, CostUSD: 0.0125, Success: true, StructuredData: json.RawMessage(`{"invoiceId":"inv-1"}`)}, + ) + Expect(err).NotTo(HaveOccurred()) + parts := decodedDataLines(recorder.Body.String()) + Expect(partTypes(parts)).To(Equal([]string{ + "start", "start-step", + "reasoning-start", "reasoning-delta", "reasoning-delta", "reasoning-end", + "text-start", "text-delta", "text-end", + "tool-input-available", "tool-approval-request", "tool-output-available", + "text-start", "text-delta", "text-end", + "data-result", "finish-step", "finish", + })) + Expect(parts[2]).To(HaveKeyWithValue("id", "reasoning-0")) + Expect(parts[6]).To(HaveKeyWithValue("id", "text-1")) + Expect(parts[9]).To(SatisfyAll( + HaveKeyWithValue("toolCallId", "call-1"), + HaveKeyWithValue("toolName", "invoice_get"), + HaveKeyWithValue("dynamic", true), + )) + Expect(parts[10]).To(SatisfyAll( + HaveKeyWithValue("approvalId", "call-1"), + HaveKeyWithValue("toolCallId", "call-1"), + )) + Expect(parts[11]["output"]).To(Equal(map[string]any{"output": `{"status":"draft"}`})) + Expect(parts[15]["data"]).To(Equal(map[string]any{"invoiceId": "inv-1"})) + Expect(parts[17]["messageMetadata"]).To(Equal(map[string]any{ + "providerSessionId": "session-1", + "model": "claude-sonnet", + "usage": map[string]any{ + "inputTokens": 100.0, "outputTokens": 40.0, "reasoningTokens": 10.0, + "cacheReadTokens": 5.0, "cacheWriteTokens": 0.0, "totalTokens": 155.0, + }, + "cost": 0.0125, + "contextTokens": 100.0, + "success": true, + })) + }) + + It("turns a Captain error event into a closed, valid UI stream", func() { + recorder, err := recordEvents( + api.Event{Kind: api.EventText, Text: "partial"}, + api.Event{Kind: api.EventError, Error: "provider disconnected"}, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(partTypes(decodedDataLines(recorder.Body.String()))).To(Equal([]string{ + "start", "start-step", "text-start", "text-delta", "text-end", + "error", "finish-step", "finish", + })) + Expect(decodedDataLines(recorder.Body.String())[5]).To(HaveKeyWithValue("errorText", "provider disconnected")) + Expect(recorder.Body.String()).To(HaveSuffix("data: [DONE]\n\n")) + }) + + It("finishes a suspended turn with its approval card still pending", func() { + recorder, err := recordEvents( + api.Event{Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_update", Input: map[string]any{"id": "inv-1"}}, + api.Event{Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_update"}, + ) + Expect(err).NotTo(HaveOccurred()) + Expect(partTypes(decodedDataLines(recorder.Body.String()))).To(Equal([]string{ + "start", "start-step", "tool-input-available", "tool-approval-request", "finish-step", "finish", + })) + }) + + It("carries durable Captain approval state across the AI SDK boundary", func() { + approval := pendingApprovalState(api.ToolApprovalRequest{ + ToolCallID: "call-1", Tool: "invoice_update", Input: json.RawMessage(`{"id":"inv-1"}`), + }) + recorder, err := recordEvents( + api.Event{Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_update", Input: map[string]any{"id": "inv-1"}}, + api.Event{Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_update"}, + api.Event{Kind: api.EventResult, Success: true, ToolApproval: approval}, + ) + Expect(err).NotTo(HaveOccurred()) + parts := decodedDataLines(recorder.Body.String()) + Expect(partTypes(parts)).To(Equal([]string{ + "start", "start-step", "tool-input-available", "tool-approval-request", + "data-tool-approval", "finish-step", "finish", + })) + Expect(parts[4]["data"]).To(HaveKeyWithValue("calls", HaveLen(1))) + }) + + DescribeTable("rejects approval state that does not match the streamed pending tools", + func(state *api.ToolApprovalState, message string) { + recorder, err := recordEvents( + api.Event{Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_update", Input: map[string]any{"id": "inv-1"}}, + api.Event{Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_update"}, + api.Event{Kind: api.EventResult, Success: true, ToolApproval: state}, + ) + Expect(err).To(MatchError(message)) + Expect(partTypes(decodedDataLines(recorder.Body.String()))).To(ContainElement("error")) + }, + Entry("different call ID", pendingApprovalState(api.ToolApprovalRequest{ + ToolCallID: "call-2", Tool: "invoice_update", Input: json.RawMessage(`{"id":"inv-1"}`), + }), `approval state call "call-2" has no matching streamed approval request`), + Entry("different tool name", pendingApprovalState(api.ToolApprovalRequest{ + ToolCallID: "call-1", Tool: "invoice_delete", Input: json.RawMessage(`{"id":"inv-1"}`), + }), `approval state call "call-1" names "invoice_delete", want "invoice_update"`), + Entry("different tool input", pendingApprovalState(api.ToolApprovalRequest{ + ToolCallID: "call-1", Tool: "invoice_update", Input: json.RawMessage(`{"id":"inv-2"}`), + }), `approval state call "call-1" input does not match the streamed tool request`), + Entry("extra pending call", pendingApprovalState( + api.ToolApprovalRequest{ToolCallID: "call-1", Tool: "invoice_update", Input: json.RawMessage(`{"id":"inv-1"}`)}, + api.ToolApprovalRequest{ToolCallID: "call-2", Tool: "invoice_delete", Input: json.RawMessage(`{"id":"inv-2"}`)}, + ), `approval state call "call-2" has no matching streamed approval request`), + ) + + It("rejects an approval state missing a streamed pending tool", func() { + state := pendingApprovalState(api.ToolApprovalRequest{ + ToolCallID: "call-1", Tool: "invoice_update", Input: json.RawMessage(`{"id":"inv-1"}`), + }) + _, err := recordEvents( + api.Event{Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_update", Input: map[string]any{"id": "inv-1"}}, + api.Event{Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_update"}, + api.Event{Kind: api.EventToolUse, ToolCallID: "call-2", Tool: "invoice_delete", Input: map[string]any{"id": "inv-2"}}, + api.Event{Kind: api.EventPermission, ToolCallID: "call-2", Tool: "invoice_delete"}, + api.Event{Kind: api.EventResult, Success: true, ToolApproval: state}, + ) + Expect(err).To(MatchError(`streamed approval request "call-2" is absent from the approval state`)) + }) + + DescribeTable("fails loud and emits an error part for malformed correlation", + func(events []api.Event, message string) { + recorder, err := recordEvents(events...) + Expect(err).To(MatchError(message)) + parts := decodedDataLines(recorder.Body.String()) + Expect(partTypes(parts)).To(ContainElement("error")) + Expect(parts[len(parts)-1]).To(HaveKeyWithValue("type", "finish")) + Expect(recorder.Body.String()).To(HaveSuffix("data: [DONE]\n\n")) + }, + Entry("missing tool call id", []api.Event{{Kind: api.EventToolUse, Tool: "invoice_get"}}, `tool use "invoice_get" has no tool call id`), + Entry("duplicate tool call", []api.Event{ + {Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_get"}, + {Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_get"}, + }, `duplicate tool call id "call-1"`), + Entry("orphan permission", []api.Event{{Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_get"}}, `permission for tool call "call-1" has no matching tool use`), + Entry("orphan result", []api.Event{{Kind: api.EventToolResult, ToolCallID: "call-1", Tool: "invoice_get"}}, `result for tool call "call-1" has no matching tool use`), + Entry("duplicate permission", []api.Event{ + {Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_update"}, + {Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_update"}, + {Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_update"}, + }, `duplicate permission for tool call "call-1"`), + Entry("mismatched tool name", []api.Event{ + {Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_get"}, + {Kind: api.EventToolResult, ToolCallID: "call-1", Tool: "invoice_delete"}, + }, `result for tool call "call-1" names "invoice_delete", want "invoice_get"`), + Entry("terminal result while approval is unresolved", []api.Event{ + {Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_update"}, + {Kind: api.EventPermission, ToolCallID: "call-1", Tool: "invoice_update"}, + {Kind: api.EventResult, Success: true}, + }, `tool call "call-1" ended without a result`), + Entry("dangling tool", []api.Event{{Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_get"}}, `tool call "call-1" ended without a result or approval request`), + ) +}) diff --git a/pkg/aichat/threads.go b/pkg/aichat/threads.go new file mode 100644 index 0000000..c802106 --- /dev/null +++ b/pkg/aichat/threads.go @@ -0,0 +1,154 @@ +package aichat + +import ( + "context" + "fmt" + "sort" + "sync" + "time" +) + +type Thread struct { + ID string `json:"id"` + Title string `json:"title"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` + Messages []UIMessage `json:"messages"` + + TotalInputTokens int `json:"totalInputTokens"` + TotalOutputTokens int `json:"totalOutputTokens"` + TotalReasoningTokens int `json:"totalReasoningTokens"` + TotalCacheReadTokens int `json:"totalCacheReadTokens"` + TotalCacheWriteTokens int `json:"totalCacheWriteTokens"` + TotalCostUSD float64 `json:"totalCostUsd"` + LastContextTokens int `json:"lastContextTokens"` + ProviderSessionID string `json:"providerSessionId,omitempty"` +} + +type TurnUsage struct { + InputTokens int + OutputTokens int + ReasoningTokens int + CacheReadTokens int + CacheWriteTokens int + CostUSD float64 +} + +// ThreadStore is the persistence boundary for chat history, provider session +// identity, and cumulative usage. Implementations must be concurrency-safe. +type ThreadStore interface { + Create(context.Context, string) (*Thread, error) + List(context.Context) ([]*Thread, error) + Get(context.Context, string) (*Thread, error) + AppendMessage(context.Context, string, UIMessage) error + Delete(context.Context, string) error + SetProviderSession(context.Context, string, string) error + AddUsage(context.Context, string, TurnUsage) (*Thread, error) +} + +type memoryThreadStore struct { + mu sync.Mutex + seq int + threads map[string]*Thread +} + +func NewMemoryThreadStore() ThreadStore { + return &memoryThreadStore{threads: map[string]*Thread{}} +} + +func (s *memoryThreadStore) Create(_ context.Context, title string) (*Thread, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.seq++ + now := time.Now() + thread := &Thread{ID: fmt.Sprintf("thread-%d", s.seq), Title: title, CreatedAt: now, UpdatedAt: now, Messages: []UIMessage{}} + s.threads[thread.ID] = thread + return cloneThread(thread), nil +} + +func (s *memoryThreadStore) List(context.Context) ([]*Thread, error) { + s.mu.Lock() + defer s.mu.Unlock() + threads := make([]*Thread, 0, len(s.threads)) + for _, thread := range s.threads { + threads = append(threads, cloneThread(thread)) + } + sort.Slice(threads, func(i, j int) bool { return threads[i].UpdatedAt.After(threads[j].UpdatedAt) }) + return threads, nil +} + +func (s *memoryThreadStore) Get(_ context.Context, id string) (*Thread, error) { + s.mu.Lock() + defer s.mu.Unlock() + thread, ok := s.threads[id] + if !ok { + return nil, fmt.Errorf("thread %q not found", id) + } + return cloneThread(thread), nil +} + +func (s *memoryThreadStore) AppendMessage(_ context.Context, id string, message UIMessage) error { + s.mu.Lock() + defer s.mu.Unlock() + thread, err := s.thread(id) + if err != nil { + return err + } + thread.Messages = append(thread.Messages, message) + thread.UpdatedAt = time.Now() + return nil +} + +func (s *memoryThreadStore) Delete(_ context.Context, id string) error { + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.threads[id]; !ok { + return fmt.Errorf("thread %q not found", id) + } + delete(s.threads, id) + return nil +} + +func (s *memoryThreadStore) SetProviderSession(_ context.Context, id, sessionID string) error { + s.mu.Lock() + defer s.mu.Unlock() + thread, err := s.thread(id) + if err != nil { + return err + } + thread.ProviderSessionID = sessionID + thread.UpdatedAt = time.Now() + return nil +} + +func (s *memoryThreadStore) AddUsage(_ context.Context, id string, usage TurnUsage) (*Thread, error) { + s.mu.Lock() + defer s.mu.Unlock() + thread, err := s.thread(id) + if err != nil { + return nil, err + } + thread.TotalInputTokens += usage.InputTokens + thread.TotalOutputTokens += usage.OutputTokens + thread.TotalReasoningTokens += usage.ReasoningTokens + thread.TotalCacheReadTokens += usage.CacheReadTokens + thread.TotalCacheWriteTokens += usage.CacheWriteTokens + thread.TotalCostUSD += usage.CostUSD + thread.LastContextTokens = usage.InputTokens + thread.UpdatedAt = time.Now() + return cloneThread(thread), nil +} + +func (s *memoryThreadStore) thread(id string) (*Thread, error) { + thread, ok := s.threads[id] + if !ok { + return nil, fmt.Errorf("thread %q not found", id) + } + return thread, nil +} + +func cloneThread(thread *Thread) *Thread { + copy := *thread + copy.Messages = append([]UIMessage(nil), thread.Messages...) + return © +} diff --git a/pkg/aichat/threads_http.go b/pkg/aichat/threads_http.go new file mode 100644 index 0000000..9840cf1 --- /dev/null +++ b/pkg/aichat/threads_http.go @@ -0,0 +1,106 @@ +package aichat + +import ( + "encoding/json" + "fmt" + "io" + "net/http" +) + +func (s *Service) registerThreadRoutes(mux *http.ServeMux) { + mux.HandleFunc("POST /api/chat/threads", s.handleCreateThread) + mux.HandleFunc("GET /api/chat/threads", s.handleListThreads) + mux.HandleFunc("GET /api/chat/threads/{id}", s.handleGetThread) + mux.HandleFunc("DELETE /api/chat/threads/{id}", s.handleDeleteThread) +} + +func (s *Service) threadStore(w http.ResponseWriter) ThreadStore { + if s.options.Threads == nil { + http.Error(w, "thread persistence is not configured", http.StatusNotImplemented) + return nil + } + return s.options.Threads +} + +func (s *Service) handleCreateThread(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w) + if store == nil { + return + } + body := struct { + Title string `json:"title"` + }{} + if request.Body != nil { + if err := json.NewDecoder(request.Body).Decode(&body); err != nil && err != io.EOF { + http.Error(w, fmt.Sprintf("invalid thread request: %v", err), http.StatusBadRequest) + return + } + } + if body.Title == "" { + body.Title = "New conversation" + } + thread, err := store.Create(request.Context(), body.Title) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := writeJSON(w, http.StatusCreated, thread); err != nil { + serviceLog.Errorf("write created chat thread: %v", err) + } +} + +func (s *Service) handleListThreads(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w) + if store == nil { + return + } + threads, err := store.List(request.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := writeJSON(w, http.StatusOK, threads); err != nil { + serviceLog.Errorf("write chat thread list: %v", err) + } +} + +func (s *Service) handleGetThread(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w) + if store == nil { + return + } + thread, err := store.Get(request.Context(), request.PathValue("id")) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + if err := writeJSON(w, http.StatusOK, thread); err != nil { + serviceLog.Errorf("write chat thread %q: %v", request.PathValue("id"), err) + } +} + +func (s *Service) handleDeleteThread(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w) + if store == nil { + return + } + if err := store.Delete(request.Context(), request.PathValue("id")); err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + w.WriteHeader(http.StatusNoContent) +} + +func writeJSON(w http.ResponseWriter, status int, value any) error { + payload, err := json.Marshal(value) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return err + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + if _, err := w.Write(append(payload, '\n')); err != nil { + return fmt.Errorf("write JSON response: %w", err) + } + return nil +} diff --git a/pkg/aichat/tools.go b/pkg/aichat/tools.go new file mode 100644 index 0000000..f25cba1 --- /dev/null +++ b/pkg/aichat/tools.go @@ -0,0 +1,120 @@ +package aichat + +import ( + "context" + "fmt" + + aitools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" +) + +// ToolSet couples executable Captain definitions with their optional canonical +// frontend catalog rows. When Catalog is empty it is projected from Definitions. +type ToolSet struct { + Definitions []api.ToolDefinition + Catalog []aitools.ToolCatalogEntry +} + +// ToolProvider loads an application or MCP tool set for the current request. +type ToolProvider interface { + ToolSet(context.Context) (ToolSet, error) +} + +type ToolProviderFunc func(context.Context) (ToolSet, error) + +func (f ToolProviderFunc) ToolSet(ctx context.Context) (ToolSet, error) { return f(ctx) } + +func StaticToolProvider(definitions []api.ToolDefinition) ToolProvider { + return ToolProviderFunc(func(context.Context) (ToolSet, error) { + return ToolSet{Definitions: append([]api.ToolDefinition(nil), definitions...)}, nil + }) +} + +func CombineToolProviders(providers ...ToolProvider) ToolProvider { + return ToolProviderFunc(func(ctx context.Context) (ToolSet, error) { + combined := ToolSet{} + seenDefinitions := map[string]bool{} + seenCatalog := map[string]bool{} + for i, provider := range providers { + if provider == nil { + continue + } + set, err := provider.ToolSet(ctx) + if err != nil { + return ToolSet{}, fmt.Errorf("load tool provider %d: %w", i+1, err) + } + if err := appendToolSet(&combined, set, "application", seenDefinitions, seenCatalog); err != nil { + return ToolSet{}, err + } + } + return combined, nil + }) +} + +func (s *Service) loadTools(ctx context.Context) (ToolSet, error) { + combined := ToolSet{} + seenDefinitions := map[string]bool{} + seenCatalog := map[string]bool{} + for _, source := range []struct { + name string + provider ToolProvider + }{{name: "custom", provider: s.options.Tools}, {name: "mcp", provider: s.options.MCP}} { + if source.provider == nil { + continue + } + set, err := source.provider.ToolSet(ctx) + if err != nil { + return ToolSet{}, fmt.Errorf("load %s tools: %w", source.name, err) + } + if err := appendToolSet(&combined, set, source.name, seenDefinitions, seenCatalog); err != nil { + return ToolSet{}, err + } + } + return combined, nil +} + +func appendToolSet(combined *ToolSet, set ToolSet, source string, seenDefinitions, seenCatalog map[string]bool) error { + for _, definition := range set.Definitions { + if definition.Name == "" { + return fmt.Errorf("%s tool name is required", source) + } + if definition.Handler == nil { + return fmt.Errorf("%s tool %q handler is required", source, definition.Name) + } + if seenDefinitions[definition.Name] { + return fmt.Errorf("duplicate tool definition %q", definition.Name) + } + seenDefinitions[definition.Name] = true + combined.Definitions = append(combined.Definitions, definition) + } + catalog := set.Catalog + if len(catalog) == 0 { + catalog = make([]aitools.ToolCatalogEntry, 0, len(set.Definitions)) + for _, definition := range set.Definitions { + entry := aitools.CustomCatalogEntry(toolDefinitionForCatalog(definition), definition.Name, definition.InputSchema) + entry.Source = source + catalog = append(catalog, entry) + } + } + for _, entry := range catalog { + if entry.Name == "" { + return fmt.Errorf("%s tool catalog name is required", source) + } + if seenCatalog[entry.Name] { + return fmt.Errorf("duplicate tool catalog entry %q", entry.Name) + } + seenCatalog[entry.Name] = true + combined.Catalog = append(combined.Catalog, entry) + } + return nil +} + +func toolDefinitionForCatalog(definition api.ToolDefinition) aitools.ToolDefinition { + return aitools.ToolDefinition{ + Name: definition.Name, Description: definition.Description, InputSchema: definition.InputSchema, + Group: definition.Group, Parent: definition.Parent, Icon: definition.Icon, + DefaultPermission: definition.DefaultPermission, Strict: definition.Strict, + ReadOnlyHint: definition.ReadOnlyHint, DestructiveHint: definition.DestructiveHint, + IdempotentHint: definition.IdempotentHint, Annotations: definition.Annotations, + } +} diff --git a/pkg/aichat/wire.go b/pkg/aichat/wire.go new file mode 100644 index 0000000..86ce1a9 --- /dev/null +++ b/pkg/aichat/wire.go @@ -0,0 +1,103 @@ +// Package aichat implements the transport-facing AI SDK v6 chat protocol. +package aichat + +import ( + "encoding/json" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" +) + +// ChatRequest is the body posted by the AI SDK DefaultChatTransport. +type ChatRequest struct { + ID string `json:"id,omitempty"` + Messages []UIMessage `json:"messages"` + Model string `json:"model,omitempty"` + ReasoningEffort api.Effort `json:"reasoningEffort,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + Budget api.Budget `json:"budget,omitempty"` + ToolPreferences api.ToolPreferences `json:"toolPreferences,omitempty"` + PermissionMode api.PermissionMode `json:"permissionMode,omitempty"` + ToolApproval *api.ToolApprovalResume `json:"toolApproval,omitempty"` + + Context string `json:"context,omitempty"` + ContextItems []ChatContextItem `json:"contextItems,omitempty"` + + ThreadID string `json:"threadId,omitempty"` + ProviderSessionID string `json:"providerSessionId,omitempty"` +} + +// ChatContextItem carries app-owned structured state alongside its readable label. +type ChatContextItem struct { + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Label string `json:"label,omitempty"` + Fields map[string]string `json:"fields,omitempty"` + Payload json.RawMessage `json:"payload,omitempty"` +} + +// UIMessage is the AI SDK v6 message wire shape. +type UIMessage struct { + ID string `json:"id,omitempty"` + Role string `json:"role"` + Parts []UIPart `json:"parts"` + Metadata *MessageMetadata `json:"metadata,omitempty"` +} + +// UIPart models the input part variants Captain consumes. +type UIPart struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + + MediaType string `json:"mediaType,omitempty"` + URL string `json:"url,omitempty"` + Filename string `json:"filename,omitempty"` + AttachmentID string `json:"attachmentId,omitempty"` + + ToolName string `json:"toolName,omitempty"` + ToolCallID string `json:"toolCallId,omitempty"` + State string `json:"state,omitempty"` + Input json.RawMessage `json:"input,omitempty"` + Output json.RawMessage `json:"output,omitempty"` + ErrorText string `json:"errorText,omitempty"` + Data json.RawMessage `json:"data,omitempty"` + Approval *Approval `json:"approval,omitempty"` +} + +// Approval is the AI SDK tool approval envelope attached to a tool UI part. +type Approval struct { + ID string `json:"id"` + Approved *bool `json:"approved,omitempty"` + Reason string `json:"reason,omitempty"` +} + +// IsTool reports whether the part is a static or dynamic tool part. +func (p UIPart) IsTool() bool { + return p.Type == "dynamic-tool" || strings.HasPrefix(p.Type, "tool-") +} + +// EffectiveToolName returns a dynamic tool's explicit name or a static tool's type suffix. +func (p UIPart) EffectiveToolName() string { + if p.ToolName != "" { + return p.ToolName + } + name, ok := strings.CutPrefix(p.Type, "tool-") + if !ok { + return "" + } + return name +} + +// ModelCatalogResponse reuses Captain's canonical frontend model catalog rows. +type ModelCatalogResponse = []ai.ModelInfo + +// ModelCatalogEntry is Captain's canonical frontend model catalog row. +type ModelCatalogEntry = ai.ModelInfo + +// ToolCatalogResponse reuses Captain's canonical frontend tool catalog. +type ToolCatalogResponse = tools.ToolCatalog + +// ToolCatalogEntry is Captain's canonical frontend tool catalog row. +type ToolCatalogEntry = tools.ToolCatalogEntry diff --git a/pkg/aichat/wire_ginkgo_test.go b/pkg/aichat/wire_ginkgo_test.go new file mode 100644 index 0000000..b7a6090 --- /dev/null +++ b/pkg/aichat/wire_ginkgo_test.go @@ -0,0 +1,93 @@ +package aichat_test + +import ( + "encoding/json" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" +) + +var _ = Describe("AI SDK v6 wire types", func() { + It("decodes the DefaultChatTransport request without losing UI parts", func() { + const body = `{ + "id":"chat-1", + "messages":[{"id":"message-1","role":"assistant","parts":[ + {"type":"reasoning","text":"checking"}, + {"type":"dynamic-tool","toolName":"invoice_get","toolCallId":"call-1","state":"approval-responded","input":{"id":"inv-1"},"approval":{"id":"approval-1","approved":true}}, + {"type":"file","mediaType":"application/pdf","url":"https://example.com/invoice.pdf","filename":"invoice.pdf","attachmentId":"att_sha256_abc"} + ]}], + "model":"anthropic/claude-sonnet", + "reasoningEffort":"high", + "temperature":0, + "budget":{"cost":1.5,"maxTokens":2048,"maxTurns":4}, + "toolPreferences":{"billing":"ask","invoice_get":"on"}, + "context":"invoice editor", + "contextItems":[{"id":"inv-1","type":"invoice","label":"Invoice 1","fields":{"status":"draft"},"payload":{"total":25}}], + "threadId":"thread-1", + "providerSessionId":"session-1", + "permissionMode":"acceptEdits" + }` + + var request aichat.ChatRequest + Expect(json.Unmarshal([]byte(body), &request)).To(Succeed()) + Expect(request.ID).To(Equal("chat-1")) + Expect(request.Model).To(Equal("anthropic/claude-sonnet")) + Expect(request.ReasoningEffort).To(Equal(api.EffortHigh)) + Expect(request.Temperature).NotTo(BeNil()) + Expect(*request.Temperature).To(Equal(0.0)) + Expect(request.Budget).To(Equal(api.Budget{Cost: 1.5, MaxTokens: 2048, MaxTurns: 4})) + Expect(request.ToolPreferences).To(Equal(api.ToolPreferences{ + "billing": api.ToolModeAsk, + "invoice_get": api.ToolModeOn, + })) + Expect(request.PermissionMode).To(Equal(api.PermissionAcceptEdits)) + Expect(request.ToolApproval).To(BeNil()) + Expect(request.ThreadID).To(Equal("thread-1")) + Expect(request.ProviderSessionID).To(Equal("session-1")) + Expect(request.ContextItems).To(HaveLen(1)) + Expect(request.ContextItems[0].Payload).To(MatchJSON(`{"total":25}`)) + + Expect(request.Messages).To(HaveLen(1)) + Expect(request.Messages[0].Parts).To(HaveLen(3)) + tool := request.Messages[0].Parts[1] + Expect(tool.IsTool()).To(BeTrue()) + Expect(tool.EffectiveToolName()).To(Equal("invoice_get")) + Expect(tool.Approval).NotTo(BeNil()) + Expect(tool.Approval.Approved).NotTo(BeNil()) + Expect(*tool.Approval.Approved).To(BeTrue()) + Expect(aichat.UIPart{Type: "tool-static_lookup"}.EffectiveToolName()).To(Equal("static_lookup")) + Expect(aichat.UIPart{Type: "text"}.EffectiveToolName()).To(BeEmpty()) + }) + + It("rejects the removed string tool approval policy", func() { + var request aichat.ChatRequest + Expect(json.Unmarshal([]byte(`{"messages":[],"toolApproval":"manual"}`), &request)).To( + MatchError(ContainSubstring("cannot unmarshal string")), + ) + }) + + It("publishes stable model and tool catalog response shapes", func() { + strict := true + models := aichat.ModelCatalogResponse{{ + ID: "openai/gpt", Provider: "openai", Label: "GPT", Reasoning: true, + Temperature: true, Configured: true, ContextWindow: 128000, + InputMediaTypes: []string{"image/*"}, + }} + tools := aichat.ToolCatalogResponse{Tools: []aichat.ToolCatalogEntry{{ + Name: "invoice_get", Source: "custom", Group: "billing", + PreferenceKey: "billing", DefaultPermission: api.ToolModeAsk, + Strict: &strict, Method: "GET", Path: "/invoices/{id}", + OperationName: "invoice get", InputSchema: map[string]any{"type": "object"}, + }}} + + modelJSON, err := json.Marshal(models) + Expect(err).NotTo(HaveOccurred()) + Expect(modelJSON).To(MatchJSON(`[{"id":"openai/gpt","provider":"openai","label":"GPT","reasoning":true,"temperature":true,"configured":true,"contextWindow":128000,"inputMediaTypes":["image/*"]}]`)) + toolJSON, err := json.Marshal(tools) + Expect(err).NotTo(HaveOccurred()) + Expect(toolJSON).To(MatchJSON(`{"tools":[{"name":"invoice_get","source":"custom","group":"billing","preferenceKey":"billing","defaultPermission":"ask","strict":true,"method":"GET","path":"/invoices/{id}","operationName":"invoice get","inputSchema":{"type":"object"}}]}`)) + }) +}) diff --git a/pkg/aiflags/defaults.go b/pkg/aiflags/defaults.go new file mode 100644 index 0000000..c7f622c --- /dev/null +++ b/pkg/aiflags/defaults.go @@ -0,0 +1,185 @@ +package aiflags + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/api/registry" + "github.com/flanksource/captain/pkg/captainconfig" +) + +// ProviderDefaultView is one provider's effective saved defaults, with the legacy +// flat config projected in and validated. +type ProviderDefaultView struct { + Agent string `json:"agent"` + Model string `json:"model"` + Effort string `json:"effort"` + Configured bool `json:"configured"` +} + +// LoadDefaults reads the saved AI defaults from ~/.captain.yaml. +// +// It reports a broken config rather than swallowing it: whether to degrade to zero +// defaults and carry on is a CLI policy, not a library's call. captain's own +// commands make that choice in pkg/cli (loadSavedAI warns and continues); callers +// wanting the same behaviour pass their own AIDefaults to ResolveWith. +// +// Deliberately no logger here — commons/logger pulls ~55 packages (prometheus, +// fsnotify, …) and would cost this leaf its entire reason for existing. +func LoadDefaults() (captainconfig.AIDefaults, error) { + cfg, _, err := captainconfig.Load() + if err != nil { + return captainconfig.AIDefaults{}, err + } + return cfg.AI, nil +} + +// EffectiveDefaults resolves one provider's saved defaults: the per-provider block +// back-filled from the legacy flat keys, with the agent defaulting to the provider +// itself and the model to that agent's built-in default. +func EffectiveDefaults(saved captainconfig.AIDefaults, provider registry.Backend) (ProviderDefaultView, error) { + if provider.Provider() != provider { + return ProviderDefaultView{}, fmt.Errorf("invalid provider %q", provider) + } + configured, exists := saved.Providers[string(provider)] + legacy := saved.Provider(string(provider)) + if configured.Agent == "" { + configured.Agent = legacy.Agent + } + if configured.Model == "" { + configured.Model = legacy.Model + } + if configured.ReasoningEffort == "" { + configured.ReasoningEffort = legacy.ReasoningEffort + } + agent := registry.Backend(strings.TrimSpace(configured.Agent)) + if agent == "" { + agent = provider + } + if agent.Provider() != provider { + return ProviderDefaultView{}, fmt.Errorf("agent %q does not belong to provider %q", agent, provider) + } + model := strings.TrimSpace(configured.Model) + if model == "" { + model = DefaultModelFor(agent) + } + effort := registry.Effort(strings.TrimSpace(configured.ReasoningEffort)) + if err := registry.ValidateEffort(agent, model, effort); err != nil { + return ProviderDefaultView{}, err + } + return ProviderDefaultView{ + Agent: string(agent), Model: model, Effort: string(effort), Configured: exists, + }, nil +} + +// ApplyDefaults fills a model's unset fields from the saved per-provider defaults, +// primary and fallbacks alike. It expects an already-expanded model (see the +// package doc) and does not resolve — the caller resolves once, afterwards. +func ApplyDefaults(model registry.Model, saved captainconfig.AIDefaults) (registry.Model, error) { + var err error + if strings.TrimSpace(model.Name) != "" { + if model, err = model.Expand(); err != nil { + return registry.Model{}, err + } + } + if model, err = applyCandidateDefaults(model, saved, true); err != nil { + return registry.Model{}, err + } + for i := range model.Fallbacks { + fallback := model.Fallbacks[i] + if fallback, err = fallback.Expand(); err != nil { + return registry.Model{}, fmt.Errorf("fallback[%d]: %w", i, err) + } + // allowActive=false: a fallback must not silently become the active + // provider's model — that would make the fallback chain a no-op. + if fallback, err = applyCandidateDefaults(fallback, saved, false); err != nil { + return registry.Model{}, fmt.Errorf("fallback[%d]: %w", i, err) + } + model.Fallbacks[i] = fallback + } + return model, nil +} + +func applyCandidateDefaults(model registry.Model, saved captainconfig.AIDefaults, allowActive bool) (registry.Model, error) { + provider := model.Backend.Provider() + if provider == "" && strings.TrimSpace(model.Name) != "" { + backend, err := registry.InferBackend(model.Name) + if err != nil { + return registry.Model{}, err + } + provider = backend.Provider() + } + if provider == "" && allowActive { + provider = registry.Backend(saved.ActiveProvider()) + } + if provider == "" { + return registry.Model{}, fmt.Errorf("provider cannot be resolved for model %q", model.Name) + } + defaults, err := EffectiveDefaults(saved, provider) + if err != nil { + return registry.Model{}, err + } + // An explicit --mode owns the mechanism. Taking the saved agent here would make + // `--mode cli` against a saved claude-agent default fail as "mode cli + // contradicts backend claude-agent" — a contradiction the user never wrote. + if model.Backend == "" && model.Mode == "" { + model.Backend = registry.Backend(defaults.Agent) + } + if strings.TrimSpace(model.Name) == "" { + model.Name = defaults.Model + } + if model.Effort == registry.EffortNone { + model.Effort = registry.Effort(defaults.Effort) + } + // Permissive when the backend is still unset (--mode without --backend): the + // authoritative effort check runs against the real backend during resolution. + if err := registry.ValidateEffort(model.Backend, model.Name, model.Effort); err != nil { + return registry.Model{}, err + } + return model, nil +} + +// AllProviderDefaults resolves every provider's effective defaults. +func AllProviderDefaults(saved captainconfig.AIDefaults) (map[string]ProviderDefaultView, error) { + out := make(map[string]ProviderDefaultView, len(registry.Providers())) + for _, p := range registry.Providers() { + provider, err := p.BackendFor(registry.ModeAPI) + if err != nil { + return nil, err + } + defaults, err := EffectiveDefaults(saved, provider) + if err != nil { + return nil, err + } + out[string(provider)] = defaults + } + return out, nil +} + +// DefaultModelFor returns a hard-coded picker default per backend that seeds the +// form. CLI/agent backends use exact provider model IDs from the catalog so the +// seeded default is a selectable option. API backends have no "default" flag on +// /v1/models, so we use the most-current id we expect each provider to keep +// stable; the user can pick anything else. +// +// This stays a hand-maintained table rather than "the catalog's newest preferred +// model": these are the values `captain configure` seeds and users have saved, and +// deriving them would silently move every unconfigured user's default whenever the +// catalog snapshot updates. +func DefaultModelFor(b registry.Backend) string { + switch b { + case registry.BackendAnthropic: + return "claude-sonnet-5" + case registry.BackendClaudeCLI, registry.BackendClaudeAgent, registry.BackendClaudeCmux: + return "claude-sonnet-5" + case registry.BackendOpenAI: + return "gpt-5.6" + case registry.BackendDeepSeek: + return "deepseek-reasoner" + case registry.BackendCodexCLI, registry.BackendCodexAgent, registry.BackendCodexCmux: + return "gpt-5.6-sol" + case registry.BackendGemini, registry.BackendGeminiCLI: + return "gemini-3.5-flash" + } + return "" +} diff --git a/pkg/aiflags/flags.go b/pkg/aiflags/flags.go new file mode 100644 index 0000000..aba66a6 --- /dev/null +++ b/pkg/aiflags/flags.go @@ -0,0 +1,206 @@ +// Package aiflags is captain's model-selection flag surface, extracted so any +// clicky-based CLI can embed it and get captain's exact model semantics. +// +// Its contract with clicky is the struct tags — there is no interface to satisfy; +// embed ModelFlags anonymously and the flags bind themselves. +// +// It is deliberately a leaf: registry (the model catalog) and captainconfig +// (~/.captain.yaml) are its only captain imports, and both are yaml-only. It must +// never import pkg/api, pkg/ai, or pkg/cli — pkg/cli alone pulls ~1000 packages +// (k8s, AWS, Postgres), which no downstream should inherit to parse `--model`. +// That is affordable because api.Model is a type ALIAS for registry.Model, so the +// Model returned here is interchangeable with api.Model at every call site. +// +// # The invariant +// +// Expand, then Merge, then Resolve — once, at the end. +// +// Merging an unexpanded override onto a base silently keeps the base's backend: +// base{Backend: claude-agent}.Merge(Model{Name: "api:opus"}) still says +// claude-agent, and resolution then fails or, worse, runs the wrong runtime. +// Expand sets Backend only when the string carries a prefix, so a bare +// `--model opus` correctly inherits the base's backend while `api:opus` +// correctly overrides it. Every ladder in this package and its callers follows +// that order; departing from it is how a model string loses its mode again. +package aiflags + +import ( + "fmt" + "strconv" + "strings" + + "github.com/flanksource/captain/pkg/api/registry" + "github.com/flanksource/captain/pkg/captainconfig" +) + +// ModelFlags is the model-selection flag surface: the 1:1 flag projection of +// registry.Model and nothing else. Session ids, API keys and budgets are +// deliberately absent — they belong to the request/config, not to a model, and +// keeping them out is what lets any command embed this without inheriting flags +// it has no business owning. +// +// Embed it ANONYMOUSLY — clicky's binder only recurses anonymous fields (it does +// so at any depth), and a named field would silently bind nothing. +// +// Every field is string/[]string/bool on purpose: clicky's binder switches on the +// concrete type and registers NO flag at all for a named type, so Effort and Mode +// cannot be registry.Effort/registry.RuntimeMode here. They are parsed and +// validated in Resolve instead. +// +// Flag names are global to a command — clicky does not prefix embedded structs — +// so this struct owns --model, --fallback, --backend, --mode, --effort, +// --temperature and --no-cache outright. An embedder that redeclares any of them +// panics cobra at init. +// +// It deliberately claims NO single-letter shorthands. Shorthands are per-command +// UX and the letters are scarce: -m here meant `gavel commit` could not embed this +// at all, because -m is its --message (the git convention). A struct meant to be +// embedded everywhere cannot squat them. +// +// No field carries a `default:` tag: defaults only materialize through cobra +// binding, so a directly-constructed ModelFlags would disagree with a bound one. +// Zero means unset, everywhere. +type ModelFlags struct { + Model string `flag:"model" help:"Model name(s), e.g. claude-sonnet-5, a compact selector like agent:opus:high, or a comma-separated primary,fallback list (defaults to the value saved by 'captain configure')"` + Fallback []string `flag:"fallback" help:"Model to try if the primary is unavailable (repeatable; comma-separated allowed)"` + Backend string `flag:"backend" help:"Force backend: anthropic|gemini|openai|deepseek|claude-cli|claude-agent|claude-cmux|codex-cli|codex-agent|codex-cmux|gemini-cli (default: inferred from model or saved by 'captain configure')"` + Mode string `flag:"mode" help:"Runtime mechanism: api|cli|agent|cmux (sdk aliases agent). Combined with the model's family to pick a backend; contradicting --backend or a mode prefix on --model fails loud"` + Effort string `flag:"effort" help:"Reasoning effort: low|medium|high|xhigh|max|ultra (model-dependent)"` + Temperature string `flag:"temperature" help:"Sampling temperature (0.0-2.0)"` + NoCache bool `flag:"no-cache" help:"Disable response caching"` +} + +// ToModel projects the flags onto a Model and expands any compact selector, but +// does NOT resolve it against the catalog. That is the merge-safe form: callers +// layering flags over a spec must Merge before resolving (see the package doc). +func (f ModelFlags) ToModel() (registry.Model, error) { + m := registry.Model{ + Name: strings.TrimSpace(f.Model), + Backend: registry.Backend(strings.TrimSpace(f.Backend)), + Effort: registry.Effort(strings.TrimSpace(f.Effort)), + NoCache: f.NoCache, + } + if m.Backend != "" && !m.Backend.Valid() { + return registry.Model{}, fmt.Errorf("invalid --backend %q (valid: %s)", f.Backend, registry.BackendList()) + } + if err := m.Effort.Validate(); err != nil { + return registry.Model{}, fmt.Errorf("invalid --effort %q: %w", f.Effort, err) + } + // Normalize the mode BEFORE it can reach Model.validateMode: that check + // compares Mode against Backend.Mode() by equality, so an un-normalized "sdk" + // alongside backend claude-agent would report a contradiction that isn't real. + if s := strings.TrimSpace(f.Mode); s != "" { + mode, ok := registry.ParseRuntimeMode(s) + if !ok { + return registry.Model{}, fmt.Errorf("invalid --mode %q (valid: %s)", f.Mode, registry.RuntimeModeList()) + } + m.Mode = mode + } + temp, err := f.Temp() + if err != nil { + return registry.Model{}, err + } + m.Temperature = temp + m.Fallbacks = FallbackModels(f.Fallback) + + if m.Name != "" { + if m, err = m.Expand(); err != nil { + return registry.Model{}, err + } + } + return m, nil +} + +// Resolve is the one-call path: flags + ~/.captain.yaml → a fully resolved Model. +// A broken config surfaces as an error; callers that prefer to warn and carry on +// with zero defaults load their own and use ResolveWith. +func (f ModelFlags) Resolve() (registry.Model, error) { + saved, err := LoadDefaults() + if err != nil { + return registry.Model{}, err + } + return f.ResolveWith(saved) +} + +// ResolveWith is the pure core: no ambient I/O, no globals. Saved defaults arrive +// as a parameter so tests and spec-overlaying callers can drive it directly. +func (f ModelFlags) ResolveWith(saved captainconfig.AIDefaults) (registry.Model, error) { + m, err := f.ToModel() + if err != nil { + return registry.Model{}, err + } + if !f.NoCache && saved.NoCache { + m.NoCache = true + } + if m, err = ApplyDefaults(m, saved); err != nil { + return registry.Model{}, err + } + return registry.ResolveModel(m) +} + +// Overlay layers the flags over an already-structured base (a spec's model), +// flags winning field by field, and resolves the result once. +// +// This is the entry point for callers that have both a config/spec and flags — +// which is most of them. Doing it by hand is how a caller ends up merging an +// unexpanded name onto a populated backend and losing the mode. +func (f ModelFlags) Overlay(base registry.Model) (registry.Model, error) { + saved, err := LoadDefaults() + if err != nil { + return registry.Model{}, err + } + return f.OverlayWith(base, saved) +} + +// OverlayWith is Overlay with saved defaults injected. +func (f ModelFlags) OverlayWith(base registry.Model, saved captainconfig.AIDefaults) (registry.Model, error) { + over, err := f.ToModel() + if err != nil { + return registry.Model{}, err + } + merged, err := ApplyDefaults(base.Merge(over), saved) + if err != nil { + return registry.Model{}, err + } + return registry.ResolveModel(merged) +} + +// Temp parses --temperature. A nil result means unset: an explicit 0 and "unset" +// must stay distinguishable, so providers can tell "sample deterministically" +// from "use the model's default". +func (f ModelFlags) Temp() (*float64, error) { + v, err := parseFloat("temperature", f.Temperature) + if err != nil || v == 0 { + return nil, err + } + if v < 0 || v > 2 { + return nil, fmt.Errorf("invalid --temperature %v (valid: 0.0-2.0)", v) + } + return &v, nil +} + +// FallbackModels turns repeatable/comma-separated --fallback values into models. +func FallbackModels(flags []string) registry.ModelList { + var out registry.ModelList + for _, flag := range flags { + for _, name := range strings.Split(flag, ",") { + if name = strings.TrimSpace(name); name != "" { + out = append(out, registry.Model{Name: name}) + } + } + } + return out +} + +// parseFloat parses a numeric string flag, failing loud rather than silently +// coercing malformed input to zero. +func parseFloat(name, val string) (float64, error) { + if strings.TrimSpace(val) == "" { + return 0, nil + } + f, err := strconv.ParseFloat(val, 64) + if err != nil { + return 0, fmt.Errorf("invalid --%s %q: %w", name, val, err) + } + return f, nil +} diff --git a/pkg/api/aliases.go b/pkg/api/aliases.go new file mode 100644 index 0000000..7ad3856 --- /dev/null +++ b/pkg/api/aliases.go @@ -0,0 +1,66 @@ +package api + +import "github.com/flanksource/captain/pkg/api/registry" + +// Model identity lives in pkg/api/registry — a leaf package, because decoding a +// Spec parses model strings and the parser therefore cannot sit above pkg/api. +// These aliases keep api.Model / api.Backend / api.Effort the names the rest of +// captain uses: an alias is the identical type, methods included, so nothing +// downstream (including pkg/ai's own re-exports) has to know registry exists. + +type ( + Backend = registry.Backend + Effort = registry.Effort + Model = registry.Model + ModelList = registry.ModelList +) + +const ( + BackendAnthropic = registry.BackendAnthropic + BackendGemini = registry.BackendGemini + BackendOpenAI = registry.BackendOpenAI + BackendDeepSeek = registry.BackendDeepSeek + BackendClaudeCLI = registry.BackendClaudeCLI + BackendCodexCLI = registry.BackendCodexCLI + BackendGeminiCLI = registry.BackendGeminiCLI + BackendClaudeAgent = registry.BackendClaudeAgent + BackendCodexAgent = registry.BackendCodexAgent + BackendClaudeCmux = registry.BackendClaudeCmux + BackendCodexCmux = registry.BackendCodexCmux + + AnthropicProvider = registry.AnthropicProvider + OpenAIProvider = registry.OpenAIProvider + GeminiProvider = registry.GeminiProvider + DeepSeekProvider = registry.DeepSeekProvider + + EffortNone = registry.EffortNone + EffortLow = registry.EffortLow + EffortMedium = registry.EffortMedium + EffortHigh = registry.EffortHigh + EffortXHigh = registry.EffortXHigh + EffortMax = registry.EffortMax + EffortUltra = registry.EffortUltra + + CodexAutoReviewModel = registry.CodexAutoReviewModel +) + +// ErrInferBackend marks the "can't infer a backend from this model name" failure +// so callers can enrich it (e.g. with "did you mean" model suggestions). +var ErrInferBackend = registry.ErrInferBackend + +// AllBackends lists every supported backend in canonical order. +func AllBackends() []Backend { return registry.AllBackends() } + +// BackendList renders AllBackends as a comma-separated string for help/error text. +func BackendList() string { return registry.BackendList() } + +// AuthEnvVars returns the environment variables consulted for a backend's API +// key, in priority order. +func AuthEnvVars(b Backend) []string { return registry.AuthEnvVars(b) } + +// InferBackend resolves the backend from a model name prefix, failing loud when +// the name matches nothing. +func InferBackend(model string) (Backend, error) { return registry.InferBackend(model) } + +// AllEfforts lists the non-empty effort tiers in ascending order. +func AllEfforts() []Effort { return registry.AllEfforts() } diff --git a/pkg/api/attachment.go b/pkg/api/attachment.go new file mode 100644 index 0000000..3bcc4e1 --- /dev/null +++ b/pkg/api/attachment.go @@ -0,0 +1,97 @@ +package api + +import ( + "encoding/hex" + "fmt" + "net/url" + "strings" +) + +const AttachmentIDPrefix = "sha256:" + +// AttachmentRef is the serializable reference to one multimodal prompt input. +// Exactly one source is present before resolution; resolved references use ID. +type AttachmentRef struct { + ID string `json:"id,omitempty" yaml:"id,omitempty" pretty:"label=ID"` + Path string `json:"path,omitempty" yaml:"path,omitempty" pretty:"label=Path"` + URL string `json:"url,omitempty" yaml:"url,omitempty" pretty:"label=URL"` + Filename string `json:"filename,omitempty" yaml:"filename,omitempty" pretty:"label=Filename"` + MediaType string `json:"mediaType,omitempty" yaml:"mediaType,omitempty" pretty:"label=Media Type"` + Size int64 `json:"size,omitempty" yaml:"size,omitempty" pretty:"label=Size"` + SHA256 string `json:"sha256,omitempty" yaml:"sha256,omitempty" pretty:"label=SHA-256"` + + preparedBytes []byte + preparedPath string +} + +// AttachmentContent is runtime-only immutable attachment content. +type AttachmentContent struct { + Bytes []byte + Path string +} + +func (a AttachmentRef) Validate() error { + sources := 0 + for _, source := range []string{a.ID, a.Path, a.URL} { + if strings.TrimSpace(source) != "" { + sources++ + } + } + if sources != 1 { + return fmt.Errorf("attachment requires exactly one source: id, path, or url") + } + if a.ID != "" { + if err := validateAttachmentID(a.ID); err != nil { + return err + } + } + if a.URL != "" { + u, err := url.Parse(a.URL) + if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") { + return fmt.Errorf("attachment url must use http or https: %q", a.URL) + } + } + if a.Size < 0 { + return fmt.Errorf("attachment size cannot be negative") + } + if a.SHA256 != "" { + if err := validateSHA256(a.SHA256); err != nil { + return err + } + } + return nil +} + +func (a AttachmentRef) IsPrepared() bool { + return a.ID != "" && (len(a.preparedBytes) > 0 || a.preparedPath != "") +} + +func (a AttachmentRef) PreparedContent() (AttachmentContent, bool) { + if !a.IsPrepared() { + return AttachmentContent{}, false + } + return AttachmentContent{Bytes: a.preparedBytes, Path: a.preparedPath}, true +} + +func (a AttachmentRef) WithPreparedContent(content AttachmentContent) AttachmentRef { + a.preparedBytes = content.Bytes + a.preparedPath = content.Path + return a +} + +func validateAttachmentID(id string) error { + if !strings.HasPrefix(id, AttachmentIDPrefix) { + return fmt.Errorf("attachment id must start with %q", AttachmentIDPrefix) + } + return validateSHA256(strings.TrimPrefix(id, AttachmentIDPrefix)) +} + +func validateSHA256(value string) error { + if len(value) != 64 { + return fmt.Errorf("attachment sha256 must contain 64 hexadecimal characters") + } + if _, err := hex.DecodeString(value); err != nil { + return fmt.Errorf("attachment sha256 is invalid: %w", err) + } + return nil +} diff --git a/pkg/api/attachment_ginkgo_test.go b/pkg/api/attachment_ginkgo_test.go new file mode 100644 index 0000000..fcfc470 --- /dev/null +++ b/pkg/api/attachment_ginkgo_test.go @@ -0,0 +1,44 @@ +package api_test + +import ( + "strings" + + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("AttachmentRef", func() { + It("requires exactly one source", func() { + ref := api.AttachmentRef{Path: "invoice.pdf", URL: "https://example.com/invoice.pdf"} + Expect(ref.Validate()).To(MatchError(ContainSubstring("exactly one source"))) + }) + + It("accepts a durable attachment id", func() { + ref := api.AttachmentRef{ID: "sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"} + Expect(ref.Validate()).To(Succeed()) + }) +}) + +var _ = Describe("Prompt attachment validation", func() { + It("accepts an attachment-only prompt", func() { + prompt := api.Prompt{Attachments: []api.AttachmentRef{{Path: "invoice.pdf"}}} + Expect(prompt.Validate()).To(Succeed()) + }) + + It("does not classify an attachment-only verification request as verify-only", func() { + spec := api.Spec{ + Prompt: api.Prompt{Attachments: []api.AttachmentRef{{Path: "invoice.pdf"}}}, + Workflow: &api.Workflow{Verify: &api.Verify{}}, + } + Expect(spec.IsVerifyOnly()).To(BeFalse()) + }) + + It("includes durable attachment descriptors in the cache identity", func() { + first := api.Prompt{User: "compare", Attachments: []api.AttachmentRef{{ID: api.AttachmentIDPrefix + strings.Repeat("a", 64), MediaType: "image/png"}}} + second := api.Prompt{User: "compare", Attachments: []api.AttachmentRef{{ID: api.AttachmentIDPrefix + strings.Repeat("b", 64), MediaType: "image/png"}}} + Expect(first.CacheIdentity()).NotTo(Equal(second.CacheIdentity())) + Expect(first.CacheIdentity()).NotTo(ContainSubstring("base64")) + }) +}) diff --git a/pkg/api/attachment_suite_test.go b/pkg/api/attachment_suite_test.go new file mode 100644 index 0000000..96fb3c3 --- /dev/null +++ b/pkg/api/attachment_suite_test.go @@ -0,0 +1,13 @@ +package api_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAttachments(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "API Attachments Suite") +} diff --git a/pkg/api/cost.go b/pkg/api/cost.go index 981d781..c67c8f8 100644 --- a/pkg/api/cost.go +++ b/pkg/api/cost.go @@ -2,6 +2,14 @@ package api // Usage is the token breakdown for one or more model calls. Canonical home; // pkg/ai re-exports it via `type Usage = api.Usage`. +// +// The buckets are DISJOINT and additive: InputTokens excludes cache reads, +// OutputTokens excludes reasoning, and TotalTokens is their plain sum. Providers +// whose wire format reports overlapping counts (OpenAI/Codex fold cache into +// prompt tokens and reasoning into completion tokens; Gemini folds cache into +// prompt tokens) MUST normalize at their parse boundary via NetInputTokens / +// NetOutputTokens before populating this struct, so cost pricing and totals do +// not double-count. Anthropic already reports disjoint buckets natively. type Usage struct { InputTokens int `json:"inputTokens" yaml:"inputTokens" pretty:"label=Input,table"` OutputTokens int `json:"outputTokens" yaml:"outputTokens" pretty:"label=Output,table"` @@ -10,11 +18,34 @@ type Usage struct { CacheWriteTokens int `json:"cacheWriteTokens,omitempty" yaml:"cacheWriteTokens,omitempty" pretty:"label=Cache Write,table"` } -// TotalTokens sums every token bucket. +// TotalTokens sums every token bucket. Correct only under the disjoint-bucket +// invariant documented on Usage. func (u Usage) TotalTokens() int { return u.InputTokens + u.OutputTokens + u.ReasoningTokens + u.CacheReadTokens + u.CacheWriteTokens } +// NetInputTokens returns input tokens with the cached prompt subset removed, per +// the disjoint-bucket contract (InputTokens must exclude cache reads). If the +// cached count is absent or inconsistent (exceeds input), the input is returned +// unchanged rather than clamped to zero, so malformed provider data cannot erase +// real input tokens. +func NetInputTokens(input, cached int) int { + if cached <= 0 || input-cached < 0 { + return input + } + return input - cached +} + +// NetOutputTokens returns output tokens with the reasoning subset removed, per +// the disjoint-bucket contract (OutputTokens must exclude reasoning). Uses the +// same inconsistency guard as NetInputTokens. +func NetOutputTokens(output, reasoning int) int { + if reasoning <= 0 || output-reasoning < 0 { + return output + } + return output - reasoning +} + // Cost is the token + money accounting for a model call. Canonical home; pkg/ai // re-exports it via `type Cost = api.Cost`. type Cost struct { @@ -30,14 +61,27 @@ type Cost struct { ReasoningCost float64 `json:"reasoningCost,omitempty" yaml:"reasoningCost,omitempty" pretty:"label=Reasoning $,table"` CacheReadCost float64 `json:"cacheReadCost,omitempty" yaml:"cacheReadCost,omitempty" pretty:"label=Cache Read $,table"` CacheWriteCost float64 `json:"cacheWriteCost,omitempty" yaml:"cacheWriteCost,omitempty" pretty:"label=Cache Write $,table"` + + // ProviderCostUSD, when > 0, is the authoritative total the model provider + // reported for this call (e.g. claude-cli total_cost_usd). Total() returns it + // in preference to the list-price bucket sum, so provider billing wins over a + // recomputed estimate. The per-bucket costs are retained for display; this + // value is not rendered as its own column since Total already reflects it. + ProviderCostUSD float64 `json:"providerCostUSD,omitempty" yaml:"providerCostUSD,omitempty"` } -// Total is the combined cost in USD across every billable bucket. +// Total is the combined cost in USD. It prefers the provider-reported total when +// present, falling back to the sum of the list-price buckets. func (c Cost) Total() float64 { + if c.ProviderCostUSD > 0 { + return c.ProviderCostUSD + } return c.InputCost + c.OutputCost + c.ReasoningCost + c.CacheReadCost + c.CacheWriteCost } -// Add returns the field-wise sum of two costs, keeping the receiver's Model. +// Add returns the field-wise sum of two costs, keeping the receiver's Model. The +// provider-reported totals sum too, so a rollup's Total() stays authoritative +// across a session (which, in practice, uses a single provider throughout). func (c Cost) Add(other Cost) Cost { return Cost{ Model: c.Model, @@ -52,6 +96,7 @@ func (c Cost) Add(other Cost) Cost { ReasoningCost: c.ReasoningCost + other.ReasoningCost, CacheReadCost: c.CacheReadCost + other.CacheReadCost, CacheWriteCost: c.CacheWriteCost + other.CacheWriteCost, + ProviderCostUSD: c.ProviderCostUSD + other.ProviderCostUSD, } } diff --git a/pkg/api/enums.go b/pkg/api/enums.go index 83089ac..4300e8d 100644 --- a/pkg/api/enums.go +++ b/pkg/api/enums.go @@ -3,186 +3,40 @@ // Setup, Prompt) and the Spec that composes them. It never imports pkg/ai; // pkg/ai re-exports the enum/value types here via aliases, so this package is // the single source of truth for Backend/Effort/Cost. +// +// Model identity itself (Backend, Effort, Model, the compact grammar, the model +// catalog) lives one level down in pkg/api/registry and is re-exported here by +// alias — see aliases.go. Spec decoding parses model strings, so the parser has +// to sit below this package. package api import ( - "errors" "fmt" "strings" ) -// ErrInferBackend marks the "can't infer a backend from this model name" failure -// so callers can enrich it (e.g. with "did you mean" model suggestions). -var ErrInferBackend = errors.New("unable to infer backend from model name") - -// Backend is the provider/runtime that serves a request. This is the canonical -// definition; pkg/ai re-exports it via `type Backend = api.Backend`. -type Backend string - -const ( - BackendAnthropic Backend = "anthropic" - BackendGemini Backend = "gemini" - BackendOpenAI Backend = "openai" - BackendDeepSeek Backend = "deepseek" - BackendClaudeCLI Backend = "claude-cli" - BackendCodexCLI Backend = "codex-cli" - BackendGeminiCLI Backend = "gemini-cli" - BackendClaudeAgent Backend = "claude-agent" - // BackendClaudeCmux / BackendCodexCmux drive an interactive claude/codex TUI - // inside a tmux/cmux surface (the cmux provider), tailing the session JSONL. - // They are selected explicitly, not inferred from a model name. - BackendClaudeCmux Backend = "claude-cmux" - BackendCodexCmux Backend = "codex-cmux" -) - -// AllBackends lists every supported backend in canonical order — the single -// source of truth behind Valid, BackendList, and the help/error strings. -func AllBackends() []Backend { - return []Backend{ - BackendAnthropic, BackendGemini, BackendOpenAI, BackendDeepSeek, - BackendClaudeCLI, BackendClaudeAgent, BackendClaudeCmux, - BackendCodexCLI, BackendCodexCmux, BackendGeminiCLI, - } -} - -// Valid reports whether b is one of the supported backends. -func (b Backend) Valid() bool { - for _, x := range AllBackends() { - if b == x { - return true - } - } - return false -} - -// Kind classifies a backend as "api" (called directly over HTTP with an API key) -// or "cli" (delegated to an installed coding-agent binary with its own auth). -func (b Backend) Kind() string { - switch b { - case BackendAnthropic, BackendGemini, BackendOpenAI, BackendDeepSeek: - return "api" - default: - return "cli" - } -} - -// AuthEnvVars returns the environment variables consulted for a backend's API -// key, in priority order. Some CLI backends can use a parent provider key, while -// cmux backends are keyless and rely on the local CLI login. -func AuthEnvVars(b Backend) []string { - switch b { - case BackendAnthropic, BackendClaudeCLI, BackendClaudeAgent: - return []string{"ANTHROPIC_API_KEY"} - case BackendOpenAI, BackendCodexCLI: - return []string{"OPENAI_API_KEY"} - case BackendDeepSeek: - return []string{"DEEPSEEK_API_KEY"} - case BackendGemini, BackendGeminiCLI: - return []string{"GEMINI_API_KEY", "GOOGLE_API_KEY"} - default: - return nil - } -} - -// BackendList renders AllBackends as a comma-separated string for help/error text. -func BackendList() string { - parts := make([]string, len(AllBackends())) - for i, b := range AllBackends() { - parts[i] = string(b) - } - return strings.Join(parts, ", ") -} - -// InferBackend resolves the backend from a model name prefix, failing loud when -// the name matches nothing (the caller must then pass an explicit backend). -func InferBackend(model string) (Backend, error) { - m := strings.ToLower(model) - - // CLI backends (check before API backends to avoid prefix conflicts). - switch { - case strings.HasPrefix(m, "claude-agent-"): - return BackendClaudeAgent, nil - case strings.HasPrefix(m, "claude-code-"): - return BackendClaudeCLI, nil - case strings.HasPrefix(m, "codex"): - return BackendCodexCLI, nil - case strings.HasPrefix(m, "gemini-cli-"): - return BackendGeminiCLI, nil - } - - switch { - case strings.HasPrefix(m, "claude-"): - return BackendAnthropic, nil - case strings.HasPrefix(m, "gemini-"), strings.HasPrefix(m, "models/gemini-"): - return BackendGemini, nil - case strings.HasPrefix(m, "grok-"): - return BackendCodexCLI, nil - case strings.HasPrefix(m, "gpt-"), strings.HasPrefix(m, "o1"), strings.HasPrefix(m, "o3"), strings.HasPrefix(m, "o4"): - return BackendOpenAI, nil - case strings.HasPrefix(m, "deepseek"): - return BackendDeepSeek, nil - } - - return "", fmt.Errorf("%w: %s (pass an explicit backend: %s)", ErrInferBackend, model, BackendList()) -} - -// Effort is the per-request reasoning effort. captain owns this enum (it adds -// the "xhigh" tier that clicky's aichat.Effort lacks); "" means backend default. -type Effort string - -const ( - EffortNone Effort = "" - EffortLow Effort = "low" - EffortMedium Effort = "medium" - EffortHigh Effort = "high" - EffortXHigh Effort = "xhigh" -) - -// AllEfforts lists the non-empty effort tiers in ascending order. -func AllEfforts() []Effort { - return []Effort{EffortLow, EffortMedium, EffortHigh, EffortXHigh} -} - -// Valid reports whether e is a recognised effort tier (including none/""). -func (e Effort) Valid() bool { - switch e { - case EffortNone, EffortLow, EffortMedium, EffortHigh, EffortXHigh: - return true - default: - return false - } -} - -// Validate fails loud on an unknown effort tier, naming the valid set. -func (e Effort) Validate() error { - if e.Valid() { - return nil - } - return fmt.Errorf("invalid reasoning effort %q; want one of: low, medium, high, xhigh", e) -} - // SchemaStrictness governs what captain does when a structured-output response -// fails validation against the request's JSON schema. "" disables post-response -// validation (the default — the schema is still sent to the model, just never -// checked on the way back). +// fails validation against the request's JSON schema. "" uses the backend +// default, while "none" explicitly disables post-response validation. type SchemaStrictness string const ( - SchemaStrictnessNone SchemaStrictness = "" - SchemaStrictnessWarning SchemaStrictness = "warning" - SchemaStrictnessError SchemaStrictness = "error" - SchemaStrictnessRetry SchemaStrictness = "retry" + SchemaStrictnessNone SchemaStrictness = "" + SchemaStrictnessDisabled SchemaStrictness = "none" + SchemaStrictnessWarning SchemaStrictness = "warning" + SchemaStrictnessError SchemaStrictness = "error" + SchemaStrictnessRetry SchemaStrictness = "retry" ) // AllSchemaStrictness lists the non-empty strictness modes. func AllSchemaStrictness() []SchemaStrictness { - return []SchemaStrictness{SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry} + return []SchemaStrictness{SchemaStrictnessDisabled, SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry} } // Valid reports whether s is a recognised strictness mode (including none/""). func (s SchemaStrictness) Valid() bool { switch s { - case SchemaStrictnessNone, SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry: + case SchemaStrictnessNone, SchemaStrictnessDisabled, SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry: return true default: return false @@ -194,7 +48,7 @@ func (s SchemaStrictness) Validate() error { if s.Valid() { return nil } - return fmt.Errorf("invalid schemaStrictness %q; want one of: warning, error, retry", s) + return fmt.Errorf("invalid schemaStrictness %q; want one of: none, warning, error, retry", s) } // VerifyScope narrows a workflow's verification to the changed files vs the @@ -233,21 +87,34 @@ func (s VerifyScope) Validate() error { type ToolMode string const ( - ToolModeEnabled ToolMode = "enabled" - ToolModeAsk ToolMode = "ask" - ToolModeDisabled ToolMode = "disabled" + ToolModeOn ToolMode = "on" + ToolModeAsk ToolMode = "ask" + ToolModeOff ToolMode = "off" + ToolModeAuto ToolMode = "auto" ) -// Valid reports whether m is a recognised tool mode. -func (m ToolMode) Valid() bool { - switch m { - case ToolModeEnabled, ToolModeAsk, ToolModeDisabled: - return true +// NormalizeToolMode canonicalizes a mode. +func NormalizeToolMode(m ToolMode) (ToolMode, bool) { + switch ToolMode(strings.ToLower(strings.TrimSpace(string(m)))) { + case ToolModeOn: + return ToolModeOn, true + case ToolModeAsk: + return ToolModeAsk, true + case ToolModeOff: + return ToolModeOff, true + case ToolModeAuto: + return ToolModeAuto, true default: - return false + return "", false } } +// Valid reports whether m is a recognised tool mode. +func (m ToolMode) Valid() bool { + _, ok := NormalizeToolMode(m) + return ok +} + // ToolPolicy is the runtime-spec policy map value for one tool. It keeps the // wire shape close to coding-agent UX: auto, ask, allow, deny. type ToolPolicy string diff --git a/pkg/api/enums_test.go b/pkg/api/enums_test.go index 5776623..836eab8 100644 --- a/pkg/api/enums_test.go +++ b/pkg/api/enums_test.go @@ -5,70 +5,20 @@ import ( "testing" ) -func TestInferBackend(t *testing.T) { - cases := map[string]Backend{ - "claude-sonnet-4-6": BackendAnthropic, - "claude-agent-opus": BackendClaudeAgent, - "claude-code-x": BackendClaudeCLI, - "gpt-5.5": BackendOpenAI, - "o3": BackendOpenAI, - "gemini-2.5-pro": BackendGemini, - "gemini-cli-pro": BackendGeminiCLI, - "codex-gpt-5-codex": BackendCodexCLI, - "grok-2": BackendCodexCLI, - "deepseek-chat": BackendDeepSeek, - } - for model, want := range cases { - got, err := InferBackend(model) - if err != nil || got != want { - t.Errorf("InferBackend(%q) = %q, %v; want %q", model, got, err, want) - } - } - if _, err := InferBackend("totally-unknown"); err == nil { - t.Error("InferBackend(unknown) should fail loud") - } -} - -func TestEffortValidateIncludesXHigh(t *testing.T) { - for _, e := range []Effort{EffortNone, EffortLow, EffortMedium, EffortHigh, EffortXHigh} { - if err := e.Validate(); err != nil { - t.Errorf("Effort(%q).Validate() = %v, want nil", e, err) - } - } - if err := Effort("max").Validate(); err == nil { - t.Error("Effort(max).Validate() should fail (only low/medium/high/xhigh)") - } - if want := []Effort{EffortLow, EffortMedium, EffortHigh, EffortXHigh}; !reflect.DeepEqual(AllEfforts(), want) { - t.Errorf("AllEfforts() = %v, want %v", AllEfforts(), want) - } -} - func TestSchemaStrictnessValidate(t *testing.T) { - for _, s := range []SchemaStrictness{SchemaStrictnessNone, SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry} { + for _, s := range []SchemaStrictness{SchemaStrictnessNone, SchemaStrictnessDisabled, SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry} { if err := s.Validate(); err != nil { t.Errorf("SchemaStrictness(%q).Validate() = %v, want nil", s, err) } } if err := SchemaStrictness("strict").Validate(); err == nil { - t.Error("SchemaStrictness(strict).Validate() should fail (only warning/error/retry)") + t.Error("SchemaStrictness(strict).Validate() should fail (only none/warning/error/retry)") } - if want := []SchemaStrictness{SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry}; !reflect.DeepEqual(AllSchemaStrictness(), want) { + if want := []SchemaStrictness{SchemaStrictnessDisabled, SchemaStrictnessWarning, SchemaStrictnessError, SchemaStrictnessRetry}; !reflect.DeepEqual(AllSchemaStrictness(), want) { t.Errorf("AllSchemaStrictness() = %v, want %v", AllSchemaStrictness(), want) } } -func TestBackendKindAndAuth(t *testing.T) { - if BackendAnthropic.Kind() != "api" || BackendClaudeAgent.Kind() != "cli" { - t.Errorf("Kind() classification wrong: api=%q cli=%q", BackendAnthropic.Kind(), BackendClaudeAgent.Kind()) - } - if got := AuthEnvVars(BackendGemini); !reflect.DeepEqual(got, []string{"GEMINI_API_KEY", "GOOGLE_API_KEY"}) { - t.Errorf("AuthEnvVars(gemini) = %v", got) - } - if !BackendOpenAI.Valid() || Backend("nope").Valid() { - t.Error("Backend.Valid() wrong") - } -} - func TestPermissionModeValid(t *testing.T) { if !PermissionMode("").Valid() || !PermissionAcceptEdits.Valid() || PermissionMode("yolo").Valid() { t.Error("PermissionMode.Valid() wrong") diff --git a/pkg/api/message.go b/pkg/api/message.go new file mode 100644 index 0000000..5fb84b5 --- /dev/null +++ b/pkg/api/message.go @@ -0,0 +1,222 @@ +package api + +import ( + "encoding/json" + "fmt" + "strings" +) + +// MessageRole identifies the author of one canonical conversation message. +type MessageRole string + +const ( + RoleSystem MessageRole = "system" + RoleUser MessageRole = "user" + RoleAssistant MessageRole = "assistant" + RoleTool MessageRole = "tool" +) + +func (r MessageRole) Valid() bool { + switch r { + case RoleSystem, RoleUser, RoleAssistant, RoleTool: + return true + default: + return false + } +} + +// PartType identifies the provider-neutral content carried by a message part. +type PartType string + +const ( + PartText PartType = "text" + PartReasoning PartType = "reasoning" + PartAttachment PartType = "attachment" + PartToolRequest PartType = "tool-request" + PartToolResult PartType = "tool-result" +) + +func (t PartType) Valid() bool { + switch t { + case PartText, PartReasoning, PartAttachment, PartToolRequest, PartToolResult: + return true + default: + return false + } +} + +// ToolRequest is one assistant request to execute a named tool. ToolCallID is +// stable across the corresponding ToolResult and provider projection. +type ToolRequest struct { + ToolCallID string `json:"toolCallId" yaml:"toolCallId"` + Name string `json:"name" yaml:"name"` + Input json.RawMessage `json:"input,omitempty" yaml:"input,omitempty"` +} + +// ToolResult is the output of a prior ToolRequest. Error is set instead of +// Output when tool execution failed. +type ToolResult struct { + ToolCallID string `json:"toolCallId" yaml:"toolCallId"` + Output json.RawMessage `json:"output,omitempty" yaml:"output,omitempty"` + Error string `json:"error,omitempty" yaml:"error,omitempty"` +} + +// Part is one typed content item in a canonical Message. Type selects exactly +// one payload: Text, Attachment, ToolRequest, or ToolResult. +type Part struct { + Type PartType `json:"type" yaml:"type"` + Text string `json:"text,omitempty" yaml:"text,omitempty"` + Attachment *AttachmentRef `json:"attachment,omitempty" yaml:"attachment,omitempty"` + ToolRequest *ToolRequest `json:"toolRequest,omitempty" yaml:"toolRequest,omitempty"` + ToolResult *ToolResult `json:"toolResult,omitempty" yaml:"toolResult,omitempty"` +} + +// Message is one provider-neutral conversation entry. +type Message struct { + Role MessageRole `json:"role" yaml:"role"` + Parts []Part `json:"parts" yaml:"parts"` +} + +func (p Part) Validate() error { + if !p.Type.Valid() { + return fmt.Errorf("invalid part type %q", p.Type) + } + if err := p.validatePayloadShape(); err != nil { + return err + } + switch p.Type { + case PartText, PartReasoning: + if strings.TrimSpace(p.Text) == "" { + return fmt.Errorf("%s text is required", p.Type) + } + case PartAttachment: + if err := p.Attachment.Validate(); err != nil { + return fmt.Errorf("attachment: %w", err) + } + case PartToolRequest: + if strings.TrimSpace(p.ToolRequest.ToolCallID) == "" { + return fmt.Errorf("tool request call ID is required") + } + if strings.TrimSpace(p.ToolRequest.Name) == "" { + return fmt.Errorf("tool request name is required") + } + if len(p.ToolRequest.Input) > 0 && !json.Valid(p.ToolRequest.Input) { + return fmt.Errorf("tool request input must be valid JSON") + } + case PartToolResult: + if strings.TrimSpace(p.ToolResult.ToolCallID) == "" { + return fmt.Errorf("tool result call ID is required") + } + if len(p.ToolResult.Output) > 0 && !json.Valid(p.ToolResult.Output) { + return fmt.Errorf("tool result output must be valid JSON") + } + if len(p.ToolResult.Output) > 0 && p.ToolResult.Error != "" { + return fmt.Errorf("tool result output and error are mutually exclusive") + } + } + return nil +} + +func (p Part) validatePayloadShape() error { + textSet := p.Text != "" + attachmentSet := p.Attachment != nil + requestSet := p.ToolRequest != nil + resultSet := p.ToolResult != nil + switch p.Type { + case PartText, PartReasoning: + if attachmentSet || requestSet || resultSet { + return fmt.Errorf("%s part contains an incompatible payload", p.Type) + } + case PartAttachment: + if !attachmentSet || textSet || requestSet || resultSet { + return fmt.Errorf("attachment part must contain only an attachment") + } + case PartToolRequest: + if !requestSet || textSet || attachmentSet || resultSet { + return fmt.Errorf("tool-request part must contain only a tool request") + } + case PartToolResult: + if !resultSet || textSet || attachmentSet || requestSet { + return fmt.Errorf("tool-result part must contain only a tool result") + } + } + return nil +} + +// ValidateMessages validates role/part compatibility and tool transcript +// correlation. Tool results must directly follow the assistant message whose +// stable call IDs they resolve. +func ValidateMessages(messages []Message) error { + if len(messages) == 0 { + return fmt.Errorf("at least one message is required") + } + requestIDs := map[string]bool{} + resultIDs := map[string]bool{} + for i, message := range messages { + if !message.Role.Valid() { + return fmt.Errorf("message %d: invalid role %q", i+1, message.Role) + } + if len(message.Parts) == 0 { + return fmt.Errorf("message %d (%s): at least one part is required", i+1, message.Role) + } + for j, part := range message.Parts { + if err := part.Validate(); err != nil { + return fmt.Errorf("message %d part %d: %w", i+1, j+1, err) + } + if !roleAllowsPart(message.Role, part.Type) { + return fmt.Errorf("message %d part %d: %s message cannot contain %s", i+1, j+1, message.Role, part.Type) + } + if part.Type == PartToolRequest { + if requestIDs[part.ToolRequest.ToolCallID] { + return fmt.Errorf("message %d part %d: duplicate tool request call ID %q", i+1, j+1, part.ToolRequest.ToolCallID) + } + requestIDs[part.ToolRequest.ToolCallID] = true + } + if part.Type == PartToolResult { + if resultIDs[part.ToolResult.ToolCallID] { + return fmt.Errorf("message %d part %d: duplicate tool result call ID %q", i+1, j+1, part.ToolResult.ToolCallID) + } + resultIDs[part.ToolResult.ToolCallID] = true + } + } + if message.Role == RoleTool { + if err := validateToolMessage(messages, i); err != nil { + return fmt.Errorf("message %d: %w", i+1, err) + } + } + } + return nil +} + +func roleAllowsPart(role MessageRole, part PartType) bool { + switch role { + case RoleSystem: + return part == PartText + case RoleUser: + return part == PartText || part == PartAttachment + case RoleAssistant: + return part == PartText || part == PartReasoning || part == PartToolRequest + case RoleTool: + return part == PartToolResult + default: + return false + } +} + +func validateToolMessage(messages []Message, index int) error { + if index == 0 || messages[index-1].Role != RoleAssistant { + return fmt.Errorf("tool message must immediately follow an assistant message") + } + preceding := map[string]bool{} + for _, part := range messages[index-1].Parts { + if part.Type == PartToolRequest && part.ToolRequest != nil { + preceding[part.ToolRequest.ToolCallID] = true + } + } + for _, part := range messages[index].Parts { + if part.Type == PartToolResult && part.ToolResult != nil && !preceding[part.ToolResult.ToolCallID] { + return fmt.Errorf("tool call %q was not requested by the preceding assistant message", part.ToolResult.ToolCallID) + } + } + return nil +} diff --git a/pkg/api/message_ginkgo_test.go b/pkg/api/message_ginkgo_test.go new file mode 100644 index 0000000..976a5a3 --- /dev/null +++ b/pkg/api/message_ginkgo_test.go @@ -0,0 +1,77 @@ +package api_test + +import ( + "encoding/json" + "strings" + + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Canonical messages", func() { + validMessages := func() []api.Message { + return []api.Message{ + {Role: api.RoleSystem, Parts: []api.Part{{Type: api.PartText, Text: "Be precise."}}}, + {Role: api.RoleUser, Parts: []api.Part{ + {Type: api.PartText, Text: "Inspect this image."}, + {Type: api.PartAttachment, Attachment: &api.AttachmentRef{ID: api.AttachmentIDPrefix + strings.Repeat("a", 64), MediaType: "image/png"}}, + }}, + {Role: api.RoleAssistant, Parts: []api.Part{ + {Type: api.PartReasoning, Text: "I should inspect its dimensions."}, + {Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ToolCallID: "call-1", Name: "image_dimensions", Input: json.RawMessage(`{"attachment":"a"}`)}}, + }}, + {Role: api.RoleTool, Parts: []api.Part{ + {Type: api.PartToolResult, ToolResult: &api.ToolResult{ToolCallID: "call-1", Output: json.RawMessage(`{"width":640,"height":480}`)}}, + }}, + {Role: api.RoleAssistant, Parts: []api.Part{{Type: api.PartText, Text: "It is 640 by 480."}}}, + } + } + + It("validates a multimodal tool transcript", func() { + Expect(api.ValidateMessages(validMessages())).To(Succeed()) + }) + + DescribeTable("rejects parts on incompatible roles", + func(role api.MessageRole, part api.Part, want string) { + err := api.ValidateMessages([]api.Message{{Role: role, Parts: []api.Part{part}}}) + Expect(err).To(MatchError(ContainSubstring(want))) + }, + Entry("system reasoning", api.RoleSystem, api.Part{Type: api.PartReasoning, Text: "thought"}, "system message cannot contain reasoning"), + Entry("user reasoning", api.RoleUser, api.Part{Type: api.PartReasoning, Text: "thought"}, "user message cannot contain reasoning"), + Entry("assistant attachment", api.RoleAssistant, api.Part{Type: api.PartAttachment, Attachment: &api.AttachmentRef{Path: "image.png"}}, "assistant message cannot contain attachment"), + Entry("tool text", api.RoleTool, api.Part{Type: api.PartText, Text: "result"}, "tool message cannot contain text"), + ) + + It("requires a tool result to immediately follow and match an assistant request", func() { + messages := validMessages() + messages[3].Parts[0].ToolResult.ToolCallID = "unknown" + Expect(api.ValidateMessages(messages)).To(MatchError(ContainSubstring(`tool call "unknown" was not requested by the preceding assistant message`))) + + messages = validMessages() + messages = append(messages[:3], api.Message{Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "continue"}}}, messages[3]) + Expect(api.ValidateMessages(messages)).To(MatchError(ContainSubstring("tool message must immediately follow an assistant message"))) + }) + + It("requires stable unique tool call IDs", func() { + messages := validMessages() + messages[2].Parts[1].ToolRequest.ToolCallID = "" + Expect(api.ValidateMessages(messages)).To(MatchError(ContainSubstring("tool request call ID is required"))) + + messages = validMessages() + messages[2].Parts = append(messages[2].Parts, api.Part{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: "call-1", Name: "other", Input: json.RawMessage(`{}`), + }}) + Expect(api.ValidateMessages(messages)).To(MatchError(ContainSubstring(`duplicate tool request call ID "call-1"`))) + }) + + It("keeps conversation and single-turn prompt as mutually exclusive request modes", func() { + spec := api.Spec{Model: api.Model{Name: "gpt-5", Backend: api.BackendOpenAI}, Messages: validMessages()} + Expect(spec.Validate()).To(Succeed()) + Expect(spec.IsVerifyOnly()).To(BeFalse()) + + spec.Prompt.User = "silent fallback" + Expect(spec.Validate()).To(MatchError(ContainSubstring("prompt body and messages are mutually exclusive"))) + }) +}) diff --git a/pkg/api/model.go b/pkg/api/model.go deleted file mode 100644 index c22ac42..0000000 --- a/pkg/api/model.go +++ /dev/null @@ -1,154 +0,0 @@ -package api - -import ( - "fmt" - "strings" - - "github.com/flanksource/captain/pkg/collections" -) - -// Model identifies which LLM serves a request plus the per-request inference -// knobs. Maps onto the legacy ai.Config.Model + ai.Request.{Temperature, -// ReasoningEffort}. -type Model struct { - // Name is the catalog model slug, e.g. "claude-sonnet-4-6"; it drives backend - // inference and pricing lookup. - Name string `json:"model" yaml:"model" jsonschema:"required" pretty:"label=Model"` - - // ID is the fully-qualified provider id when it differs from Name, e.g. the - // genkit "anthropic/claude-sonnet-4-6" or a codex slug. Empty means use Name. - ID string `json:"id,omitempty" yaml:"id,omitempty" pretty:"label=ID"` - - // Backend overrides backend inference from Name; empty means InferBackend(Name). - Backend Backend `json:"backend,omitempty" yaml:"backend,omitempty" pretty:"label=Backend"` - - // Temperature is the sampling temperature in [0,2]. A pointer so an explicit - // 0.0 is distinguishable from "unset" (fail loud, not a silent default). - Temperature *float64 `json:"temperature,omitempty" yaml:"temperature,omitempty" pretty:"label=Temp"` - - // Effort is the reasoning effort for thinking-capable models. - Effort Effort `json:"effort,omitempty" yaml:"effort,omitempty" jsonschema:"enum=,enum=low,enum=medium,enum=high,enum=xhigh" pretty:"label=Effort"` - - // NoCache disables model response caching for this run. - NoCache bool `json:"noCache,omitempty" yaml:"noCache,omitempty" pretty:"label=No Cache"` - - // Fallbacks are alternative models tried in order when the primary fails with a - // transient/unavailable error or its provider cannot be constructed. Each is a - // full Model (own backend/effort/temperature); a fallback's own nested Fallbacks - // are ignored. Populate it directly, via a comma-separated Name (see ExpandCSV), - // or via the --fallback flag / fallbacks: frontmatter. - Fallbacks []Model `json:"fallbacks,omitempty" yaml:"fallbacks,omitempty" pretty:"label=Fallbacks"` -} - -// ResolveBackend returns Backend when set, otherwise infers it from Name. -func (m Model) ResolveBackend() (Backend, error) { - if m.Backend != "" { - if !m.Backend.Valid() { - return "", fmt.Errorf("invalid backend %q (valid: %s)", m.Backend, BackendList()) - } - return m.Backend, nil - } - return InferBackend(m.Name) -} - -// Temp returns the temperature and whether it was explicitly set (non-nil), so -// providers can distinguish an intentional 0.0 from "unset, use the default". -func (m Model) Temp() (float64, bool) { - if m.Temperature == nil { - return 0, false - } - return *m.Temperature, true -} - -// Validate checks the model name is present, the knobs are in range, and every -// fallback is itself a well-formed model. -func (m Model) Validate() error { - if m.Name == "" { - return fmt.Errorf("model name is required") - } - if err := m.validateKnobs(); err != nil { - return err - } - for i, fb := range m.Fallbacks { - if fb.Name == "" { - return fmt.Errorf("fallback[%d]: model name is required", i) - } - if err := fb.validateKnobs(); err != nil { - return fmt.Errorf("fallback[%d] %q: %w", i, fb.Name, err) - } - } - return nil -} - -// validateKnobs range-checks the per-request inference knobs shared by a primary -// model and each fallback (backend/temperature/effort), independent of Name. -func (m Model) validateKnobs() error { - if m.Backend != "" && !m.Backend.Valid() { - return fmt.Errorf("invalid backend %q (valid: %s)", m.Backend, BackendList()) - } - if m.Temperature != nil && (*m.Temperature < 0 || *m.Temperature > 2) { - return fmt.Errorf("invalid temperature %v (valid: 0.0-2.0)", *m.Temperature) - } - return m.Effort.Validate() -} - -// ExpandCSV moves any comma-separated tail of Name into name-only Fallbacks -// (prepended, so CSV order is preserved ahead of an explicit Fallbacks list), -// leaving Name as the single primary. It is idempotent: a single-model Name is -// returned with only whitespace trimmed, and a second call is a no-op. -func (m Model) ExpandCSV() Model { - names := splitCSV(m.Name) - if len(names) == 0 { - return m - } - m.Name = names[0] - if len(names) == 1 { - return m - } - tail := make([]Model, 0, len(names)-1) - for _, n := range names[1:] { - tail = append(tail, Model{Name: n}) - } - m.Fallbacks = append(tail, m.Fallbacks...) - return m -} - -// Candidates returns the ordered models to try: the ExpandCSV primary first, then -// each fallback. Fallbacks inherit the primary's Temperature/Effort/NoCache when -// unset, keep their own Name/Backend (an empty Backend is inferred at construction -// from the fallback's own Name), clear ID, and have nested Fallbacks dropped. A -// length of 1 means "no fallback". -func (m Model) Candidates() []Model { - m = m.ExpandCSV() - primary := m - primary.Fallbacks = nil - out := make([]Model, 0, collections.SafeAdd(1, len(m.Fallbacks))) - out = append(out, primary) - for _, fb := range m.Fallbacks { - fb.Fallbacks = nil - fb.ID = "" - if fb.Temperature == nil { - fb.Temperature = m.Temperature - } - if fb.Effort == "" { - fb.Effort = m.Effort - } - if !fb.NoCache { - fb.NoCache = m.NoCache - } - out = append(out, fb) - } - return out -} - -// splitCSV splits a comma-separated string into trimmed, non-empty parts. -func splitCSV(s string) []string { - parts := strings.Split(s, ",") - out := make([]string, 0, len(parts)) - for _, p := range parts { - if t := strings.TrimSpace(p); t != "" { - out = append(out, t) - } - } - return out -} diff --git a/pkg/api/permissions.go b/pkg/api/permissions.go index 0a43beb..99198c4 100644 --- a/pkg/api/permissions.go +++ b/pkg/api/permissions.go @@ -64,6 +64,11 @@ func (p Permissions) Validate() error { return fmt.Errorf("invalid preset %q (valid: edit, bare)", preset) } } + for tool, mode := range p.Tools.Modes { + if !mode.Valid() { + return fmt.Errorf("invalid tool mode %q for tool %q (valid: on, ask, off, auto)", mode, tool) + } + } for tool, policy := range p.Tools.Policies() { if !policy.Valid() { return fmt.Errorf("invalid tool policy %q for tool %q (valid: auto, ask, allow, deny)", policy, tool) @@ -105,11 +110,11 @@ func (t Tools) Policies() map[string]ToolPolicy { continue } switch mode { - case ToolModeEnabled: + case ToolModeOn: out[tool] = ToolPolicyAuto case ToolModeAsk: out[tool] = ToolPolicyAsk - case ToolModeDisabled: + case ToolModeOff: out[tool] = ToolPolicyDeny } } @@ -215,7 +220,7 @@ func (t *Tools) applyPolicy(tool string, policy ToolPolicy) { if t.Modes == nil { t.Modes = map[string]ToolMode{} } - t.Modes[tool] = ToolModeEnabled + t.Modes[tool] = ToolModeOn } } diff --git a/pkg/api/permissions_test.go b/pkg/api/permissions_test.go index 394a0aa..7b07e8e 100644 --- a/pkg/api/permissions_test.go +++ b/pkg/api/permissions_test.go @@ -16,7 +16,7 @@ func TestPermissions_JSONPolicyShape(t *testing.T) { Deny: []string{"Bash"}, Modes: map[string]ToolMode{ "WebSearch": ToolModeAsk, - "Write": ToolModeEnabled, + "Write": ToolModeOn, }, }, MCP: MCP{ @@ -56,7 +56,7 @@ func TestPermissions_JSONLegacyInput(t *testing.T) { "tools": { "allow": ["Read"], "deny": ["Bash"], - "modes": {"Edit": "ask", "Write": "enabled"} + "modes": {"Edit": "ask", "Write": "on"} }, "mcp": {"servers": ["filesystem", "gavel"], "gavel": "disabled"}, "plugins": ["/plugins"], diff --git a/pkg/api/prompt.go b/pkg/api/prompt.go index bdff782..3172c9d 100644 --- a/pkg/api/prompt.go +++ b/pkg/api/prompt.go @@ -3,6 +3,7 @@ package api import ( "encoding/json" "fmt" + "strings" ) // Prompt is the instruction payload: the user prompt plus optional system @@ -11,7 +12,7 @@ import ( // Source,StructuredOutput,Metadata}. type Prompt struct { // User is the user prompt. (ai.Request.Prompt) - User string `json:"user" yaml:"user" jsonschema:"required" pretty:"label=Prompt"` + User string `json:"user,omitempty" yaml:"user,omitempty" jsonschema:"required" pretty:"label=Prompt"` // System is the system prompt. (ai.Request.SystemPrompt) System string `json:"system,omitempty" yaml:"system,omitempty" pretty:"label=System"` // AppendSystem is appended to the default system prompt. (ai.Request.AppendSystemPrompt) @@ -30,22 +31,44 @@ type Prompt struct { // SchemaJSON are mutually exclusive. SchemaJSON json.RawMessage `json:"schemaJSON,omitempty" yaml:"schemaJSON,omitempty" pretty:"-"` // SchemaStrictness governs how a response that fails JSON-schema validation is - // handled: "" (default) skips post-response validation; "warning" logs and - // continues; "error" fails; "retry" re-asks the model once with the validation + // handled: "" uses the backend default (Anthropic structured output retries, + // others skip post-validation); "none" always skips; "warning" logs and + // continues; "error" fails; "retry" re-asks the model with the validation // error, then fails. Only meaningful alongside a schema (Schema or SchemaJSON). SchemaStrictness SchemaStrictness `json:"schemaStrictness,omitempty" yaml:"schemaStrictness,omitempty" pretty:"-"` // Metadata is arbitrary caller metadata. (ai.Request.Metadata) Metadata map[string]string `json:"metadata,omitempty" yaml:"metadata,omitempty" pretty:"label=Metadata"` + // Attachments are ordered multimodal inputs sent with the user prompt. + Attachments []AttachmentRef `json:"attachments,omitempty" yaml:"attachments,omitempty" pretty:"label=Attachments"` +} + +func (p Prompt) CacheIdentity() string { + var identity strings.Builder + identity.WriteString(p.User) + for _, attachment := range p.Attachments { + identity.WriteByte('\n') + identity.WriteString(attachment.ID) + identity.WriteByte('|') + identity.WriteString(attachment.MediaType) + identity.WriteByte('|') + identity.WriteString(attachment.SHA256) + } + return identity.String() } // HasSchema reports whether the prompt requests structured output by either // mechanism (a reflected Go struct or a pre-built JSON schema). func (p Prompt) HasSchema() bool { return p.Schema != nil || len(p.SchemaJSON) > 0 } -// Validate requires a non-empty user prompt. +// Validate requires prompt text or at least one attachment. func (p Prompt) Validate() error { - if p.User == "" { - return fmt.Errorf("prompt text is required") + if p.User == "" && len(p.Attachments) == 0 { + return fmt.Errorf("prompt text is required when no attachment is supplied") + } + for i, attachment := range p.Attachments { + if err := attachment.Validate(); err != nil { + return fmt.Errorf("attachment %d: %w", i+1, err) + } } if err := p.SchemaStrictness.Validate(); err != nil { return err diff --git a/pkg/api/prompt_unmarshal.go b/pkg/api/prompt_unmarshal.go new file mode 100644 index 0000000..d9cb05a --- /dev/null +++ b/pkg/api/prompt_unmarshal.go @@ -0,0 +1,50 @@ +package api + +import ( + "bytes" + "encoding/json" + + "gopkg.in/yaml.v3" +) + +// UnmarshalJSON accepts either a bare string — shorthand for {user: } — +// or the full object form. This lets a config scalar carry just the user prompt +// (`prompt: "review strictly"`) while still supporting the structured form +// (`prompt: {user: ..., system: ...}`). +func (p *Prompt) UnmarshalJSON(data []byte) error { + if trimmed := bytes.TrimSpace(data); len(trimmed) > 0 && trimmed[0] == '"' { + var user string + if err := json.Unmarshal(trimmed, &user); err != nil { + return err + } + *p = Prompt{User: user} + return nil + } + type promptAlias Prompt + var a promptAlias + if err := json.Unmarshal(data, &a); err != nil { + return err + } + *p = Prompt(a) + return nil +} + +// UnmarshalYAML mirrors UnmarshalJSON: a scalar node is the user prompt, a +// mapping node is the full object. +func (p *Prompt) UnmarshalYAML(value *yaml.Node) error { + if value.Kind == yaml.ScalarNode && value.Tag != "!!null" { + var user string + if err := value.Decode(&user); err != nil { + return err + } + *p = Prompt{User: user} + return nil + } + type promptAlias Prompt + var a promptAlias + if err := value.Decode(&a); err != nil { + return err + } + *p = Prompt(a) + return nil +} diff --git a/pkg/api/prompt_unmarshal_test.go b/pkg/api/prompt_unmarshal_test.go new file mode 100644 index 0000000..0932934 --- /dev/null +++ b/pkg/api/prompt_unmarshal_test.go @@ -0,0 +1,76 @@ +package api + +import ( + "encoding/json" + "testing" + + "gopkg.in/yaml.v3" +) + +func TestPrompt_UnmarshalJSON_Shorthand(t *testing.T) { + cases := []struct { + name string + in string + want Prompt + }{ + {"bare string", `"review strictly"`, Prompt{User: "review strictly"}}, + {"object form", `{"user":"u","system":"s"}`, Prompt{User: "u", System: "s"}}, + {"empty string", `""`, Prompt{User: ""}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var p Prompt + if err := json.Unmarshal([]byte(tc.in), &p); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if p.User != tc.want.User || p.System != tc.want.System { + t.Errorf("got %+v, want %+v", p, tc.want) + } + }) + } +} + +func TestPrompt_UnmarshalYAML_Shorthand(t *testing.T) { + cases := []struct { + name string + in string + want Prompt + }{ + {"scalar", `review strictly`, Prompt{User: "review strictly"}}, + {"object", "user: u\nsystem: s", Prompt{User: "u", System: "s"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var p Prompt + if err := yaml.Unmarshal([]byte(tc.in), &p); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if p.User != tc.want.User || p.System != tc.want.System { + t.Errorf("got %+v, want %+v", p, tc.want) + } + }) + } +} + +// TestSpec_PromptShorthand pins that the shorthand works through the enclosing +// Spec — `prompt: "text"` alongside an inlined model — for both encoders. +func TestSpec_PromptShorthand(t *testing.T) { + t.Run("json", func(t *testing.T) { + var s Spec + if err := json.Unmarshal([]byte(`{"model":"m","prompt":"hi there"}`), &s); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if s.Model.Name != "m" || s.Prompt.User != "hi there" { + t.Errorf("got model=%q prompt.user=%q", s.Model.Name, s.Prompt.User) + } + }) + t.Run("yaml", func(t *testing.T) { + var s Spec + if err := yaml.Unmarshal([]byte("model: m\nprompt: hi there\n"), &s); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if s.Model.Name != "m" || s.Prompt.User != "hi there" { + t.Errorf("got model=%q prompt.user=%q", s.Model.Name, s.Prompt.User) + } + }) +} diff --git a/pkg/api/registry/backend.go b/pkg/api/registry/backend.go new file mode 100644 index 0000000..675594d --- /dev/null +++ b/pkg/api/registry/backend.go @@ -0,0 +1,153 @@ +// Package registry owns captain's model identity: the provider descriptors, the +// Backend/Effort/RuntimeMode enums, the Model spec type, the compact model +// grammar, and the embedded model catalog. +// +// It is a leaf — it imports nothing else from captain. That is deliberate: +// pkg/api decodes specs (and therefore parses model strings) during +// unmarshalling, so the parser cannot live above pkg/api without a cycle, and +// splitting the knowledge across both packages is exactly what let the compact +// grammar and the selector grammar drift apart. pkg/api re-exports everything +// here via aliases, so api.Model and api.Backend remain the names callers use. +package registry + +import ( + "errors" + "fmt" + "strings" +) + +// ErrInferBackend marks the "can't infer a backend from this model name" failure +// so callers can enrich it (e.g. with "did you mean" model suggestions). +var ErrInferBackend = errors.New("unable to infer backend from model name") + +// Backend is the provider/runtime that serves a request: exactly one +// (Provider, RuntimeMode) pair. This is the canonical definition; pkg/api +// re-exports it via `type Backend = registry.Backend`. +// +// Its string values are frozen — they are persisted in specs, session rows, and +// the webapp wire format. +type Backend string + +const ( + BackendAnthropic Backend = "anthropic" + BackendGemini Backend = "gemini" + BackendOpenAI Backend = "openai" + BackendDeepSeek Backend = "deepseek" + BackendClaudeCLI Backend = "claude-cli" + BackendCodexCLI Backend = "codex-cli" + BackendGeminiCLI Backend = "gemini-cli" + BackendClaudeAgent Backend = "claude-agent" + BackendCodexAgent Backend = "codex-agent" + // BackendClaudeCmux / BackendCodexCmux drive an interactive claude/codex TUI + // inside a tmux/cmux surface (the cmux provider), tailing the session JSONL. + // They are selected explicitly, not inferred from a model name. + BackendClaudeCmux Backend = "claude-cmux" + BackendCodexCmux Backend = "codex-cmux" +) + +const ( + AnthropicProvider = BackendAnthropic + OpenAIProvider = BackendOpenAI + GeminiProvider = BackendGemini + DeepSeekProvider = BackendDeepSeek +) + +// AllBackends lists every supported backend in canonical order — the single +// source of truth behind Valid, BackendList, and the help/error strings. +func AllBackends() []Backend { + return []Backend{ + BackendAnthropic, BackendGemini, BackendOpenAI, BackendDeepSeek, + BackendClaudeCLI, BackendClaudeAgent, BackendClaudeCmux, + BackendCodexCLI, BackendCodexAgent, BackendCodexCmux, BackendGeminiCLI, + } +} + +// Valid reports whether b is one of the supported backends. +func (b Backend) Valid() bool { + _, _, ok := ProviderFor(b) + return ok +} + +// Mode returns the runtime mechanism this backend represents, or "" if invalid. +func (b Backend) Mode() RuntimeMode { + _, mode, _ := ProviderFor(b) + return mode +} + +// ModelProvider returns the descriptor for the family that owns this backend, +// or nil when the backend is invalid. +func (b Backend) ModelProvider() *Provider { + p, _, _ := ProviderFor(b) + return p +} + +// Kind classifies a backend as "api" (called directly over HTTP with an API key) +// or "cli" (delegated to an installed coding-agent binary with its own auth). +func (b Backend) Kind() string { + if b.Mode() == ModeAPI { + return "api" + } + return "cli" +} + +// Provider returns the direct API provider family that owns a runtime adapter. +// An invalid backend returns the empty value. +func (b Backend) Provider() Backend { + p, _, ok := ProviderFor(b) + if !ok { + return "" + } + backend, err := p.BackendFor(ModeAPI) + if err != nil { + return "" + } + return backend +} + +// Family returns the model family a backend serves (claude | codex | gemini | +// deepseek), or "" for an unrecognised backend. +func (b Backend) Family() string { + p, _, ok := ProviderFor(b) + if !ok { + return "" + } + return p.AgentName +} + +// AuthEnvVars returns the environment variables consulted for a backend's API +// key, in priority order. Some CLI backends can use a parent provider key, while +// cmux backends are keyless and rely on the local CLI login. +func AuthEnvVars(b Backend) []string { + p, mode, ok := ProviderFor(b) + if !ok { + return nil + } + if caps, ok := p.Caps(mode); ok && caps.Keyless { + return nil + } + return p.SupportedEnvVars() +} + +// BackendList renders AllBackends as a comma-separated string for help/error text. +func BackendList() string { + parts := make([]string, len(AllBackends())) + for i, b := range AllBackends() { + parts[i] = string(b) + } + return strings.Join(parts, ", ") +} + +// InferBackend resolves the backend a model name implies, failing loud when no +// provider claims the name (the caller must then pass an explicit backend). +// +// It reads the same claim table as the rest of the parser. It used to carry its +// own prefix switch, which knew about grok but not sora or the codex codenames, +// while pkg/ai's two claim tables each knew a different subset — so the answer +// depended on which entry point you came through. +func InferBackend(model string) (Backend, error) { + p, _, mode, ok := ProviderForToken(model) + if !ok { + return "", fmt.Errorf("%w: %s (pass an explicit backend: %s)", ErrInferBackend, model, BackendList()) + } + return p.BackendFor(mode) +} diff --git a/pkg/api/registry/backend_test.go b/pkg/api/registry/backend_test.go new file mode 100644 index 0000000..36d47d5 --- /dev/null +++ b/pkg/api/registry/backend_test.go @@ -0,0 +1,92 @@ +package registry + +import ( + "reflect" + "testing" +) + +func TestInferBackend(t *testing.T) { + cases := map[string]Backend{ + "claude-sonnet-4-6": BackendAnthropic, + "claude-agent-opus": BackendClaudeAgent, + "claude-code-x": BackendClaudeCLI, + "opus-4-8": BackendAnthropic, + "sonnet": BackendAnthropic, + "gpt-5.5": BackendOpenAI, + "o3": BackendOpenAI, + "gemini-2.5-pro": BackendGemini, + "gemini-cli-pro": BackendGeminiCLI, + "codex-gpt-5-codex": BackendCodexCLI, + "deepseek-chat": BackendDeepSeek, + } + for model, want := range cases { + got, err := InferBackend(model) + if err != nil || got != want { + t.Errorf("InferBackend(%q) = %q, %v; want %q", model, got, err, want) + } + } + if _, err := InferBackend("totally-unknown"); err == nil { + t.Error("InferBackend(unknown) should fail loud") + } + // grok is no longer claimed by any provider (the codex CLI's grok mode was + // removed). Pin the removal: a silent return of grok claiming would otherwise + // resurrect a backend captain no longer routes. + if _, err := InferBackend("grok-2"); err == nil { + t.Error("InferBackend(grok-2) should fail loud: grok mode was removed") + } +} + +func TestEffortValidateIncludesCodexTiers(t *testing.T) { + for _, e := range []Effort{EffortNone, EffortLow, EffortMedium, EffortHigh, EffortXHigh, EffortMax, EffortUltra} { + if err := e.Validate(); err != nil { + t.Errorf("Effort(%q).Validate() = %v, want nil", e, err) + } + } + if err := Effort("extreme").Validate(); err == nil { + t.Error("Effort(extreme).Validate() should fail") + } + if want := []Effort{EffortLow, EffortMedium, EffortHigh, EffortXHigh, EffortMax, EffortUltra}; !reflect.DeepEqual(AllEfforts(), want) { + t.Errorf("AllEfforts() = %v, want %v", AllEfforts(), want) + } +} + +func TestBackendKindAndAuth(t *testing.T) { + if BackendAnthropic.Kind() != "api" || BackendClaudeAgent.Kind() != "cli" { + t.Errorf("Kind() classification wrong: api=%q cli=%q", BackendAnthropic.Kind(), BackendClaudeAgent.Kind()) + } + if got := AuthEnvVars(BackendGemini); !reflect.DeepEqual(got, []string{"GEMINI_API_KEY", "GOOGLE_API_KEY"}) { + t.Errorf("AuthEnvVars(gemini) = %v", got) + } + if !BackendOpenAI.Valid() || Backend("nope").Valid() { + t.Error("Backend.Valid() wrong") + } +} + +func TestBackendProviderAndAgents(t *testing.T) { + tests := map[Backend]Backend{ + BackendAnthropic: AnthropicProvider, + BackendClaudeCLI: AnthropicProvider, + BackendClaudeAgent: AnthropicProvider, + BackendClaudeCmux: AnthropicProvider, + BackendOpenAI: OpenAIProvider, + BackendCodexCLI: OpenAIProvider, + BackendCodexAgent: OpenAIProvider, + BackendCodexCmux: OpenAIProvider, + BackendGemini: GeminiProvider, + BackendGeminiCLI: GeminiProvider, + BackendDeepSeek: DeepSeekProvider, + } + for backend, want := range tests { + if got := backend.Provider(); got != want { + t.Errorf("%s.Provider() = %q, want %q", backend, got, want) + } + } + // Provider.Backends replaces AgentsForProvider. The two lists it unifies + // disagreed on order (AgentsForProvider put the CLI before the agent, the + // wildcard fan-out the reverse); the user-visible wildcard order wins. + if got := OpenAI.Backends(); !reflect.DeepEqual(got, []Backend{ + BackendOpenAI, BackendCodexAgent, BackendCodexCLI, BackendCodexCmux, + }) { + t.Fatalf("OpenAI.Backends() = %v", got) + } +} diff --git a/pkg/api/registry/cost.go b/pkg/api/registry/cost.go new file mode 100644 index 0000000..153cd92 --- /dev/null +++ b/pkg/api/registry/cost.go @@ -0,0 +1,75 @@ +package registry + +import ( + "regexp" + "strings" +) + +// ModelCost is a model's published list price in USD per million tokens, as +// carried in the generated catalog. It is the single source of truth for what a +// model charges. Captain previously kept a hand-written Claude family table +// beside the catalog, which classified on the substring "opus" alone and so +// billed every Opus at the 4.1 rate ($15/$75) long after 4.5 cut it to $5/$25. +type ModelCost struct { + Input float64 `json:"input,omitempty"` + Output float64 `json:"output,omitempty"` + CacheRead float64 `json:"cacheRead,omitempty"` + CacheWrite float64 `json:"cacheWrite,omitempty"` +} + +// snapshotSuffix matches the trailing -YYYYMMDD providers append to a pinned +// model snapshot ("claude-opus-4-5-20251101"). +var snapshotSuffix = regexp.MustCompile(`-\d{8}$`) + +// CostFor resolves a model's list price, accepting aliases, provider +// namespaces, dated snapshot ids, and version lines. ok is false when the +// catalog snapshot carries no price for the model; callers must report that as +// unknown rather than as zero, and must never substitute a related model's +// price — pricing one version off another is the defect this replaced. +func CostFor(model string) (ModelCost, bool) { + p, token, _, ok := ProviderForToken(model) + if !ok { + return ModelCost{}, false + } + token = resolveAlias(strings.TrimSpace(token)) + if c, ok := costOf(p.lookupExact(token)); ok { + return c, true + } + // An unpinned snapshot prices as the line it snapshots: providers publish one + // price per model version, not per dated build. + if bare := snapshotSuffix.ReplaceAllString(token, ""); bare != token { + if c, ok := costOf(p.lookupExact(bare)); ok { + return c, true + } + } + // Fall back to matching on the parsed version, which catches spellings the + // catalog does not hold verbatim — notably OpenRouter's dotted + // "claude-opus-4.5" against the registry's dashed "claude-opus-4-5". The + // version must match exactly: the prefix matching used for routing would + // resolve "claude-opus-4" onto 4.8 and bill a retired model at a rate it + // never charged. + identity, ok := p.ParseIdentity(token) + if !ok || identity.Version == "" { + return ModelCost{}, false + } + version := normalizeModelVersion(identity.Version) + for _, m := range knownModels { + if m.Provider != identity.Provider || m.Family != identity.Family { + continue + } + if normalizeModelVersion(m.Version) != version { + continue + } + if c, ok := costOf(m, true); ok { + return c, true + } + } + return ModelCost{}, false +} + +func costOf(m KnownModel, found bool) (ModelCost, bool) { + if !found || m.Cost == nil { + return ModelCost{}, false + } + return *m.Cost, true +} diff --git a/pkg/api/registry/cost_test.go b/pkg/api/registry/cost_test.go new file mode 100644 index 0000000..ec9cbeb --- /dev/null +++ b/pkg/api/registry/cost_test.go @@ -0,0 +1,120 @@ +package registry + +import "testing" + +// publishedRates are the vendors' list prices in USD per million tokens, taken +// from the provider pricing pages rather than from captain's own catalog, so +// this test fails if a regenerated models.json ever drifts from reality. +var publishedRates = map[string]ModelCost{ + "claude-opus-5": {Input: 5, Output: 25, CacheRead: 0.5, CacheWrite: 6.25}, + "claude-opus-4-8": {Input: 5, Output: 25, CacheRead: 0.5, CacheWrite: 6.25}, + "claude-sonnet-5": {Input: 2, Output: 10, CacheRead: 0.2, CacheWrite: 2.5}, + "claude-sonnet-4-6": {Input: 3, Output: 15, CacheRead: 0.3, CacheWrite: 3.75}, + "claude-fable-5": {Input: 10, Output: 50, CacheRead: 1, CacheWrite: 12.5}, + "claude-haiku-4-5": {Input: 1, Output: 5, CacheRead: 0.1, CacheWrite: 1.25}, +} + +// TestCostForMatchesPublishedRates is the regression guard for the family-table +// pricing defect: classifying on the substring "opus" alone billed every Opus at +// $15/$75 (the 4.1 rate) and priced Fable — which matched no family at all — at +// the Sonnet fallback. Each id here must carry its own version's price. +func TestCostForMatchesPublishedRates(t *testing.T) { + for id, want := range publishedRates { + t.Run(id, func(t *testing.T) { + got, ok := CostFor(id) + if !ok { + t.Fatalf("CostFor(%q) found no price", id) + } + if got != want { + t.Errorf("CostFor(%q) = %+v, want %+v", id, got, want) + } + }) + } +} + +// TestCostForDistinguishesVersionsAndFamilies pins that pricing is per-model, +// not per-family: the two Opus generations and the four Claude families must not +// collapse onto one rate. +func TestCostForDistinguishesVersionsAndFamilies(t *testing.T) { + opus, _ := CostFor("claude-opus-5") + sonnetOld, _ := CostFor("claude-sonnet-4-6") + sonnetNew, _ := CostFor("claude-sonnet-5") + fable, _ := CostFor("claude-fable-5") + + if opus.Input == 15 || opus.Output == 75 { + t.Errorf("claude-opus-5 priced at the retired Opus 4.1 rate: %+v", opus) + } + if sonnetNew == sonnetOld { + t.Errorf("Sonnet 5 and Sonnet 4.6 share one rate %+v; versions must price apart", sonnetNew) + } + if fable == sonnetOld { + t.Errorf("Fable priced at the Sonnet fallback %+v instead of its own rate", fable) + } +} + +// TestCostForAcceptsIDSpellings covers the id forms that reach pricing: catalog +// ids, provider namespaces, codename aliases, dated provider snapshots, and +// OpenRouter's dotted version spelling. +func TestCostForAcceptsIDSpellings(t *testing.T) { + opus45, ok := CostFor("claude-opus-4-5") + if !ok { + t.Fatal("claude-opus-4-5 must be priced") + } + for _, spelling := range []string{ + "anthropic/claude-opus-4-5", + "claude-opus-4-5-20251101", + "anthropic/claude-opus-4.5", + } { + got, ok := CostFor(spelling) + if !ok { + t.Errorf("CostFor(%q) found no price", spelling) + continue + } + if got != opus45 { + t.Errorf("CostFor(%q) = %+v, want the claude-opus-4-5 rate %+v", spelling, got, opus45) + } + } + + viaAlias, ok := CostFor("sol") + if !ok { + t.Fatal("codename alias sol must resolve to a priced model") + } + if exact, _ := CostFor("gpt-5.6-sol"); viaAlias != exact { + t.Errorf("alias sol priced %+v, want gpt-5.6-sol rate %+v", viaAlias, exact) + } +} + +// TestCostForRejectsUnpricedIDs pins that a miss stays a miss. Returning a +// nearby model's rate is what made the old table wrong; an unknown id must be +// reported as unknown so callers can show "unpriced" instead of a fabricated +// number. +func TestCostForRejectsUnpricedIDs(t *testing.T) { + for _, id := range []string{ + "", + "totally-unknown-model-zzz", + "claude-opus", // family with no version: must not inherit the flagship rate + "gpt", // ditto on OpenAI + "claude-opus-99-does-not-exist", + // Retired, and priced $15/$75 when it shipped. Version matching must be + // exact, or the routing rules would resolve it onto 4.8's $5/$25. + "claude-opus-4", + } { + if got, ok := CostFor(id); ok { + t.Errorf("CostFor(%q) = %+v, want no price", id, got) + } + } +} + +// TestEveryCatalogModelIsPriced guards the generator: a models.json row with no +// cost block would silently price that model at zero everywhere. +func TestEveryCatalogModelIsPriced(t *testing.T) { + for _, m := range knownModels { + if m.Cost == nil { + t.Errorf("catalog model %q carries no cost block", m.ID) + continue + } + if m.Cost.Input <= 0 || m.Cost.Output <= 0 { + t.Errorf("catalog model %q has non-positive price %+v", m.ID, *m.Cost) + } + } +} diff --git a/pkg/api/registry/effort.go b/pkg/api/registry/effort.go new file mode 100644 index 0000000..cd31f82 --- /dev/null +++ b/pkg/api/registry/effort.go @@ -0,0 +1,40 @@ +package registry + +import "fmt" + +// Effort is the per-request reasoning effort. captain owns this enum (including +// Codex's xhigh/max/ultra tiers); "" means backend default. +type Effort string + +const ( + EffortNone Effort = "" + EffortLow Effort = "low" + EffortMedium Effort = "medium" + EffortHigh Effort = "high" + EffortXHigh Effort = "xhigh" + EffortMax Effort = "max" + EffortUltra Effort = "ultra" +) + +// AllEfforts lists the non-empty effort tiers in ascending order. +func AllEfforts() []Effort { + return []Effort{EffortLow, EffortMedium, EffortHigh, EffortXHigh, EffortMax, EffortUltra} +} + +// Valid reports whether e is a recognised effort tier (including none/""). +func (e Effort) Valid() bool { + switch e { + case EffortNone, EffortLow, EffortMedium, EffortHigh, EffortXHigh, EffortMax, EffortUltra: + return true + default: + return false + } +} + +// Validate fails loud on an unknown effort tier, naming the valid set. +func (e Effort) Validate() error { + if e.Valid() { + return nil + } + return fmt.Errorf("invalid reasoning effort %q; want one of: low, medium, high, xhigh, max, ultra", e) +} diff --git a/pkg/api/registry/effort_support.go b/pkg/api/registry/effort_support.go new file mode 100644 index 0000000..7b27d02 --- /dev/null +++ b/pkg/api/registry/effort_support.go @@ -0,0 +1,54 @@ +package registry + +import ( + "fmt" + "strings" +) + +// ModelEfforts returns the effort tiers a backend/model pair accepts, when the +// catalog knows the combination. ok is false for models the catalog has never +// heard of, which keeps effort validation permissive for brand-new ids. +// +// The lookup is exact: callers pass an already-resolved id, and resolving again +// here would answer for a sibling model (see RegistryModelDef). +func ModelEfforts(backend Backend, model string) (supported []Effort, defaultEffort Effort, ok bool) { + p, mode, found := ProviderFor(backend) + if !found { + return nil, EffortNone, false + } + m, found := p.Lookup(model) + if !found || !p.availableFor(m, mode) { + return nil, EffortNone, false + } + return append([]Effort(nil), m.SupportedEfforts...), m.DefaultEffort, true +} + +// ValidateEffort enforces model-aware effort tiers without requiring the catalog +// to be exhaustive: an unknown model accepts any valid tier, while a known one +// accepts only the tiers the catalog records for it. +func ValidateEffort(backend Backend, model string, effort Effort) error { + if err := effort.Validate(); err != nil { + return err + } + if effort == EffortNone { + return nil + } + supported, _, known := ModelEfforts(backend, model) + if !known { + return nil + } + if len(supported) == 0 { + return fmt.Errorf("model %q on %s does not support a reasoning effort", model, backend) + } + for _, candidate := range supported { + if candidate == effort { + return nil + } + } + values := make([]string, 0, len(supported)) + for _, candidate := range supported { + values = append(values, string(candidate)) + } + return fmt.Errorf("model %q on %s does not support reasoning effort %q; want one of: %s", + model, backend, effort, strings.Join(values, ", ")) +} diff --git a/pkg/api/registry/embed.go b/pkg/api/registry/embed.go new file mode 100644 index 0000000..2505d13 --- /dev/null +++ b/pkg/api/registry/embed.go @@ -0,0 +1,74 @@ +package registry + +import ( + _ "embed" + "encoding/json" + "fmt" +) + +// modelsJSON is the generated model catalog produced by `task models:update` +// (pkg/ai/internal/gen-model-registry). It is the models.dev snapshot with that +// generator's patches.json overlaid. Regenerate it after editing patches.json. +// +//go:embed models.json +var modelsJSON []byte + +// KnownModel is one row of the generated catalog: a provider-native model id +// plus the capability metadata captain needs to route, price, and gate it. +// +// Aliases and SupersededBy are data, not code, on purpose. They used to be +// hardcoded switches in pkg/ai, which meant pkg/api's spec decoder could not see +// them and rejected models the pkg/ai selector accepted — `--model agent:sol` +// errored while the same value in prompt frontmatter ran. Keeping them here is +// what lets a single parser serve both entry points. +type KnownModel struct { + ID string `json:"id"` + Provider string `json:"provider"` + Family string `json:"family"` + Version string `json:"version"` + Label string `json:"label"` + ReleaseDate string `json:"releaseDate,omitempty"` + Reasoning bool `json:"reasoning,omitempty"` + Temperature bool `json:"temperature,omitempty"` + ContextWindow int `json:"contextWindow,omitempty"` + InputMediaTypes []string `json:"inputMediaTypes,omitempty"` + // Cost is the model's published list price. Nil means the catalog snapshot + // carries no price for this id — a miss, never a free model. + Cost *ModelCost `json:"cost,omitempty"` + Preferred bool `json:"preferred,omitempty"` + AdaptiveThinking bool `json:"adaptiveThinking,omitempty"` + Availability []string `json:"availability,omitempty"` + SupportedEfforts []Effort `json:"supportedEfforts,omitempty"` + DefaultEffort Effort `json:"defaultEffort,omitempty"` + Priority int `json:"priority,omitempty"` + + // Aliases are extra input tokens that resolve to this model — the codenames + // "sol"/"terra"/"luna" on the gpt-5.6-* line. Input conveniences only: an + // alias is never returned to a caller. + Aliases []string `json:"aliases,omitempty"` + // SupersededBy names the model that replaces this one. A superseded id still + // parses, so old specs keep loading, but resolves to its successor. + SupersededBy string `json:"supersededBy,omitempty"` +} + +// knownModels is the parsed model catalog. It is a package-level var (not an +// init) so Go's initialization ordering guarantees it is populated before the +// provider descriptors, which index it. +var knownModels = mustParseModels(modelsJSON) + +func mustParseModels(data []byte) []KnownModel { + var models []KnownModel + if err := json.Unmarshal(data, &models); err != nil { + panic(fmt.Sprintf("pkg/api/registry: invalid embedded models.json: %v", err)) + } + if len(models) == 0 { + panic("pkg/api/registry: embedded models.json is empty") + } + return models +} + +// KnownModels returns every catalog row. The slice is copied so callers cannot +// mutate the registry. +func KnownModels() []KnownModel { + return append([]KnownModel(nil), knownModels...) +} diff --git a/pkg/api/registry/embed_test.go b/pkg/api/registry/embed_test.go new file mode 100644 index 0000000..289b559 --- /dev/null +++ b/pkg/api/registry/embed_test.go @@ -0,0 +1,82 @@ +package registry + +import "testing" + +// TestEmbeddedRegistryLoads guards the go:embed + JSON parse path: the generated +// models.json must load into a non-empty slice whose every entry is well-formed +// and provider-tagged with a known provider. +func TestEmbeddedRegistryLoads(t *testing.T) { + if len(knownModels) == 0 { + t.Fatal("knownModels is empty; models.json failed to load") + } + for _, m := range knownModels { + if m.ID == "" || m.Provider == "" || m.Family == "" { + t.Errorf("entry %+v is missing id/provider/family", m) + } + if _, ok := ProviderByName(m.Provider); !ok { + t.Errorf("entry %q has unknown provider %q", m.ID, m.Provider) + } + } +} + +// TestRegistryDerivedCapabilities locks the capability flags the generator +// derives from models.dev: reasoning, temperature support, and the adaptive-vs- +// enabled thinking schema. The opus-4-8/4-7 adaptive values guard the 400 that +// legacy `thinking:{type:enabled}` triggers on those models. +func TestRegistryDerivedCapabilities(t *testing.T) { + cases := []struct { + id string + wantReasoning bool + wantTemp bool + wantPreferred bool + wantAdaptive bool + }{ + {"claude-opus-5", true, false, true, true}, + {"claude-sonnet-5", true, false, true, true}, + {"claude-fable-5", true, false, true, true}, + {"claude-opus-4-8", true, false, true, true}, + {"claude-opus-4-7", true, false, false, true}, + {"claude-sonnet-4-6", true, true, false, false}, + {"claude-haiku-4-5", true, true, true, false}, + } + for _, tc := range cases { + t.Run(tc.id, func(t *testing.T) { + m, ok := Anthropic.Lookup(tc.id) + if !ok { + t.Fatalf("model %q not found in registry", tc.id) + } + if m.Reasoning != tc.wantReasoning { + t.Errorf("Reasoning = %v, want %v", m.Reasoning, tc.wantReasoning) + } + if m.Temperature != tc.wantTemp { + t.Errorf("Temperature = %v, want %v", m.Temperature, tc.wantTemp) + } + if m.Preferred != tc.wantPreferred { + t.Errorf("Preferred = %v, want %v", m.Preferred, tc.wantPreferred) + } + if m.AdaptiveThinking != tc.wantAdaptive { + t.Errorf("AdaptiveThinking = %v, want %v", m.AdaptiveThinking, tc.wantAdaptive) + } + }) + } +} + +// TestRegistryDataCarriesAliasesAndSuccessors pins the knowledge that used to be +// hardcoded in pkg/ai (normalizeCodexVariantAlias, isSupersededRegistryExact) and +// is now catalog data, reachable from every entry point. +func TestRegistryDataCarriesAliasesAndSuccessors(t *testing.T) { + for alias, want := range map[string]string{"sol": "gpt-5.6-sol", "terra": "gpt-5.6-terra", "luna": "gpt-5.6-luna"} { + if got := resolveAlias(alias); got != want { + t.Errorf("resolveAlias(%q) = %q, want %q", alias, got, want) + } + } + for _, retired := range []string{"claude-sonnet-4-5", "claude-sonnet-4-5-20250929"} { + m, ok := Anthropic.Lookup(retired) + if !ok { + t.Fatalf("retired model %q missing from registry", retired) + } + if m.SupersededBy != "claude-sonnet-4-6" { + t.Errorf("%q supersededBy = %q, want claude-sonnet-4-6", retired, m.SupersededBy) + } + } +} diff --git a/pkg/api/registry/errors.go b/pkg/api/registry/errors.go new file mode 100644 index 0000000..b8f9456 --- /dev/null +++ b/pkg/api/registry/errors.go @@ -0,0 +1,155 @@ +package registry + +import ( + "context" + "errors" + "net" + "net/url" + "regexp" + "strings" +) + +// ErrorClass is what went wrong with a provider call, normalized across +// providers so retry, fallback, and budget logic can reason about it without +// re-reading error prose. +type ErrorClass string + +const ( + // ErrorNone means "not a recognized provider failure". + ErrorNone ErrorClass = "" + // ErrorRateLimit is a 429 / quota exhaustion. Retryable, ideally after a wait. + ErrorRateLimit ErrorClass = "rate_limit" + // ErrorOverloaded is a 503/529 or explicit overload. Retryable. + ErrorOverloaded ErrorClass = "overloaded" + // ErrorNetwork is a transport failure: dial, reset, EOF, deadline. Retryable. + ErrorNetwork ErrorClass = "network" + // ErrorContextLength means the request exceeded the model's context window. + // NOT retryable: the identical request fails identically, so a retry only + // burns tokens. Recoverable by trimming input or moving to a larger model. + ErrorContextLength ErrorClass = "context_length" + // ErrorAuth is a missing/rejected credential. Not retryable. + ErrorAuth ErrorClass = "auth" + // ErrorModelUnavailable is a provider-confirmed model selection failure. Not + // retryable on the same model; an explicit fallback may still recover. + ErrorModelUnavailable ErrorClass = "model_unavailable" + // ErrorInvalidRequest is a malformed request. Not retryable. + ErrorInvalidRequest ErrorClass = "invalid_request" +) + +// Retryable reports whether the same request may succeed if tried again. +func (c ErrorClass) Retryable() bool { + switch c { + case ErrorRateLimit, ErrorOverloaded, ErrorNetwork: + return true + default: + return false + } +} + +// httpStatusError is implemented by errors that carry a real HTTP status. Status +// is read structurally where available, because reading it out of prose is how +// captain used to classify "1429 tokens" as a rate limit. +type httpStatusError interface{ StatusCode() int } + +// Status codes must stand alone. `strings.Contains(msg, "429")` — the previous +// implementation — also matched "1429 tokens", "id=4290", and "request 5031", +// turning ordinary messages into spurious retries. +// +// \b is not enough either: it treats "-" and "." as boundaries, so "gpt-429-turbo" +// and version strings like "4.29.1" would still match. The delimiter set +// therefore excludes word characters, "." and "-", leaving whitespace, +// punctuation, and string edges. +const ( + statusPrefix = `(^|[^\w.-])` + statusSuffix = `([^\w.-]|$)` +) + +var ( + reRateLimitStatus = regexp.MustCompile(statusPrefix + `429` + statusSuffix) + reOverloadedStatus = regexp.MustCompile(statusPrefix + `(503|529)` + statusSuffix) + reAuthStatus = regexp.MustCompile(statusPrefix + `(401|403)` + statusSuffix) +) + +// ClassifyError applies the provider-independent heuristics: structured signals +// first (wrapped sentinels, net.Error, url.Error, HTTP status), then bounded +// text matching. Providers refine what is left via Provider.ClassifyError. +func ClassifyError(err error) ErrorClass { + if err == nil { + return ErrorNone + } + + // 1. Structure beats prose. + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return ErrorNetwork + } + var status httpStatusError + if errors.As(err, &status) { + switch code := status.StatusCode(); { + case code == 429: + return ErrorRateLimit + case code == 503 || code == 529: + return ErrorOverloaded + case code == 401 || code == 403: + return ErrorAuth + case code == 404: + return ErrorModelUnavailable + case code >= 500: + return ErrorOverloaded + case code == 400 || code == 422: + return ErrorInvalidRequest + } + } + var netErr net.Error + if errors.As(err, &netErr) { + return ErrorNetwork + } + var urlErr *url.Error + if errors.As(err, &urlErr) { + return ErrorNetwork + } + + // 2. Bounded text matching for providers that only return prose. + msg := strings.ToLower(err.Error()) + switch { + case containsAny(msg, "context length", "context window", "maximum context", + "context_length_exceeded", "too many tokens", "prompt is too long", + "request too large", "string too long"): + return ErrorContextLength + case reRateLimitStatus.MatchString(msg), strings.Contains(msg, "rate limit"), + strings.Contains(msg, "rate_limit"), strings.Contains(msg, "quota exceeded"), + strings.Contains(msg, "resource_exhausted"), strings.Contains(msg, "too many requests"): + return ErrorRateLimit + case reOverloadedStatus.MatchString(msg), strings.Contains(msg, "overloaded"), + strings.Contains(msg, "server_error"), strings.Contains(msg, "service unavailable"), + strings.Contains(msg, "try again later"): + return ErrorOverloaded + case containsAny(msg, "timeout", "timed out", "connection reset", "connection refused", + "broken pipe", "unexpected eof", "no such host", "tls handshake"): + return ErrorNetwork + case reAuthStatus.MatchString(msg), containsAny(msg, "unauthorized", "invalid api key", + "invalid_api_key", "authentication", "permission denied", "missing api key", + "no api key", "api key is not set", "api_key is not set"): + return ErrorAuth + } + return ErrorNone +} + +// ClassifyError refines the shared verdict with this provider's own error +// vocabulary. Providers set classifyErr only where their wording is genuinely +// their own; everything common stays in the shared heuristics. +func (p *Provider) ClassifyError(err error) ErrorClass { + base := ClassifyError(err) + if err == nil || p.classifyErr == nil { + return base + } + return p.classifyErr(err, base) +} + +func containsAny(haystack string, needles ...string) bool { + for _, needle := range needles { + if strings.Contains(haystack, needle) { + return true + } + } + return false +} diff --git a/pkg/api/registry/errors_test.go b/pkg/api/registry/errors_test.go new file mode 100644 index 0000000..ec5511d --- /dev/null +++ b/pkg/api/registry/errors_test.go @@ -0,0 +1,136 @@ +package registry + +import ( + "context" + "errors" + "fmt" + "net" + "net/url" + "testing" +) + +// statusErr is a provider error carrying a real HTTP status. +type statusErr struct { + code int + msg string +} + +func (e statusErr) Error() string { return e.msg } +func (e statusErr) StatusCode() int { return e.code } + +// TestClassifyErrorReadsStructureNotProse pins that a real status code decides +// the class even when the message says nothing about it. +func TestClassifyErrorReadsStructureNotProse(t *testing.T) { + cases := []struct { + code int + want ErrorClass + }{ + {429, ErrorRateLimit}, + {503, ErrorOverloaded}, + {529, ErrorOverloaded}, + {500, ErrorOverloaded}, + {401, ErrorAuth}, + {403, ErrorAuth}, + {404, ErrorModelUnavailable}, + {400, ErrorInvalidRequest}, + } + for _, tc := range cases { + t.Run(fmt.Sprint(tc.code), func(t *testing.T) { + err := fmt.Errorf("generate: %w", statusErr{code: tc.code, msg: "provider said no"}) + if got := ClassifyError(err); got != tc.want { + t.Errorf("ClassifyError(status %d) = %q, want %q", tc.code, got, tc.want) + } + }) + } +} + +// TestClassifyErrorDoesNotMatchNumbersInProse is the regression this refactor +// exists to prevent. The previous implementation did +// strings.Contains(msg, "429"), so any message that merely contained those +// digits — a token count, an id, a byte offset — was treated as a rate limit and +// retried for nothing. +func TestClassifyErrorDoesNotMatchNumbersInProse(t *testing.T) { + falsePositives := []string{ + "request rejected: 1429 tokens exceeds nothing in particular", + "model gpt-429-turbo is fine", + "trace id=4290 completed", + "finished after 5031 ms", + "embedding dimension 15039 mismatch", + } + for _, msg := range falsePositives { + t.Run(msg, func(t *testing.T) { + if got := ClassifyError(errors.New(msg)); got.Retryable() { + t.Errorf("ClassifyError(%q) = %q, which retries; digits inside prose are not a status code", msg, got) + } + }) + } + + // A real status in prose still classifies: the boundary is the point, not + // refusing to read text at all. + for msg, want := range map[string]ErrorClass{ + "provider returned 429 Too Many Requests": ErrorRateLimit, + "upstream 503 service unavailable": ErrorOverloaded, + "HTTP 401 unauthorized": ErrorAuth, + } { + if got := ClassifyError(errors.New(msg)); got != want { + t.Errorf("ClassifyError(%q) = %q, want %q", msg, got, want) + } + } +} + +func TestClassifyErrorNetworkAndContext(t *testing.T) { + cases := map[string]struct { + err error + want ErrorClass + }{ + "deadline": {context.DeadlineExceeded, ErrorNetwork}, + "canceled": {context.Canceled, ErrorNetwork}, + "url error": {&url.Error{Op: "Post", URL: "https://api.example", Err: errors.New("dial tcp: connect: connection refused")}, ErrorNetwork}, + "net error": {&net.OpError{Op: "dial", Err: errors.New("connection refused")}, ErrorNetwork}, + "reset": {errors.New("read: connection reset by peer"), ErrorNetwork}, + "context len": {errors.New("This model's maximum context length is 200000 tokens"), ErrorContextLength}, + "ctx len code": {errors.New("context_length_exceeded"), ErrorContextLength}, + "prompt long": {errors.New("prompt is too long: 250000 tokens > 200000 maximum"), ErrorContextLength}, + "nil": {nil, ErrorNone}, + "unknown": {errors.New("something else entirely"), ErrorNone}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + if got := ClassifyError(tc.err); got != tc.want { + t.Errorf("ClassifyError(%v) = %q, want %q", tc.err, got, tc.want) + } + }) + } +} + +// TestContextLengthIsNotRetryable: an oversized request fails identically on +// every retry, so retrying it only spends tokens to reach the same error. +func TestContextLengthIsNotRetryable(t *testing.T) { + if ErrorContextLength.Retryable() { + t.Error("ErrorContextLength must not be retryable: the same request fails the same way") + } + // A context-length message that also contains a token count must not be + // rescued into a retryable class by the digits it carries. + err := errors.New("prompt is too long: 429000 tokens > 200000 maximum") + if got := ClassifyError(err); got != ErrorContextLength { + t.Errorf("ClassifyError(%v) = %q, want %q; the token count must not win over the real cause", err, got, ErrorContextLength) + } +} + +func TestRetryableClasses(t *testing.T) { + retryable := map[ErrorClass]bool{ + ErrorRateLimit: true, + ErrorOverloaded: true, + ErrorNetwork: true, + ErrorContextLength: false, + ErrorAuth: false, + ErrorModelUnavailable: false, + ErrorInvalidRequest: false, + ErrorNone: false, + } + for class, want := range retryable { + if got := class.Retryable(); got != want { + t.Errorf("%q.Retryable() = %v, want %v", class, got, want) + } + } +} diff --git a/pkg/api/registry/generation.go b/pkg/api/registry/generation.go new file mode 100644 index 0000000..3461439 --- /dev/null +++ b/pkg/api/registry/generation.go @@ -0,0 +1,99 @@ +package registry + +// defaultMaxOutputTokens is the visible-answer budget Anthropic requires on every +// request; the extended-thinking budget is added on top of it. +const defaultMaxOutputTokens = 4096 + +// GenerationConfig builds the provider-native generation config for one request, +// translating reasoning effort and temperature into this provider's controls and +// gating them by the model's catalog capabilities: +// +// - effort is dropped for models that do not support reasoning; +// - temperature is dropped for models that do not support it; +// - Anthropic reasoning uses the adaptive schema (thinking:{type:adaptive} + +// output_config.effort) for adaptive models and the legacy enabled schema +// (thinking:{type:enabled, budget_tokens}) otherwise. +// +// Returns nil when there is nothing to send. +func (p *Provider) GenerationConfig(mode RuntimeMode, model string, effort Effort, maxTokens int, temperature *float64) map[string]any { + caps := p.modelCaps(mode, model) + cfg := map[string]any{} + if temperature != nil && caps.Temperature { + cfg["temperature"] = *temperature + } + if p.genConfig != nil { + p.genConfig(cfg, caps, effort, maxTokens) + } + if len(cfg) == 0 { + return nil + } + return cfg +} + +// modelCaps resolves a model token to its catalog row so its capability flags can +// gate the generation config. An unresolved model reports no capabilities. +func (p *Provider) modelCaps(mode RuntimeMode, model string) KnownModel { + exact, ok := p.ResolveExact(mode, model) + if !ok { + return KnownModel{} + } + m, _ := p.lookupExact(exact) + return m +} + +// openaiGenerationConfig sends the reasoning tier as-is; captain's effort enum is +// validated against the model before execution. +func openaiGenerationConfig(cfg map[string]any, caps KnownModel, effort Effort, _ int) { + if caps.Reasoning && effort != EffortNone { + cfg["reasoning_effort"] = string(effort) + } +} + +func googleGenerationConfig(cfg map[string]any, caps KnownModel, effort Effort, _ int) { + if caps.Reasoning && effort != EffortNone { + cfg["thinkingConfig"] = map[string]any{"thinkingBudget": thinkingBudget(effort)} + } +} + +// anthropicGenerationConfig always sends max_tokens, which Anthropic requires. +// The thinking budget counts inside max_tokens, so room for the visible answer is +// reserved on top of it — but only when thinking is actually sent. +func anthropicGenerationConfig(cfg map[string]any, caps KnownModel, effort Effort, maxTokens int) { + base := maxTokens + if base <= 0 { + base = defaultMaxOutputTokens + } + if !caps.Reasoning || effort == EffortNone { + cfg["max_tokens"] = base + return + } + budget := thinkingBudget(effort) + cfg["max_tokens"] = base + budget + if caps.AdaptiveThinking { + cfg["thinking"] = map[string]any{"type": "adaptive"} + cfg["output_config"] = map[string]any{"effort": string(effort)} + return + } + cfg["thinking"] = map[string]any{"type": "enabled", "budget_tokens": budget} +} + +// DeepSeek selects reasoning by model id (deepseek-reasoner vs deepseek-chat) +// rather than a per-request knob, so it contributes no effort config — hence no +// genConfig hook on that descriptor. + +// thinkingBudget maps a reasoning effort to an extended-thinking token budget +// (shared by Anthropic enabled-thinking and Gemini). xhigh is captain's top tier. +func thinkingBudget(e Effort) int { + switch e { + case EffortLow: + return 2048 + case EffortMedium: + return 8192 + case EffortHigh: + return 24576 + case EffortXHigh, EffortMax, EffortUltra: + return 32768 + default: + return 0 + } +} diff --git a/pkg/api/registry/identity.go b/pkg/api/registry/identity.go new file mode 100644 index 0000000..0cfa8f8 --- /dev/null +++ b/pkg/api/registry/identity.go @@ -0,0 +1,464 @@ +package registry + +import ( + "sort" + "strings" +) + +// ModelIdentity is captain's parsed model key. It is deliberately provider +// native: aliases and mode prefixes are input conveniences only, never the value +// sent to a backend. +type ModelIdentity struct { + Provider string + Family string + Version string +} + +// ProviderForToken returns the provider that claims a model token, along with +// the token stripped of any provider namespace and the mode its own prefix +// implies. It is the single claim step: what InferBackend, splitModelProvider, +// and selectorModelFamily each used to answer differently. +// +// An explicit namespace ("anthropic/…", "googleai/…") wins. Otherwise the +// providers are tried in canonical order and the first claim wins; failing that, +// a multi-segment id is retried on its last segment so proxied names such as +// "openrouter/anthropic/claude-x" still resolve. +func ProviderForToken(token string) (*Provider, string, RuntimeMode, bool) { + token = strings.TrimSpace(token) + if token == "" { + return nil, "", "", false + } + // Claim on the alias' target: a codename such as "sol" names no family by + // itself. Doing this here rather than deep in the resolver is what makes + // `--model sol` and `--model agent:sol` agree — the selector grammar used to + // resolve aliases before claiming while the compact grammar did not, so the + // two disagreed about whether the model existed at all. + if aliased := resolveAlias(token); aliased != token { + token = aliased + } + if i := strings.IndexByte(token, '/'); i >= 0 { + if p, ok := ProviderByName(token[:i]); ok { + rest := strings.TrimPrefix(token[i+1:], "models/") + mode, claimed := p.claim(rest) + if !claimed { + mode = ModeAPI + } + return p, rest, mode, true + } + } + canonical := canonicalModelToken(token) + for _, p := range Providers() { + if mode, ok := p.claim(canonical); ok { + return p, strings.TrimPrefix(token, "models/"), mode, true + } + } + // Retry on the last path segment: an id proxied through another namespace + // ("openrouter/anthropic/claude-x") keeps its full name but resolves off the + // family it actually names. + if i := strings.LastIndexByte(token, '/'); i >= 0 { + for _, p := range Providers() { + if mode, ok := p.claim(canonicalModelToken(token[i+1:])); ok { + return p, token, mode, true + } + } + } + return nil, "", "", false +} + +// StripProviderPrefix removes any known provider namespace from a model id. +func StripProviderPrefix(model string) string { + p, token, _, ok := ProviderForToken(model) + if !ok { + return strings.TrimPrefix(strings.TrimSpace(model), "models/") + } + return strings.TrimPrefix(strings.TrimSpace(p.bareID(token)), "models/") +} + +func canonicalModelToken(model string) string { + model = strings.ToLower(strings.TrimSpace(model)) + model = strings.TrimPrefix(model, "models/") + model = strings.ReplaceAll(model, "_", "-") + model = strings.ReplaceAll(model, " ", "-") + return model +} + +// resolveAlias maps a codename onto its exact model id ("sol" → "gpt-5.6-sol"). +// Aliases are registry data, so every entry point sees them. +func resolveAlias(model string) string { + needle := canonicalModelToken(model) + for _, m := range knownModels { + for _, alias := range m.Aliases { + if canonicalModelToken(alias) == needle { + return m.ID + } + } + } + return strings.TrimSpace(model) +} + +// ParseIdentity parses a model token into this provider's family/version tuple. +// It is deliberately permissive about input prefixes. +func (p *Provider) ParseIdentity(model string) (ModelIdentity, bool) { + token := canonicalModelToken(p.bareID(model)) + for _, trim := range p.identityTrim { + token = strings.TrimPrefix(token, trim) + } + for _, empty := range p.emptyTokens { + if token == empty && p.emptyFamily != "" { + return ModelIdentity{Provider: p.Name, Family: p.emptyFamily}, true + } + } + family, version, ok := splitKnownFamily(token, p.families) + return ModelIdentity{Provider: p.Name, Family: family, Version: version}, ok +} + +func splitKnownFamily(token string, families []string) (family, version string, ok bool) { + for _, candidate := range families { + if token == candidate { + return candidate, "", true + } + prefix := candidate + "-" + if strings.HasPrefix(token, prefix) { + return candidate, normalizeModelVersion(strings.TrimPrefix(token, prefix)), true + } + } + return "", "", false +} + +// availableFor reports whether a catalog row is offered on a mode. The registry +// annotates rows with "api" or "codex" availability; the codex modes are the +// only ones that read the codex list. +func (p *Provider) availableFor(m KnownModel, mode RuntimeMode) bool { + if len(m.Availability) == 0 { + return true + } + target := "api" + if p == OpenAI && mode != ModeAPI { + target = "codex" + } + for _, available := range m.Availability { + if strings.EqualFold(strings.TrimSpace(available), target) { + return true + } + } + return false +} + +func (p *Provider) lookupExact(model string) (KnownModel, bool) { + needle := canonicalModelToken(StripProviderPrefix(model)) + for _, m := range knownModels { + if m.Provider != p.Name { + continue + } + if canonicalModelToken(m.ID) == needle { + return m, true + } + } + return KnownModel{}, false +} + +// ResolveExact resolves a user-written model token into the exact id this +// provider's mode should receive. It accepts aliases, namespaces, superseded +// ids, and family shorthands, and never returns any of them. +// +// ok is false when the token could not be resolved to a catalog row; the token +// is still returned (namespace-stripped) so an unknown-but-explicit model still +// reaches its backend — that is how new provider models work before the catalog +// snapshot catches up. +func (p *Provider) ResolveExact(mode RuntimeMode, model string) (string, bool) { + model = strings.TrimSpace(model) + if model == "" { + return "", false + } + // A bare coding-agent sentinel ("codex") means "this provider's current + // model". It is not honoured on the API mode, where "claude" stays a literal. + if mode != ModeAPI && p.AgentName != "" && strings.EqualFold(model, p.AgentName) { + if m, ok := p.latestModel(mode, ""); ok { + return m.ID, true + } + return model, false + } + model = resolveAlias(model) + if m, ok := p.lookupExact(model); ok { + if m.SupersededBy != "" { + if successor, ok := p.lookupExact(m.SupersededBy); ok && p.availableFor(successor, mode) { + return successor.ID, true + } + } else if p.availableFor(m, mode) { + return m.ID, true + } + } + identity, ok := p.ParseIdentity(model) + if !ok { + return StripProviderPrefix(model), false + } + if shouldResolveLatestVersionLine(identity.Version) { + if m, ok := p.latestModelForVersionLine(mode, identity); ok { + return m.ID, true + } + } + if m, ok := p.resolveIdentity(mode, identity); ok { + return m.ID, true + } + if identity.Version != "" { + return StripProviderPrefix(model), false + } + if m, ok := p.latestModel(mode, identity.Family); ok { + return m.ID, true + } + return StripProviderPrefix(model), false +} + +// Availability distinguishes an unknown model from a catalog model that is +// intentionally not offered on a mode. +func (p *Provider) Availability(mode RuntimeMode, model string) (known, available bool) { + m, ok := p.lookupExact(resolveAlias(model)) + if !ok { + return false, false + } + return true, p.availableFor(m, mode) +} + +// Lookup returns the catalog row for a model id on this provider, accepting +// aliases and provider namespaces. +func (p *Provider) Lookup(model string) (KnownModel, bool) { + return p.lookupExact(resolveAlias(model)) +} + +func (p *Provider) resolveIdentity(mode RuntimeMode, identity ModelIdentity) (KnownModel, bool) { + candidates := make([]KnownModel, 0) + fallback := make([]KnownModel, 0) + for _, m := range knownModels { + if m.Provider != identity.Provider || m.Family != identity.Family || !p.availableFor(m, mode) { + continue + } + if identity.Version == "" || modelVersionMatches(m.Version, identity.Version) { + if m.Preferred { + candidates = append(candidates, m) + } else if identity.Version != "" { + fallback = append(fallback, m) + } + } + } + if len(candidates) == 0 { + candidates = fallback + } + if len(candidates) == 0 { + return KnownModel{}, false + } + sortModels(candidates) + return candidates[0], true +} + +func (p *Provider) latestModel(mode RuntimeMode, family string) (KnownModel, bool) { + candidates := make([]KnownModel, 0) + for _, m := range knownModels { + if !m.Preferred || m.Provider != p.Name || !p.availableFor(m, mode) { + continue + } + if family != "" && m.Family != family { + continue + } + candidates = append(candidates, m) + } + if len(candidates) == 0 { + return KnownModel{}, false + } + sortModels(candidates) + return candidates[0], true +} + +func (p *Provider) latestModelForVersionLine(mode RuntimeMode, identity ModelIdentity) (KnownModel, bool) { + if identity.Version == "" { + return KnownModel{}, false + } + major := strings.Split(normalizeModelVersion(identity.Version), ".")[0] + if major == "" { + return KnownModel{}, false + } + candidates := make([]KnownModel, 0) + for _, m := range knownModels { + if m.Provider != identity.Provider || m.Family != identity.Family || !p.availableFor(m, mode) { + continue + } + if modelVersionMatches(m.Version, major) { + candidates = append(candidates, m) + } + } + if len(candidates) == 0 { + return KnownModel{}, false + } + sortModelsForResolution(candidates) + return candidates[0], true +} + +func shouldResolveLatestVersionLine(version string) bool { + version = normalizeModelVersion(version) + if version == "" { + return false + } + for _, part := range strings.Split(version, ".") { + if !isNumericVersionPart(part) { + return false + } + } + return true +} + +func sortModels(models []KnownModel) { + sort.SliceStable(models, func(i, j int) bool { + left, right := models[i], models[j] + if priority := comparePriority(left, right); priority != 0 { + return priority < 0 + } + if left.ReleaseDate == right.ReleaseDate { + return compareModelVersions(left.ID, right.ID) > 0 + } + return left.ReleaseDate > right.ReleaseDate + }) +} + +func sortModelsForResolution(models []KnownModel) { + sort.SliceStable(models, func(i, j int) bool { + left, right := models[i], models[j] + if priority := comparePriority(left, right); priority != 0 { + return priority < 0 + } + if left.ReleaseDate == right.ReleaseDate && left.Preferred != right.Preferred { + return left.Preferred + } + if left.ReleaseDate == right.ReleaseDate { + return compareModelVersions(left.ID, right.ID) > 0 + } + return left.ReleaseDate > right.ReleaseDate + }) +} + +// comparePriority orders two rows by explicit priority, returning <0 when left +// sorts first, >0 when right does, and 0 when priority does not decide. Priority +// 0 means "unset" and always loses to an explicit priority. +func comparePriority(left, right KnownModel) int { + if left.Priority == right.Priority || (left.Priority == 0 && right.Priority == 0) { + return 0 + } + if left.Priority == 0 { + return 1 + } + if right.Priority == 0 { + return -1 + } + return left.Priority - right.Priority +} + +func modelVersionMatches(registryVersion, requested string) bool { + registryVersion = normalizeModelVersion(registryVersion) + requested = normalizeModelVersion(requested) + if requested == "" { + return true + } + if registryVersion == requested { + return true + } + return strings.HasPrefix(registryVersion, requested+".") || strings.HasPrefix(registryVersion, requested+"-") +} + +func normalizeModelVersion(version string) string { + version = strings.Trim(strings.ToLower(strings.TrimSpace(version)), "-.") + if version == "" { + return "" + } + parts := strings.Split(version, "-") + if len(parts) == 1 { + return parts[0] + } + i := 0 + out := "" + if isNumericVersionPart(parts[0]) { + out = parts[0] + i = 1 + if len(parts) > 1 && isNumericVersionPart(parts[1]) { + out += "." + parts[1] + i = 2 + } + } + if i < len(parts) { + if out != "" { + out += "-" + } + out += strings.Join(parts[i:], "-") + } + return out +} + +func isNumericVersionPart(part string) bool { + if part == "" { + return false + } + for _, r := range part { + if r < '0' || r > '9' { + return false + } + } + return true +} + +// compareModelVersions orders two model ids by their embedded numeric version +// tokens ("claude-sonnet-4-6" > "claude-sonnet-4-5"). +func compareModelVersions(left, right string) int { + lv := modelVersionNumbers(left) + rv := modelVersionNumbers(right) + if len(lv) == 0 || len(rv) == 0 { + return 0 + } + maxLen := len(lv) + if len(rv) > maxLen { + maxLen = len(rv) + } + for i := 0; i < maxLen; i++ { + l, r := 0, 0 + if i < len(lv) { + l = lv[i] + } + if i < len(rv) { + r = rv[i] + } + if l != r { + return l - r + } + } + return 0 +} + +func modelVersionNumbers(id string) []int { + id = StripProviderPrefix(id) + parts := strings.Split(strings.ToLower(id), "-") + var out []int + for _, part := range parts { + if !isModelVersionToken(part) { + continue + } + for _, piece := range strings.Split(part, ".") { + if piece == "" { + continue + } + n := 0 + for _, r := range piece { + n = n*10 + int(r-'0') + } + out = append(out, n) + } + } + return out +} + +func isModelVersionToken(value string) bool { + if value == "" { + return false + } + for _, r := range value { + if (r < '0' || r > '9') && r != '.' { + return false + } + } + return true +} diff --git a/pkg/api/registry/media.go b/pkg/api/registry/media.go new file mode 100644 index 0000000..01584d8 --- /dev/null +++ b/pkg/api/registry/media.go @@ -0,0 +1,97 @@ +package registry + +import ( + "slices" + "strings" +) + +// MediaTypesFor returns the attachment types a model accepts on a mode: the +// model's own catalog types intersected with the adapter's ceiling. An adapter +// cannot carry what it cannot send, and a model cannot read what it does not +// understand, so the answer is the intersection of the two. +func (p *Provider) MediaTypesFor(mode RuntimeMode, model string) []string { + caps, ok := p.Caps(mode) + if !ok { + return []string{} + } + modelTypes := caps.MediaTypes + if m, found := p.Lookup(model); found && len(m.InputMediaTypes) > 0 { + modelTypes = m.InputMediaTypes + } + return ClampMediaTypes(caps.MediaTypes, modelTypes) +} + +// AdapterMediaTypes returns a mode's attachment ceiling. +func (p *Provider) AdapterMediaTypes(mode RuntimeMode) []string { + caps, ok := p.Caps(mode) + if !ok { + return []string{} + } + return append([]string{}, caps.MediaTypes...) +} + +// ClampMediaTypes intersects a model's declared media types with an adapter's +// supported set, expanding wildcards on either side. +func ClampMediaTypes(adapterTypes, modelTypes []string) []string { + if len(adapterTypes) == 0 || len(modelTypes) == 0 { + return []string{} + } + out := make([]string, 0, len(modelTypes)) + for _, modelType := range modelTypes { + modelType = NormalizeMediaType(modelType) + if modelType == "" { + continue + } + for _, adapterType := range adapterTypes { + mediaType, ok := intersectMediaTypes(modelType, adapterType) + if ok && !slices.Contains(out, mediaType) { + out = append(out, mediaType) + } + } + } + return out +} + +func intersectMediaTypes(modelType, adapterType string) (string, bool) { + if modelType == adapterType { + return modelType, true + } + if strings.HasSuffix(modelType, "/*") && MediaTypeAccepted([]string{modelType}, adapterType) { + return adapterType, true + } + if strings.HasSuffix(adapterType, "/*") && MediaTypeAccepted([]string{adapterType}, modelType) { + return modelType, true + } + return "", false +} + +// NormalizeMediaType canonicalizes the shorthand forms the catalog uses. +// text/* normalizes to empty: prompt text is not an attachment. +func NormalizeMediaType(value string) string { + switch strings.ToLower(strings.TrimSpace(value)) { + case "image", "image/*": + return "image/*" + case "audio", "audio/*": + return "audio/*" + case "video", "video/*": + return "video/*" + case "pdf", "application/pdf": + return "application/pdf" + case "text", "text/*": + return "" + default: + return strings.ToLower(strings.TrimSpace(value)) + } +} + +// MediaTypeAccepted reports whether a media type matches any accepted pattern, +// honouring "type/*" wildcards. +func MediaTypeAccepted(accepted []string, mediaType string) bool { + mediaType = strings.ToLower(strings.TrimSpace(mediaType)) + for _, pattern := range accepted { + if pattern == mediaType || strings.HasSuffix(pattern, "/*") && strings.HasPrefix(mediaType, strings.TrimSuffix(pattern, "*")) { + return true + } + } + return false +} diff --git a/pkg/api/registry/mode.go b/pkg/api/registry/mode.go new file mode 100644 index 0000000..1ed93d7 --- /dev/null +++ b/pkg/api/registry/mode.go @@ -0,0 +1,58 @@ +package registry + +import "strings" + +// RuntimeMode is the mechanism that serves a model: called directly over HTTP, +// driven through an installed CLI binary, run via an agent SDK subprocess, or +// piloted in an interactive TUI inside a cmux surface. +// +// A Backend is exactly a (Provider, RuntimeMode) pair — "claude-agent" is +// anthropic×agent. Backend remains the serialized form (it is written to specs, +// session rows, and the webapp wire format); RuntimeMode is the axis captain +// reasons about, and previously existed only in the frontend's RuntimeModePicker +// while Go smeared it across Backend.Kind, isMode, isSelectorPrefix, +// backendForMode, selectorBackend, and runtimeLogIdentity. +type RuntimeMode string + +const ( + ModeAPI RuntimeMode = "api" + ModeCLI RuntimeMode = "cli" + ModeAgent RuntimeMode = "agent" + ModeCmux RuntimeMode = "cmux" +) + +// AllRuntimeModes lists the modes in canonical order: the order a wildcard +// selector ("*:opus") fans out over, most-capable mechanism first. +// +// The two hand-written lists this replaces disagreed here — AgentsForProvider +// listed the CLI before the agent, wildcardBackends the reverse — and only the +// wildcard order was user-visible. That order wins. +func AllRuntimeModes() []RuntimeMode { + return []RuntimeMode{ModeAPI, ModeAgent, ModeCLI, ModeCmux} +} + +// ParseRuntimeMode normalizes a mode token from the compact grammar. "sdk" is +// accepted as an input alias for agent but is never emitted. +func ParseRuntimeMode(s string) (RuntimeMode, bool) { + switch RuntimeMode(strings.ToLower(strings.TrimSpace(s))) { + case ModeAPI: + return ModeAPI, true + case ModeCLI: + return ModeCLI, true + case ModeAgent, "sdk": + return ModeAgent, true + case ModeCmux: + return ModeCmux, true + default: + return "", false + } +} + +// RuntimeModeList renders the modes as comma-separated text for help/errors. +func RuntimeModeList() string { + parts := make([]string, len(AllRuntimeModes())) + for i, m := range AllRuntimeModes() { + parts[i] = string(m) + } + return strings.Join(parts, ", ") +} diff --git a/pkg/api/registry/model.go b/pkg/api/registry/model.go new file mode 100644 index 0000000..426289d --- /dev/null +++ b/pkg/api/registry/model.go @@ -0,0 +1,276 @@ +package registry + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/collections" +) + +// CodexAutoReviewModel is the internal model used by Codex approval reviewers. +// Its transcripts are implementation noise rather than user sessions. +const CodexAutoReviewModel = "codex-auto-review" + +// Model identifies which LLM serves a request plus the per-request inference +// knobs. Maps onto the legacy ai.Config.Model + ai.Request.{Temperature, +// ReasoningEffort}. +type Model struct { + // Name is the catalog model slug, e.g. "claude-sonnet-4-6"; it drives backend + // inference and pricing lookup. + Name string `json:"model,omitempty" yaml:"model,omitempty" jsonschema:"required" pretty:"label=Model"` + + // ID is the fully-qualified provider id when it differs from Name, e.g. the + // genkit "anthropic/claude-sonnet-4-6" or a codex slug. Empty means use Name. + ID string `json:"id,omitempty" yaml:"id,omitempty" pretty:"label=ID"` + + // Backend overrides backend inference from Name; empty means InferBackend(Name). + Backend Backend `json:"backend,omitempty" yaml:"backend,omitempty" pretty:"label=Backend"` + + // Temperature is the sampling temperature in [0,2]. A pointer so an explicit + // 0.0 is distinguishable from "unset" (fail loud, not a silent default). + Temperature *float64 `json:"temperature,omitempty" yaml:"temperature,omitempty" pretty:"label=Temp"` + + // Effort is the reasoning effort for thinking-capable models. + Effort Effort `json:"effort,omitempty" yaml:"effort,omitempty" jsonschema:"enum=,enum=low,enum=medium,enum=high,enum=xhigh,enum=max,enum=ultra" pretty:"label=Effort"` + + // NoCache disables model response caching for this run. + NoCache bool `json:"noCache,omitempty" yaml:"noCache,omitempty" pretty:"label=No Cache"` + + // Fallbacks are alternative models tried in order when the primary fails with a + // transient/unavailable error or its provider cannot be constructed. Each is a + // full Model (own backend/effort/temperature); a fallback's own nested Fallbacks + // are ignored. Populate it directly, via a comma-separated Name (see ExpandCSV), + // or via the --fallback flag / fallbacks: frontmatter. Each entry may be a + // compact string ("agent:opus:high") or the object form (see ModelList). + Fallbacks ModelList `json:"fallbacks,omitempty" yaml:"fallbacks,omitempty" pretty:"label=Fallbacks"` + + // Mode is the runtime mechanism: api | cli | agent | cmux. It is both an input + // and a derived value — `{model: sonnet, mode: agent}` is the object form of + // the compact "agent:sonnet" — and Backend is exactly (provider, Mode). A Mode + // that contradicts an explicit Backend fails Validate rather than being + // silently reconciled. + Mode RuntimeMode `json:"mode,omitempty" yaml:"mode,omitempty" jsonschema:"enum=api,enum=cli,enum=agent,enum=cmux" pretty:"label=Mode"` + + // The fields below are capabilities of the resolved provider×mode, filled in + // by Resolve. They are outputs: writing them in a spec does not change what + // the adapter can do, so Validate rejects any value that contradicts the + // resolved truth rather than letting a spec claim a capability it lacks. + // The runtime interface assertions remain authoritative at execution time. + + // Streaming reports that the adapter streams incremental events. + Streaming bool `json:"streaming,omitempty" yaml:"streaming,omitempty" jsonschema:"readOnly" pretty:"label=Streaming"` + // MediaTypes are the attachment types this model accepts on this adapter — + // the model's own declared types clamped by the adapter's ceiling. + MediaTypes []string `json:"mediaTypes,omitempty" yaml:"mediaTypes,omitempty" jsonschema:"readOnly" pretty:"label=Media Types"` + // Resume reports that a prior session can be continued by id. + Resume bool `json:"resume,omitempty" yaml:"resume,omitempty" jsonschema:"readOnly" pretty:"label=Resume"` + // Interrupt reports that a running turn can be interrupted. + Interrupt bool `json:"interrupt,omitempty" yaml:"interrupt,omitempty" jsonschema:"readOnly" pretty:"label=Interrupt"` + // Steer reports that a running turn accepts mid-flight steering. + Steer bool `json:"steer,omitempty" yaml:"steer,omitempty" jsonschema:"readOnly" pretty:"label=Steer"` + + // Provider is the descriptor that owns this model. Never serialized: it holds + // the whole catalog, so emitting it would inline the registry into every spec. + // Reach it from a decoded Model with Backend.ModelProvider(). + Provider *Provider `json:"-" yaml:"-"` +} + +// Capabilities fills in Mode, Provider, and the capability flags from the +// resolved backend. Resolve applies it; callers that build a Model by hand can +// call it to get the same enrichment. +func (m Model) Capabilities() Model { + p, mode, ok := ProviderFor(m.Backend) + if !ok { + return m + } + caps, ok := p.Caps(mode) + if !ok { + return m + } + m.Provider = p + m.Mode = mode + m.Streaming = caps.Streaming + m.Resume = caps.Resume + m.Interrupt = caps.Interrupt + m.Steer = caps.Steer + m.MediaTypes = p.MediaTypesFor(mode, m.Name) + return m +} + +// ResolveBackend returns Backend when set, otherwise infers it from Name. +func (m Model) ResolveBackend() (Backend, error) { + if m.Backend != "" { + if !m.Backend.Valid() { + return "", fmt.Errorf("invalid backend %q (valid: %s)", m.Backend, BackendList()) + } + return m.Backend, nil + } + return InferBackend(m.Name) +} + +// Temp returns the temperature and whether it was explicitly set (non-nil), so +// providers can distinguish an intentional 0.0 from "unset, use the default". +func (m Model) Temp() (float64, bool) { + if m.Temperature == nil { + return 0, false + } + return *m.Temperature, true +} + +// Validate checks the model name is present, the knobs are in range, and every +// fallback is itself a well-formed model. +func (m Model) Validate() error { + if m.Name == "" { + return fmt.Errorf("model name is required") + } + if err := m.validateKnobs(); err != nil { + return err + } + for i, fb := range m.Fallbacks { + if fb.Name == "" { + return fmt.Errorf("fallback[%d]: model name is required", i) + } + if err := fb.validateKnobs(); err != nil { + return fmt.Errorf("fallback[%d] %q: %w", i, fb.Name, err) + } + } + return nil +} + +// validateKnobs range-checks the per-request inference knobs shared by a primary +// model and each fallback (backend/mode/temperature/effort), independent of Name. +func (m Model) validateKnobs() error { + if m.Backend != "" && !m.Backend.Valid() { + return fmt.Errorf("invalid backend %q (valid: %s)", m.Backend, BackendList()) + } + if err := m.validateMode(); err != nil { + return err + } + if m.Temperature != nil && (*m.Temperature < 0 || *m.Temperature > 2) { + return fmt.Errorf("invalid temperature %v (valid: 0.0-2.0)", *m.Temperature) + } + return m.Effort.Validate() +} + +// validateMode rejects a mode that is unknown, or that contradicts an explicit +// backend. Backend is exactly (provider, mode), so `backend: anthropic` with +// `mode: agent` names two different runtimes — reconciling it silently would +// pick one of them behind the user's back. +func (m Model) validateMode() error { + if m.Mode == "" { + return nil + } + if _, ok := ParseRuntimeMode(string(m.Mode)); !ok { + return fmt.Errorf("invalid mode %q (valid: %s)", m.Mode, RuntimeModeList()) + } + if m.Backend == "" { + return nil + } + if actual := m.Backend.Mode(); actual != m.Mode { + return fmt.Errorf("mode %q contradicts backend %q, which is mode %q; set one or the other", + m.Mode, m.Backend, actual) + } + return nil +} + +// Merge overlays o's set fields onto m, returning the result. It backs Spec +// merging in pkg/api, where a later layer overrides an earlier one field by field. +func (m Model) Merge(o Model) Model { + if o.Name != "" { + m.Name = o.Name + } + if o.ID != "" { + m.ID = o.ID + } + if o.Backend != "" { + m.Backend = o.Backend + } + if o.Temperature != nil { + m.Temperature = o.Temperature + } + if o.Effort != "" { + m.Effort = o.Effort + } + if o.Mode != "" { + m.Mode = o.Mode + } + if o.NoCache { + m.NoCache = true + } + if len(o.Fallbacks) > 0 { + m.Fallbacks = o.Fallbacks + } + return m +} + +// ExpandCSV moves any comma-separated tail of Name into name-only Fallbacks +// (prepended, so CSV order is preserved ahead of an explicit Fallbacks list), +// leaving Name as the single primary. It is idempotent: a single-model Name is +// returned with only whitespace trimmed, and a second call is a no-op. +func (m Model) ExpandCSV() Model { + names := splitCSV(m.Name) + if len(names) == 0 { + return m + } + m.Name = names[0] + if len(names) == 1 { + return m + } + tail := make(ModelList, 0, len(names)-1) + for _, n := range names[1:] { + tail = append(tail, Model{Name: n}) + } + m.Fallbacks = append(tail, m.Fallbacks...) + return m +} + +// Candidates returns the ordered models to try: the ExpandCSV primary first, then +// each fallback. Fallbacks inherit the primary's Temperature/NoCache when unset, +// and Effort only when they belong to the same provider family. They keep their +// own Name/Backend (an empty Backend is inferred at construction from the fallback's +// own Name), clear ID, and have nested Fallbacks dropped. A length of 1 means "no fallback". +func (m Model) Candidates() []Model { + m = m.ExpandCSV() + primary := m + primary.Fallbacks = nil + out := make([]Model, 0, collections.SafeAdd(1, len(m.Fallbacks))) + out = append(out, primary) + for _, fb := range m.Fallbacks { + fb.Fallbacks = nil + fb.ID = "" + if fb.Temperature == nil { + fb.Temperature = m.Temperature + } + if fb.Effort == "" && modelProvider(fb) == modelProvider(m) { + fb.Effort = m.Effort + } + if !fb.NoCache { + fb.NoCache = m.NoCache + } + out = append(out, fb) + } + return out +} + +func modelProvider(model Model) Backend { + if provider := model.Backend.Provider(); provider != "" { + return provider + } + backend, err := InferBackend(model.Name) + if err != nil { + return "" + } + return backend.Provider() +} + +// splitCSV splits a comma-separated string into trimmed, non-empty parts. +func splitCSV(s string) []string { + parts := strings.Split(s, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + if t := strings.TrimSpace(p); t != "" { + out = append(out, t) + } + } + return out +} diff --git a/pkg/api/registry/model_compact.go b/pkg/api/registry/model_compact.go new file mode 100644 index 0000000..1ec28da --- /dev/null +++ b/pkg/api/registry/model_compact.go @@ -0,0 +1,191 @@ +package registry + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" + + "gopkg.in/yaml.v3" +) + +// Compact model syntax +// +// A Model can be written as a compact string instead of the object form, which +// is handy for fallback chains and terse config: +// +// opus → {model: opus} +// opus:high → {model: opus, effort: high} +// agent:opus:high → {model: opus, effort: high, backend: claude-agent} +// opus:high, sonnet:medium → primary opus:high with a sonnet:medium fallback +// +// The grammar per comma-separated element is `[mode:]model[:effort]`, where mode +// is a mechanism keyword (api | cli | agent | cmux; sdk is an alias for agent) +// combined with the model's inferred family to a concrete backend. A two-segment +// element is disambiguated by content: a known effort tail is model:effort, a +// known mode head is mode:model. + +// backendForPrefix resolves an element's prefix plus its model name to a concrete +// backend: the provider that claims the model, combined with the mechanism the +// prefix names. +func backendForPrefix(prefix, modelName string) (Backend, error) { + p, _, _, ok := ProviderForToken(modelName) + if !ok { + return "", fmt.Errorf("mode %q: %w: %s (pass an explicit backend: %s)", prefix, ErrUnknownModel, modelName, BackendList()) + } + mode, err := resolveMode(prefix, ModeAPI, p, modelName) + if err != nil { + return "", err + } + return p.BackendFor(mode) +} + +// parseCompactElement parses one `[prefix:]model[:effort]` element. +// +// It resolves the BACKEND but deliberately leaves Name as written: a decoded +// spec keeps the model the user asked for ("opus"), and only the later resolve +// step (ResolveModel) maps it onto an exact catalog id. The spec is a request, +// not a resolution, so decoding must not bake today's catalog snapshot into it. +func parseCompactElement(s string) (Model, error) { + s = strings.TrimSpace(s) + if s == "" { + return Model{}, fmt.Errorf("empty model") + } + prefix, name, effort, err := splitElement(s, "") + if err != nil { + return Model{}, err + } + if name == "" { + return Model{}, fmt.Errorf("model name required in %q", s) + } + m := Model{Name: name, Effort: effort} + if prefix == "" { + return m, nil + } + if prefix == "*" { + return Model{}, fmt.Errorf("wildcard selector %q is only valid for --multi-models", s) + } + backend, err := backendForPrefix(prefix, name) + if err != nil { + return Model{}, err + } + m.Backend = backend + return m, nil +} + +// parseCompactModel parses a full compact string: comma-separated elements where +// the first is the primary and the rest become its Fallbacks. +func parseCompactModel(s string) (Model, error) { + parts := splitCSV(s) + if len(parts) == 0 { + return Model{}, fmt.Errorf("empty model %q", s) + } + primary, err := parseCompactElement(parts[0]) + if err != nil { + return Model{}, err + } + for _, part := range parts[1:] { + fb, err := parseCompactElement(part) + if err != nil { + return Model{}, err + } + primary.Fallbacks = append(primary.Fallbacks, fb) + } + return primary, nil +} + +// Expand parses a compact Name (`[mode:]model[:effort]`, optionally with a +// comma-separated fallback tail) into concrete Name/Effort/Backend + Fallbacks, +// preserving any fields already set on the receiver that the compact form does +// not specify. It is idempotent: a plain model name is returned unchanged. +// Errors on a malformed compact string (unknown mode/effort, ambiguous form). +func (m Model) Expand() (Model, error) { + if !strings.ContainsAny(m.Name, ":,") { + m.Name = strings.TrimSpace(m.Name) + return m, nil + } + parsed, err := parseCompactModel(m.Name) + if err != nil { + return Model{}, err + } + if parsed.Effort == "" { + parsed.Effort = m.Effort + } + if parsed.Backend == "" { + parsed.Backend = m.Backend + } + parsed.ID = m.ID + parsed.Temperature = m.Temperature + if !parsed.NoCache { + parsed.NoCache = m.NoCache + } + parsed.Fallbacks = append(parsed.Fallbacks, m.Fallbacks...) + return parsed, nil +} + +// ModelList is the type of Model.Fallbacks. Each entry may be written as a +// compact string ("agent:opus:high") or the object form. It is a named slice +// rather than a method on Model itself because Model is inline-embedded in Spec: +// a value-level Unmarshaler on Model would hijack the whole Spec object and break +// field promotion. A Spec-level `model:` scalar lands in Model.Name (a string) — +// Model.Expand parses that. +type ModelList []Model + +func (l *ModelList) UnmarshalJSON(data []byte) error { + trimmed := bytes.TrimSpace(data) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) { + *l = nil + return nil + } + var raw []json.RawMessage + if err := json.Unmarshal(trimmed, &raw); err != nil { + return err + } + out := make(ModelList, 0, len(raw)) + for _, elem := range raw { + e := bytes.TrimSpace(elem) + if len(e) > 0 && e[0] == '"' { + var s string + if err := json.Unmarshal(e, &s); err != nil { + return err + } + m, err := parseCompactElement(s) + if err != nil { + return err + } + out = append(out, m) + continue + } + var m Model + if err := json.Unmarshal(e, &m); err != nil { + return err + } + out = append(out, m) + } + *l = out + return nil +} + +func (l *ModelList) UnmarshalYAML(value *yaml.Node) error { + if value.Kind != yaml.SequenceNode { + return fmt.Errorf("fallbacks must be a list, got %v", value.Kind) + } + out := make(ModelList, 0, len(value.Content)) + for _, node := range value.Content { + if node.Kind == yaml.ScalarNode && node.Tag != "!!null" { + m, err := parseCompactElement(node.Value) + if err != nil { + return err + } + out = append(out, m) + continue + } + var m Model + if err := node.Decode(&m); err != nil { + return err + } + out = append(out, m) + } + *l = out + return nil +} diff --git a/pkg/api/registry/model_compact_test.go b/pkg/api/registry/model_compact_test.go new file mode 100644 index 0000000..3966cf3 --- /dev/null +++ b/pkg/api/registry/model_compact_test.go @@ -0,0 +1,130 @@ +package registry + +import ( + "encoding/json" + "reflect" + "strings" + "testing" + + "gopkg.in/yaml.v3" +) + +func TestParseCompactElement(t *testing.T) { + cases := []struct { + in string + want Model + wantErr string + }{ + {in: "opus", want: Model{Name: "opus"}}, + {in: "claude-sonnet-4-6", want: Model{Name: "claude-sonnet-4-6"}}, + {in: "opus:high", want: Model{Name: "opus", Effort: EffortHigh}}, + {in: "sonnet:medium", want: Model{Name: "sonnet", Effort: EffortMedium}}, + {in: "agent:opus:high", want: Model{Name: "opus", Effort: EffortHigh, Backend: BackendClaudeAgent}}, + {in: "sdk:opus:high", want: Model{Name: "opus", Effort: EffortHigh, Backend: BackendClaudeAgent}}, // sdk = agent + {in: "cmux:opus", want: Model{Name: "opus", Backend: BackendClaudeCmux}}, + {in: "api:opus", want: Model{Name: "opus", Backend: BackendAnthropic}}, + {in: "cli:opus", want: Model{Name: "opus", Backend: BackendClaudeCLI}}, + {in: "cmux:codex:medium", want: Model{Name: "codex", Effort: EffortMedium, Backend: BackendCodexCmux}}, + {in: "api:gpt-5.5", want: Model{Name: "gpt-5.5", Backend: BackendOpenAI}}, + {in: " opus : high ", want: Model{Name: "opus", Effort: EffortHigh}}, + {in: "cmux:gemini-2.0", wantErr: "not supported for gemini"}, // no gemini cmux + {in: "bogusmode:opus:high", wantErr: "invalid mode"}, + {in: "opus:notaneffort", wantErr: "ambiguous"}, + {in: "a:b:c:d", wantErr: "too many"}, + {in: "", wantErr: "empty"}, + } + for _, tc := range cases { + t.Run(tc.in, func(t *testing.T) { + got, err := parseCompactElement(tc.in) + if tc.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("err = %v, want mention of %q", err, tc.wantErr) + } + return + } + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("got %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestModel_Expand(t *testing.T) { + t.Run("plain name unchanged", func(t *testing.T) { + got, err := Model{Name: "opus", Effort: EffortHigh}.Expand() + if err != nil || got.Name != "opus" || got.Effort != EffortHigh { + t.Fatalf("got %+v err %v", got, err) + } + }) + t.Run("compact name parsed", func(t *testing.T) { + got, err := Model{Name: "agent:opus:high"}.Expand() + if err != nil { + t.Fatal(err) + } + if got.Name != "opus" || got.Effort != EffortHigh || got.Backend != BackendClaudeAgent { + t.Errorf("got %+v", got) + } + }) + t.Run("csv tail becomes fallbacks", func(t *testing.T) { + got, err := Model{Name: "opus:high, sonnet:medium, api:gpt-5.5"}.Expand() + if err != nil { + t.Fatal(err) + } + if got.Name != "opus" || got.Effort != EffortHigh { + t.Errorf("primary = %+v", got) + } + if len(got.Fallbacks) != 2 { + t.Fatalf("fallbacks = %+v", got.Fallbacks) + } + if got.Fallbacks[0].Name != "sonnet" || got.Fallbacks[0].Effort != EffortMedium { + t.Errorf("fb0 = %+v", got.Fallbacks[0]) + } + if got.Fallbacks[1].Name != "gpt-5.5" || got.Fallbacks[1].Backend != BackendOpenAI { + t.Errorf("fb1 = %+v", got.Fallbacks[1]) + } + }) + t.Run("preserves explicit fields not in compact", func(t *testing.T) { + got, err := Model{Name: "opus:high", Backend: BackendClaudeCmux}.Expand() + if err != nil { + t.Fatal(err) + } + if got.Backend != BackendClaudeCmux { + t.Errorf("backend = %q, want preserved claude-cmux", got.Backend) + } + }) + t.Run("bad compact errors", func(t *testing.T) { + if _, err := (Model{Name: "opus:nope"}).Expand(); err == nil { + t.Fatal("expected error") + } + }) +} + +func TestModelList_Unmarshal(t *testing.T) { + t.Run("json mixed string and object", func(t *testing.T) { + var l ModelList + if err := json.Unmarshal([]byte(`["agent:opus:high", {"model":"sonnet","effort":"medium"}]`), &l); err != nil { + t.Fatal(err) + } + if len(l) != 2 { + t.Fatalf("len = %d", len(l)) + } + if l[0].Name != "opus" || l[0].Backend != BackendClaudeAgent || l[0].Effort != EffortHigh { + t.Errorf("l0 = %+v", l[0]) + } + if l[1].Name != "sonnet" || l[1].Effort != EffortMedium { + t.Errorf("l1 = %+v", l[1]) + } + }) + t.Run("yaml mixed", func(t *testing.T) { + var l ModelList + if err := yaml.Unmarshal([]byte("- api:opus\n- model: sonnet\n effort: high\n"), &l); err != nil { + t.Fatal(err) + } + if len(l) != 2 || l[0].Name != "opus" || l[0].Backend != BackendAnthropic || l[1].Name != "sonnet" { + t.Errorf("got %+v", l) + } + }) +} diff --git a/pkg/api/model_test.go b/pkg/api/registry/model_test.go similarity index 83% rename from pkg/api/model_test.go rename to pkg/api/registry/model_test.go index ca0f1a6..ced4108 100644 --- a/pkg/api/model_test.go +++ b/pkg/api/registry/model_test.go @@ -1,4 +1,4 @@ -package api +package registry import ( "encoding/json" @@ -9,6 +9,8 @@ import ( "gopkg.in/yaml.v3" ) +func floatPtr(f float64) *float64 { return &f } + func modelNames(models []Model) []string { if len(models) == 0 { return nil @@ -61,21 +63,23 @@ func TestModel_Candidates(t *testing.T) { Temperature: floatPtr(0.3), NoCache: true, Fallbacks: []Model{ - {Name: "gpt-4o"}, // inherits primary effort/temp/noCache + {Name: "gpt-4o"}, + {Name: "claude-haiku-4-5"}, {Name: "gemini-2.0-flash", Effort: EffortLow, ID: "drop-me", Fallbacks: []Model{{Name: "nested"}}}, }, } got := primary.Candidates() - if names := modelNames(got); !reflect.DeepEqual(names, []string{"claude-sonnet-5", "gpt-4o", "gemini-2.0-flash"}) { + if names := modelNames(got); !reflect.DeepEqual(names, []string{"claude-sonnet-5", "gpt-4o", "claude-haiku-4-5", "gemini-2.0-flash"}) { t.Fatalf("candidate order = %v", names) } if got[0].Fallbacks != nil { t.Errorf("primary candidate should not carry Fallbacks, got %v", got[0].Fallbacks) } - // gpt-4o inherits the primary's effort/temperature/noCache. - if got[1].Effort != EffortHigh { - t.Errorf("fallback effort = %q, want inherited %q", got[1].Effort, EffortHigh) + // A cross-provider fallback keeps an empty effort so its provider default can + // apply independently, while transport-neutral knobs still inherit. + if got[1].Effort != EffortNone { + t.Errorf("cross-provider fallback effort = %q, want provider default", got[1].Effort) } if got[1].Temperature == nil || *got[1].Temperature != 0.3 { t.Errorf("fallback temperature = %v, want inherited 0.3", got[1].Temperature) @@ -83,15 +87,18 @@ func TestModel_Candidates(t *testing.T) { if !got[1].NoCache { t.Errorf("fallback NoCache = false, want inherited true") } + if got[2].Effort != EffortHigh { + t.Errorf("same-provider fallback effort = %q, want inherited %q", got[2].Effort, EffortHigh) + } // gemini keeps its own effort, drops ID and nested fallbacks. - if got[2].Effort != EffortLow { - t.Errorf("fallback own effort = %q, want %q", got[2].Effort, EffortLow) + if got[3].Effort != EffortLow { + t.Errorf("fallback own effort = %q, want %q", got[3].Effort, EffortLow) } - if got[2].ID != "" { - t.Errorf("fallback ID = %q, want cleared", got[2].ID) + if got[3].ID != "" { + t.Errorf("fallback ID = %q, want cleared", got[3].ID) } - if got[2].Fallbacks != nil { - t.Errorf("nested fallbacks should be dropped, got %v", got[2].Fallbacks) + if got[3].Fallbacks != nil { + t.Errorf("nested fallbacks should be dropped, got %v", got[3].Fallbacks) } } diff --git a/pkg/api/registry/models.json b/pkg/api/registry/models.json new file mode 100644 index 0000000..f7675cd --- /dev/null +++ b/pkg/api/registry/models.json @@ -0,0 +1,1080 @@ +[ + { + "id": "claude-opus-5", + "provider": "anthropic", + "family": "opus", + "version": "5", + "label": "Claude Opus 5", + "releaseDate": "2026-07-24", + "reasoning": true, + "contextWindow": 1000000, + "cost": { + "input": 5, + "output": 25, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "adaptiveThinking": true, + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + { + "id": "claude-sonnet-5", + "provider": "anthropic", + "family": "sonnet", + "version": "5", + "label": "Claude Sonnet 5", + "releaseDate": "2026-06-29", + "reasoning": true, + "contextWindow": 1000000, + "cost": { + "input": 2, + "output": 10, + "cacheRead": 0.2, + "cacheWrite": 2.5 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "adaptiveThinking": true, + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + { + "id": "claude-fable-5", + "provider": "anthropic", + "family": "fable", + "version": "5", + "label": "Claude Fable 5", + "releaseDate": "2026-06-07", + "reasoning": true, + "contextWindow": 1000000, + "cost": { + "input": 10, + "output": 50, + "cacheRead": 1, + "cacheWrite": 12.5 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "adaptiveThinking": true, + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + { + "id": "claude-opus-4-8", + "provider": "anthropic", + "family": "opus", + "version": "4.8", + "label": "Claude Opus 4.8", + "releaseDate": "2026-05-28", + "reasoning": true, + "contextWindow": 1000000, + "cost": { + "input": 5, + "output": 25, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "adaptiveThinking": true, + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + { + "id": "claude-opus-4-7", + "provider": "anthropic", + "family": "opus", + "version": "4.7", + "label": "Claude Opus 4.7", + "releaseDate": "2026-04-14", + "reasoning": true, + "contextWindow": 1000000, + "cost": { + "input": 5, + "output": 25, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "adaptiveThinking": true, + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + { + "id": "claude-sonnet-4-6", + "provider": "anthropic", + "family": "sonnet", + "version": "4.6", + "label": "Claude Sonnet 4.6", + "releaseDate": "2026-02-17", + "reasoning": true, + "temperature": true, + "contextWindow": 1000000, + "cost": { + "input": 3, + "output": 15, + "cacheRead": 0.3, + "cacheWrite": 3.75 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ] + }, + { + "id": "claude-opus-4-6", + "provider": "anthropic", + "family": "opus", + "version": "4.6", + "label": "Claude Opus 4.6", + "releaseDate": "2026-02-04", + "reasoning": true, + "temperature": true, + "contextWindow": 1000000, + "cost": { + "input": 5, + "output": 25, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ] + }, + { + "id": "claude-opus-4-5", + "provider": "anthropic", + "family": "opus", + "version": "4.5", + "label": "Claude Opus 4.5 (latest)", + "releaseDate": "2025-11-24", + "reasoning": true, + "temperature": true, + "contextWindow": 200000, + "cost": { + "input": 5, + "output": 25, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ] + }, + { + "id": "claude-opus-4-5-20251101", + "provider": "anthropic", + "family": "opus", + "version": "4.5-20251101", + "label": "Claude Opus 4.5", + "releaseDate": "2025-11-24", + "reasoning": true, + "temperature": true, + "contextWindow": 200000, + "cost": { + "input": 5, + "output": 25, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ] + }, + { + "id": "claude-haiku-4-5", + "provider": "anthropic", + "family": "haiku", + "version": "4.5", + "label": "Claude Haiku 4.5 (latest)", + "releaseDate": "2025-10-15", + "reasoning": true, + "temperature": true, + "contextWindow": 200000, + "cost": { + "input": 1, + "output": 5, + "cacheRead": 0.1, + "cacheWrite": 1.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true + }, + { + "id": "claude-haiku-4-5-20251001", + "provider": "anthropic", + "family": "haiku", + "version": "4.5-20251001", + "label": "Claude Haiku 4.5", + "releaseDate": "2025-10-15", + "reasoning": true, + "temperature": true, + "contextWindow": 200000, + "cost": { + "input": 1, + "output": 5, + "cacheRead": 0.1, + "cacheWrite": 1.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ] + }, + { + "id": "claude-sonnet-4-5", + "provider": "anthropic", + "family": "sonnet", + "version": "4.5", + "label": "Claude Sonnet 4.5 (latest)", + "releaseDate": "2025-09-29", + "reasoning": true, + "temperature": true, + "contextWindow": 1000000, + "cost": { + "input": 3, + "output": 15, + "cacheRead": 0.3, + "cacheWrite": 3.75 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "supersededBy": "claude-sonnet-4-6" + }, + { + "id": "claude-sonnet-4-5-20250929", + "provider": "anthropic", + "family": "sonnet", + "version": "4.5-20250929", + "label": "Claude Sonnet 4.5", + "releaseDate": "2025-09-29", + "reasoning": true, + "temperature": true, + "contextWindow": 1000000, + "cost": { + "input": 3, + "output": 15, + "cacheRead": 0.3, + "cacheWrite": 3.75 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "supersededBy": "claude-sonnet-4-6" + }, + { + "id": "gpt-5.6", + "provider": "openai", + "family": "gpt", + "version": "5.6", + "label": "GPT-5.6", + "releaseDate": "2026-07-09", + "reasoning": true, + "contextWindow": 1050000, + "cost": { + "input": 5, + "output": 30, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "availability": [ + "api" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + { + "id": "gpt-5.6-luna", + "provider": "openai", + "family": "gpt", + "version": "5.6-luna", + "label": "GPT-5.6 Luna", + "releaseDate": "2026-07-09", + "reasoning": true, + "contextWindow": 1050000, + "cost": { + "input": 1, + "output": 6, + "cacheRead": 0.1, + "cacheWrite": 1.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "availability": [ + "api", + "codex" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max" + ], + "priority": 3, + "aliases": [ + "luna" + ] + }, + { + "id": "gpt-5.6-sol", + "provider": "openai", + "family": "gpt", + "version": "5.6-sol", + "label": "GPT-5.6 Sol", + "releaseDate": "2026-07-09", + "reasoning": true, + "contextWindow": 1050000, + "cost": { + "input": 5, + "output": 30, + "cacheRead": 0.5, + "cacheWrite": 6.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "availability": [ + "api", + "codex" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max", + "ultra" + ], + "priority": 1, + "aliases": [ + "sol" + ] + }, + { + "id": "gpt-5.6-terra", + "provider": "openai", + "family": "gpt", + "version": "5.6-terra", + "label": "GPT-5.6 Terra", + "releaseDate": "2026-07-09", + "reasoning": true, + "contextWindow": 1050000, + "cost": { + "input": 2.5, + "output": 15, + "cacheRead": 0.25, + "cacheWrite": 3.125 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "preferred": true, + "availability": [ + "api", + "codex" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh", + "max", + "ultra" + ], + "priority": 2, + "aliases": [ + "terra" + ] + }, + { + "id": "gpt-5.5", + "provider": "openai", + "family": "gpt", + "version": "5.5", + "label": "GPT-5.5", + "releaseDate": "2026-04-23", + "reasoning": true, + "contextWindow": 1050000, + "cost": { + "input": 5, + "output": 30, + "cacheRead": 0.5 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh" + ] + }, + { + "id": "gpt-5.4", + "provider": "openai", + "family": "gpt", + "version": "5.4", + "label": "GPT-5.4", + "releaseDate": "2026-03-05", + "reasoning": true, + "contextWindow": 1050000, + "cost": { + "input": 2.5, + "output": 15, + "cacheRead": 0.25 + }, + "inputMediaTypes": [ + "image/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh" + ] + }, + { + "id": "gpt-5.2", + "provider": "openai", + "family": "gpt", + "version": "5.2", + "label": "GPT-5.2", + "releaseDate": "2025-12-11", + "reasoning": true, + "contextWindow": 400000, + "cost": { + "input": 1.75, + "output": 14, + "cacheRead": 0.175 + }, + "inputMediaTypes": [ + "image/*" + ], + "supportedEfforts": [ + "low", + "medium", + "high", + "xhigh" + ] + }, + { + "id": "gpt-5.1", + "provider": "openai", + "family": "gpt", + "version": "5.1", + "label": "GPT-5.1", + "releaseDate": "2025-11-13", + "reasoning": true, + "contextWindow": 400000, + "cost": { + "input": 1.25, + "output": 10, + "cacheRead": 0.125 + }, + "inputMediaTypes": [ + "image/*" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gpt-5", + "provider": "openai", + "family": "gpt", + "version": "5", + "label": "GPT-5", + "releaseDate": "2025-08-07", + "reasoning": true, + "contextWindow": 400000, + "cost": { + "input": 1.25, + "output": 10, + "cacheRead": 0.125 + }, + "inputMediaTypes": [ + "image/*" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3.5-flash-lite", + "provider": "google", + "family": "gemini", + "version": "3.5-flash-lite", + "label": "Gemini 3.5 Flash Lite", + "releaseDate": "2026-07-21", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.3, + "output": 2.5, + "cacheRead": 0.03 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3.6-flash", + "provider": "google", + "family": "gemini", + "version": "3.6-flash", + "label": "Gemini 3.6 Flash", + "releaseDate": "2026-07-21", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 1.5, + "output": 7.5, + "cacheRead": 0.15 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "preferred": true, + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3.5-flash", + "provider": "google", + "family": "gemini", + "version": "3.5-flash", + "label": "Gemini 3.5 Flash", + "releaseDate": "2026-05-19", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 1.5, + "output": 9, + "cacheRead": 0.15 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "preferred": true, + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-flash-latest", + "provider": "google", + "family": "gemini", + "version": "flash-latest", + "label": "Gemini Flash Latest", + "releaseDate": "2026-05-19", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 1.5, + "output": 9, + "cacheRead": 0.15 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3.1-flash-lite", + "provider": "google", + "family": "gemini", + "version": "3.1-flash-lite", + "label": "Gemini 3.1 Flash Lite", + "releaseDate": "2026-05-07", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.25, + "output": 1.5, + "cacheRead": 0.025 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-flash-lite-latest", + "provider": "google", + "family": "gemini", + "version": "flash-lite-latest", + "label": "Gemini Flash-Lite Latest", + "releaseDate": "2026-05-07", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.25, + "output": 1.5, + "cacheRead": 0.025 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-robotics-er-1.6-preview", + "provider": "google", + "family": "gemini", + "version": "robotics-er-1.6-preview", + "label": "Gemini Robotics-ER 1.6 Preview", + "releaseDate": "2026-04-14", + "reasoning": true, + "temperature": true, + "contextWindow": 131072, + "cost": { + "input": 1, + "output": 5 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*" + ] + }, + { + "id": "gemini-3.1-flash-live-preview", + "provider": "google", + "family": "gemini", + "version": "3.1-flash-live-preview", + "label": "Gemini 3.1 Flash Live Preview", + "releaseDate": "2026-03-26", + "reasoning": true, + "temperature": true, + "contextWindow": 131072, + "cost": { + "input": 0.75, + "output": 4.5 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3.1-flash-lite-preview", + "provider": "google", + "family": "gemini", + "version": "3.1-flash-lite-preview", + "label": "Gemini 3.1 Flash Lite Preview", + "releaseDate": "2026-03-03", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.25, + "output": 1.5, + "cacheRead": 0.025 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3.1-pro-preview", + "provider": "google", + "family": "gemini", + "version": "3.1-pro-preview", + "label": "Gemini 3.1 Pro Preview", + "releaseDate": "2026-02-19", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 2, + "output": 12, + "cacheRead": 0.2 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3.1-pro-preview-customtools", + "provider": "google", + "family": "gemini", + "version": "3.1-pro-preview-customtools", + "label": "Gemini 3.1 Pro Preview Custom Tools", + "releaseDate": "2026-02-19", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 2, + "output": 12, + "cacheRead": 0.2 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3-flash-preview", + "provider": "google", + "family": "gemini", + "version": "3-flash-preview", + "label": "Gemini 3 Flash Preview", + "releaseDate": "2025-12-17", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.5, + "output": 3, + "cacheRead": 0.05 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "medium", + "high" + ] + }, + { + "id": "gemini-3-pro-preview", + "provider": "google", + "family": "gemini", + "version": "3-pro-preview", + "label": "Gemini 3 Pro Preview", + "releaseDate": "2025-11-18", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 2, + "output": 12, + "cacheRead": 0.2 + }, + "inputMediaTypes": [ + "image/*", + "video/*", + "audio/*", + "application/pdf" + ], + "supportedEfforts": [ + "low", + "high" + ] + }, + { + "id": "gemini-2.5-computer-use-preview-10-2025", + "provider": "google", + "family": "gemini", + "version": "2.5-computer-use-preview-10-2025", + "label": "Gemini 2.5 Computer Use Preview 10-2025", + "releaseDate": "2025-10-07", + "reasoning": true, + "temperature": true, + "contextWindow": 131072, + "cost": { + "input": 1.25, + "output": 10 + }, + "inputMediaTypes": [ + "image/*" + ] + }, + { + "id": "gemini-2.5-flash", + "provider": "google", + "family": "gemini", + "version": "2.5-flash", + "label": "Gemini 2.5 Flash", + "releaseDate": "2025-06-17", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.3, + "output": 2.5, + "cacheRead": 0.03 + }, + "inputMediaTypes": [ + "image/*", + "audio/*", + "video/*", + "application/pdf" + ] + }, + { + "id": "gemini-2.5-flash-lite", + "provider": "google", + "family": "gemini", + "version": "2.5-flash-lite", + "label": "Gemini 2.5 Flash-Lite", + "releaseDate": "2025-06-17", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.1, + "output": 0.4, + "cacheRead": 0.01 + }, + "inputMediaTypes": [ + "image/*", + "audio/*", + "video/*", + "application/pdf" + ] + }, + { + "id": "gemini-2.5-pro", + "provider": "google", + "family": "gemini", + "version": "2.5-pro", + "label": "Gemini 2.5 Pro", + "releaseDate": "2025-06-17", + "reasoning": true, + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 1.25, + "output": 10, + "cacheRead": 0.125 + }, + "inputMediaTypes": [ + "image/*", + "audio/*", + "video/*", + "application/pdf" + ] + }, + { + "id": "gemini-2.0-flash", + "provider": "google", + "family": "gemini", + "version": "2.0-flash", + "label": "Gemini 2.0 Flash", + "releaseDate": "2024-12-11", + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.1, + "output": 0.4, + "cacheRead": 0.025 + }, + "inputMediaTypes": [ + "image/*", + "audio/*", + "video/*", + "application/pdf" + ] + }, + { + "id": "gemini-2.0-flash-lite", + "provider": "google", + "family": "gemini", + "version": "2.0-flash-lite", + "label": "Gemini 2.0 Flash-Lite", + "releaseDate": "2024-12-11", + "temperature": true, + "contextWindow": 1048576, + "cost": { + "input": 0.075, + "output": 0.3 + }, + "inputMediaTypes": [ + "image/*", + "audio/*", + "video/*", + "application/pdf" + ] + }, + { + "id": "deepseek-v4-flash", + "provider": "deepseek", + "family": "deepseek", + "version": "v4-flash", + "label": "DeepSeek V4 Flash", + "releaseDate": "2026-04-24", + "reasoning": true, + "temperature": true, + "contextWindow": 1000000, + "cost": { + "input": 0.14, + "output": 0.28, + "cacheRead": 0.0028 + }, + "preferred": true + }, + { + "id": "deepseek-v4-pro", + "provider": "deepseek", + "family": "deepseek", + "version": "v4-pro", + "label": "DeepSeek V4 Pro", + "releaseDate": "2026-04-24", + "reasoning": true, + "temperature": true, + "contextWindow": 1000000, + "cost": { + "input": 0.435, + "output": 0.87, + "cacheRead": 0.003625 + }, + "preferred": true + } +] diff --git a/pkg/api/registry/parse.go b/pkg/api/registry/parse.go new file mode 100644 index 0000000..d88395e --- /dev/null +++ b/pkg/api/registry/parse.go @@ -0,0 +1,369 @@ +package registry + +import ( + "errors" + "fmt" + "strings" +) + +// ErrUnknownModel marks "no provider claims this model name". Callers enrich it +// with "did you mean" suggestions, so it stays a wrapped sentinel. +var ErrUnknownModel = ErrInferBackend + +// The model grammar +// +// sonnet → {model: claude-sonnet-5, backend: anthropic} +// sonnet:high → + effort high +// agent:sonnet:high → + backend claude-agent +// claude-cmux:opus → an explicit backend name works as a prefix too +// *:fable → every backend of the claiming family (multi only) +// opus:high, sonnet:medium → primary opus:high with a sonnet:medium fallback +// +// One element is `[prefix:]model[:effort]`, where prefix is a runtime mode +// (api | cli | agent | cmux; sdk aliases agent), a concrete backend name, or "*". +// A two-segment element is disambiguated by content: a known effort tail is +// model:effort, a known prefix head is prefix:model. +// +// This is the only implementation of the grammar. It previously existed twice — +// parseCompactElement in pkg/api and resolveSelectorPart in pkg/ai — over two +// different claim tables, so `--model agent:sol` and `model: agent:sol` in +// frontmatter disagreed about whether the model existed. + +// ParseOptions tunes one parse. +type ParseOptions struct { + // Backend forces a backend. A prefix that resolves elsewhere is an error. + Backend Backend + // Mode forces the runtime mechanism when the element carries no prefix of its + // own — the object form of "agent:sonnet". An explicit prefix still wins. + Mode RuntimeMode + // BaseName supplies the model for a bare prefix ("cli:" with --model sonnet). + BaseName string + // AllowWildcard permits "*" fan-out. Only --multi-models sets it. + AllowWildcard bool +} + +// ParseModel parses one compact element into a concrete model. A wildcard, or +// anything that resolves to more than one model, is an error here — use +// ParseModelMulti. +func ParseModel(s string) (Model, error) { + models, err := ParseModelElement(s, ParseOptions{}) + if err != nil { + return Model{}, err + } + if len(models) != 1 { + return Model{}, fmt.Errorf("wildcard selector %q is only valid for --multi-models", s) + } + return models[0], nil +} + +// ParseModelMulti expands one element, fanning a "*" prefix out across every +// backend of the claiming family that actually offers the model. +func ParseModelMulti(s string, opts ParseOptions) ([]Model, error) { + opts.AllowWildcard = true + return ParseModelElement(s, opts) +} + +// ParseModelElement parses one `[prefix:]model[:effort]` element. +func ParseModelElement(raw string, opts ParseOptions) ([]Model, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, fmt.Errorf("empty model") + } + + prefix, name, effort, err := splitElement(raw, opts.BaseName) + if err != nil { + return nil, err + } + if name == "" { + return nil, fmt.Errorf("runtime selector %q needs a model, or --model must provide a base model", raw) + } + + if prefix == "*" { + if !opts.AllowWildcard { + return nil, fmt.Errorf("wildcard selector %q is only valid for --multi-models", raw) + } + return expandWildcard(raw, name, effort) + } + + p, _, tokenMode, ok := ProviderForToken(name) + if !ok { + // An unclaimed model name is fine as long as the caller named a backend: + // that is how a brand-new provider model works before the catalog + // snapshot knows about it. Without one, fail loud. + if opts.Backend == "" { + return nil, fmt.Errorf("%w: %s (pass an explicit backend: %s)", ErrUnknownModel, name, BackendList()) + } + forced, mode, found := ProviderFor(opts.Backend) + if !found { + return nil, fmt.Errorf("invalid backend %q (valid: %s)", opts.Backend, BackendList()) + } + p, tokenMode = forced, mode + } + + // An explicit prefix wins over Model.Mode, which wins over the mode the model + // name itself implies. + if prefix == "" && opts.Mode != "" { + tokenMode = opts.Mode + } + mode, err := resolveMode(prefix, tokenMode, p, name) + if err != nil { + return nil, err + } + backend, err := p.BackendFor(mode) + if err != nil { + return nil, err + } + if opts.Backend != "" && backend != opts.Backend { + if prefix == "" { + // A bare model whose family does not match an explicitly requested + // backend: report the family clash, not the mode. + forced, _, ok := ProviderFor(opts.Backend) + if !ok { + return nil, fmt.Errorf("invalid backend %q (valid: %s)", opts.Backend, BackendList()) + } + if forced != p { + return nil, fmt.Errorf("model %q belongs to the %s family and cannot use backend %q (%s family)", + name, p.AgentName, opts.Backend, forced.AgentName) + } + // Same family, different mode: honour the explicit backend. + backend = opts.Backend + mode = modeOf(opts.Backend) + } else { + return nil, fmt.Errorf("runtime selector %q resolves to backend %q but --backend is %q", raw, backend, opts.Backend) + } + } + + resolved, err := resolveOn(p, mode, name) + if err != nil { + return nil, fmt.Errorf("runtime selector %q: %w", raw, err) + } + if err := ValidateEffort(backend, resolved, effort); err != nil { + return nil, fmt.Errorf("runtime selector %q: %w", raw, err) + } + return []Model{Model{Name: resolved, Backend: backend, Effort: effort}.Capabilities()}, nil +} + +// splitElement pulls the optional prefix and effort suffix off an element. +func splitElement(raw, baseName string) (prefix, name string, effort Effort, err error) { + tokens := strings.Split(raw, ":") + for i := range tokens { + tokens[i] = strings.TrimSpace(tokens[i]) + } + switch len(tokens) { + case 1: + // A bare prefix ("cli") borrows the base model. + if isPrefix(tokens[0]) && baseName != "" { + return strings.ToLower(tokens[0]), strings.TrimSpace(baseName), EffortNone, nil + } + return "", tokens[0], EffortNone, nil + case 2: + switch { + case Effort(tokens[1]).Valid() && tokens[1] != "": + return "", tokens[0], Effort(tokens[1]), nil + case isPrefix(tokens[0]): + return strings.ToLower(tokens[0]), tokens[1], EffortNone, nil + default: + return "", "", EffortNone, fmt.Errorf("ambiguous compact model %q (expected model:effort or mode:model)", raw) + } + case 3: + if !isPrefix(tokens[0]) { + return "", "", EffortNone, fmt.Errorf("invalid mode %q in %q (valid: %s)", tokens[0], raw, RuntimeModeList()) + } + if !Effort(tokens[2]).Valid() || tokens[2] == "" { + return "", "", EffortNone, fmt.Errorf("invalid effort %q in %q", tokens[2], raw) + } + return strings.ToLower(tokens[0]), tokens[1], Effort(tokens[2]), nil + default: + return "", "", EffortNone, fmt.Errorf("invalid compact model %q (too many ':' segments)", raw) + } +} + +// resolveMode reconciles the element's prefix with the mode the model name +// itself implies. An explicit prefix always wins. +func resolveMode(prefix string, tokenMode RuntimeMode, p *Provider, name string) (RuntimeMode, error) { + if prefix == "" { + return tokenMode, nil + } + if b := Backend(prefix); b.Valid() { + owner, mode, _ := ProviderFor(b) + if owner != p { + return "", fmt.Errorf("model %q belongs to the %s family and cannot use backend %q (%s family)", + name, p.AgentName, b, owner.AgentName) + } + return mode, nil + } + mode, ok := ParseRuntimeMode(prefix) + if !ok { + return "", fmt.Errorf("unknown runtime selector prefix %q (valid: %s, a backend name, or *)", prefix, RuntimeModeList()) + } + return mode, nil +} + +func expandWildcard(raw, name string, effort Effort) ([]Model, error) { + p, _, _, ok := ProviderForToken(name) + if !ok { + return nil, fmt.Errorf("%w: %s (pass an explicit backend: %s)", ErrUnknownModel, name, BackendList()) + } + out := make([]Model, 0, len(p.Modes())) + for _, mode := range p.Modes() { + resolved, err := resolveOn(p, mode, name) + if err != nil { + continue + } + backend, err := p.BackendFor(mode) + if err != nil { + continue + } + if err := ValidateEffort(backend, resolved, effort); err != nil { + continue + } + out = append(out, Model{Name: resolved, Backend: backend, Effort: effort}.Capabilities()) + } + if len(out) == 0 { + return nil, fmt.Errorf("runtime selector %q is not available on any backend", raw) + } + return out, nil +} + +// resolveOn maps a model token onto the exact id a provider's mode accepts, +// failing loud when the catalog knows the model but not on that mode. +func resolveOn(p *Provider, mode RuntimeMode, name string) (string, error) { + backend, err := p.BackendFor(mode) + if err != nil { + return "", err + } + if known, available := p.Availability(mode, name); known && !available { + return "", fmt.Errorf("model %q is not available on backend %q", name, backend) + } + resolved, _ := p.ResolveExact(mode, name) + return resolved, nil +} + +// isPrefix reports whether a token can lead an element: a runtime mode, a +// concrete backend name, or the wildcard. +func isPrefix(s string) bool { + s = strings.ToLower(strings.TrimSpace(s)) + if s == "*" { + return true + } + if _, ok := ParseRuntimeMode(s); ok { + return true + } + return Backend(s).Valid() +} + +// ContainsSelector reports whether s carries a prefix selector such as +// "cmux:gpt-5.6-sol" or "*:sonnet-5". Comma-separated lists are scanned per item. +func ContainsSelector(s string) bool { + for _, part := range splitCSV(s) { + prefix, _, ok := strings.Cut(part, ":") + if !ok { + continue + } + if isPrefix(prefix) { + return true + } + } + return false +} + +func modeOf(b Backend) RuntimeMode { + _, mode, _ := ProviderFor(b) + return mode +} + +// ResolveModel resolves a Model in place: its Name is parsed as a compact +// element and its fallbacks each resolved independently. It is the entry point +// for both `--model` and spec/frontmatter decoding, which is the whole point — +// they used to run different parsers. +func ResolveModel(m Model) (Model, error) { + m = m.ExpandCSV() + if strings.TrimSpace(m.Name) == "" { + return m, nil + } + if err := m.validateMode(); err != nil { + return Model{}, err + } + resolved, err := ParseModelElement(m.Name, ParseOptions{Backend: m.Backend, Mode: m.Mode}) + if err != nil { + return Model{}, err + } + if len(resolved) != 1 { + return Model{}, fmt.Errorf("wildcard selector %q is only valid for --multi-models", m.Name) + } + out := mergeResolved(m, resolved[0]) + out.Fallbacks = make(ModelList, 0, len(m.Fallbacks)) + for _, fb := range m.Fallbacks { + if strings.TrimSpace(fb.Name) == "" { + out.Fallbacks = append(out.Fallbacks, fb) + continue + } + rfb, err := ParseModelElement(fb.Name, ParseOptions{Backend: fb.Backend}) + if err != nil { + return Model{}, err + } + if len(rfb) != 1 { + return Model{}, fmt.Errorf("wildcard selector %q is only valid for --multi-models", fb.Name) + } + out.Fallbacks = append(out.Fallbacks, mergeResolved(fb, rfb[0])) + } + return out, nil +} + +// ResolveMulti expands --multi-models values into concrete runtime model/backend +// pairs. Values may be repeated and/or comma-separated, a bare prefix ("cmux") +// borrows the base model, and duplicates are dropped. Each result inherits the +// base model's temperature/cache settings, and its effort when the element does +// not set one. +func ResolveMulti(values []string, base Model) ([]Model, error) { + base, err := ResolveModel(base) + if err != nil { + return nil, err + } + out := make([]Model, 0, len(values)) + seen := map[string]bool{} + for _, value := range values { + for _, part := range splitCSV(value) { + models, err := ParseModelMulti(part, ParseOptions{BaseName: base.Name}) + if err != nil { + return nil, err + } + for _, m := range models { + m.Temperature = base.Temperature + if m.Effort == EffortNone { + m.Effort = base.Effort + } + m.NoCache = base.NoCache + key := string(m.Backend) + "\x00" + m.Name + "\x00" + string(m.Effort) + if seen[key] { + continue + } + seen[key] = true + out = append(out, m) + } + } + } + return out, nil +} + +func mergeResolved(original, resolved Model) Model { + out := original + if strings.TrimSpace(resolved.Name) != "" { + out.Name = resolved.Name + } + if resolved.ID != "" { + out.ID = resolved.ID + } + if resolved.Backend != "" { + out.Backend = resolved.Backend + } + if resolved.Effort != EffortNone { + out.Effort = resolved.Effort + } + // Capabilities are derived, never carried over from the request: whatever the + // caller wrote for them is replaced by what the resolved adapter can do. + return out.Capabilities() +} + +// IsUnknownModel reports whether err is the "no provider claims this" failure. +func IsUnknownModel(err error) bool { return errors.Is(err, ErrUnknownModel) } diff --git a/pkg/api/registry/provider.go b/pkg/api/registry/provider.go new file mode 100644 index 0000000..7902b87 --- /dev/null +++ b/pkg/api/registry/provider.go @@ -0,0 +1,200 @@ +package registry + +import ( + "fmt" + "sort" + "strings" +) + +// ModeCapabilities is one provider×mode cell: the Backend that pair serializes +// to, plus what that adapter can actually do. +// +// The runtime interface assertions (StreamingProvider, InterruptibleProvider, +// SteerableProvider) stay authoritative at execution time; this table is the +// static declaration callers can consult before constructing a provider, and the +// reason capabilities can be answered without spinning one up. +type ModeCapabilities struct { + // Backend is the serialized enum value for this provider×mode pair, e.g. + // anthropic×agent → "claude-agent". These strings are frozen: they are + // persisted in specs, session rows, and the webapp wire format. + Backend Backend + // Streaming: the adapter can stream incremental events. + Streaming bool + // Resume: the adapter can resume a prior session by id. + Resume bool + // Interrupt: the adapter implements InterruptibleProvider. + Interrupt bool + // Steer: the adapter implements SteerableProvider. + Steer bool + // MediaTypes is the adapter's attachment ceiling. A model's own declared + // types are clamped against it — the adapter cannot carry what it cannot send. + MediaTypes []string + // Keyless marks modes that never consult EnvVars because they ride the local + // CLI's own login (cmux). + Keyless bool +} + +// modeToken forces a RuntimeMode from a model-name prefix: "claude-code-…" is +// the CLI, "codex-agent-…" is the agent SDK. Matched longest-prefix-first. +type modeToken struct { + prefix string + mode RuntimeMode +} + +// Provider describes one model family — its auth env vars, catalog and pricing +// namespaces, per-mode capabilities, and the model rows it owns. +// +// It is a struct rather than an interface on purpose: there are exactly four +// instances and they are entirely data. An interface would invite divergent +// implementations, which is precisely the failure this type exists to end — this +// knowledge used to live in InferBackend, backendForMode, selectorBackend, +// splitModelProvider, selectorModelFamily, AuthEnvVars, AgentsForProvider, +// PricingIDs, orPrefix, pricingModelID, and adapterInputMediaTypes, which +// disagreed with each other. +type Provider struct { + // Name is the canonical provider key: anthropic | openai | google | deepseek. + Name string + // AgentName is the coding-agent sentinel: a bare "claude"/"codex" on a + // non-API mode means "this provider's current model". + AgentName string + // CatalogPrefix namespaces catalog/menu/genkit ids. Google's is "googleai". + CatalogPrefix string + // PricingPrefix namespaces OpenRouter-style pricing keys. Google's is + // "google" — deliberately NOT CatalogPrefix. Deriving one from the other is + // how pricing lookups silently missed and reported $0. + PricingPrefix string + // EnvVars are the API-key environment variables in priority order. + EnvVars []string + + // modes is the per-mode capability table. A missing mode means this provider + // does not support it (google has no agent or cmux row). + modes map[RuntimeMode]ModeCapabilities + // modeTokens force a mode from a model-name prefix. + modeTokens []modeToken + // claimPrefixes are the bare model-name prefixes this provider claims. A + // claim with no modeToken hit lands on ModeAPI. + claimPrefixes []string + // identityTrim are prefixes stripped before the family split. + identityTrim []string + // families are the family names this provider's model ids are built from. + families []string + // emptyTokens are tokens that name the provider but no family; they resolve + // to emptyFamily. + emptyTokens []string + emptyFamily string + + // genConfig translates effort into this provider's native request controls. + // nil means the provider has no per-request effort knob (DeepSeek). + genConfig func(cfg map[string]any, caps KnownModel, effort Effort, maxTokens int) + // classifyErr refines the shared error classification for this provider's own + // error vocabulary. nil means the shared heuristics are enough. + classifyErr func(err error, base ErrorClass) ErrorClass +} + +// SupportedEnvVars returns the auth environment variables for this provider, in +// priority order. +func (p *Provider) SupportedEnvVars() []string { + return append([]string(nil), p.EnvVars...) +} + +// Modes lists the runtime modes this provider supports, in canonical order. +func (p *Provider) Modes() []RuntimeMode { + out := make([]RuntimeMode, 0, len(p.modes)) + for _, m := range AllRuntimeModes() { + if _, ok := p.modes[m]; ok { + out = append(out, m) + } + } + return out +} + +// Caps returns the capability row for a mode. ok is false when this provider +// does not serve that mode. +func (p *Provider) Caps(mode RuntimeMode) (ModeCapabilities, bool) { + caps, ok := p.modes[mode] + return caps, ok +} + +// BackendFor maps a mode onto the serialized Backend, failing loud when the +// provider does not serve it. +func (p *Provider) BackendFor(mode RuntimeMode) (Backend, error) { + caps, ok := p.modes[mode] + if !ok { + // AgentName, not Name: users speak in families ("gemini models"), not + // provider keys ("google models"). + return "", fmt.Errorf("mode %q is not supported for %s models (supported: %s)", + mode, p.AgentName, modeListOf(p.Modes())) + } + return caps.Backend, nil +} + +// Backends lists every Backend this provider serves, in canonical mode order. +func (p *Provider) Backends() []Backend { + out := make([]Backend, 0, len(p.modes)) + for _, m := range p.Modes() { + out = append(out, p.modes[m].Backend) + } + return out +} + +// PricingIDs returns candidate pricing keys for a model, most specific first. +func (p *Provider) PricingIDs(model string) []string { + bare := p.bareID(model) + return []string{p.PricingPrefix + "/" + bare, bare} +} + +// Models returns this provider's catalog rows. +func (p *Provider) Models() []KnownModel { + out := make([]KnownModel, 0, len(knownModels)) + for _, m := range knownModels { + if m.Provider == p.Name { + out = append(out, m) + } + } + return out +} + +// claim reports whether this provider owns a model token and which mode the +// token's own prefix implies. Mode tokens are matched longest-first so +// "codex-agent-x" resolves to the agent SDK rather than the codex CLI. +func (p *Provider) claim(token string) (RuntimeMode, bool) { + t := strings.ToLower(strings.TrimSpace(token)) + for _, mt := range p.modeTokens { + if strings.HasPrefix(t, mt.prefix) { + return mt.mode, true + } + } + for _, prefix := range p.claimPrefixes { + if strings.HasPrefix(t, prefix) { + return ModeAPI, true + } + } + return "", false +} + +// bareID strips this provider's catalog namespace (and the Gemini "models/" +// namespace) from an id. +func (p *Provider) bareID(model string) string { + model = strings.TrimSpace(model) + for _, prefix := range []string{p.CatalogPrefix + "/", p.PricingPrefix + "/", p.Name + "/", "models/"} { + model = strings.TrimPrefix(model, prefix) + } + return model +} + +func modeListOf(modes []RuntimeMode) string { + parts := make([]string, len(modes)) + for i, m := range modes { + parts[i] = string(m) + } + return strings.Join(parts, ", ") +} + +// sortModeTokens orders a provider's mode tokens longest-prefix-first so the +// most specific match wins regardless of declaration order. +func sortModeTokens(tokens []modeToken) []modeToken { + sort.SliceStable(tokens, func(i, j int) bool { + return len(tokens[i].prefix) > len(tokens[j].prefix) + }) + return tokens +} diff --git a/pkg/api/registry/providers.go b/pkg/api/registry/providers.go new file mode 100644 index 0000000..b2dce7e --- /dev/null +++ b/pkg/api/registry/providers.go @@ -0,0 +1,141 @@ +package registry + +import "strings" + +// The provider descriptors. Everything captain used to rediscover with string +// prefixes and per-backend switches is a field here. +// +// Media types come from what each adapter can actually carry; Interrupt/Steer +// mirror the InterruptibleProvider/SteerableProvider implementations in +// pkg/ai/provider (steer: claude-agent only; interrupt: claude-agent and +// codex-agent). Streaming is true everywhere because every adapter implements +// ExecuteStream; it is declared rather than assumed so a future non-streaming +// adapter has an honest place to say so. +var ( + Anthropic = &Provider{ + Name: "anthropic", + AgentName: "claude", + CatalogPrefix: "anthropic", + PricingPrefix: "anthropic", + EnvVars: []string{"ANTHROPIC_API_KEY"}, + modes: map[RuntimeMode]ModeCapabilities{ + ModeAPI: {Backend: BackendAnthropic, Streaming: true, MediaTypes: []string{"image/*"}}, + ModeCLI: {Backend: BackendClaudeCLI, Streaming: true, Resume: true}, + ModeAgent: {Backend: BackendClaudeAgent, Streaming: true, Resume: true, Interrupt: true, Steer: true, MediaTypes: []string{"image/png", "image/jpeg", "image/gif", "image/webp", "application/pdf"}}, + ModeCmux: {Backend: BackendClaudeCmux, Streaming: true, Resume: true, Keyless: true}, + }, + modeTokens: sortModeTokens([]modeToken{ + {prefix: "claude-agent", mode: ModeAgent}, + {prefix: "claude-code", mode: ModeCLI}, + }), + claimPrefixes: []string{"claude", "opus", "sonnet", "haiku", "fable"}, + identityTrim: []string{"claude-agent-", "claude-code-", "claude-"}, + families: []string{"fable", "opus", "sonnet", "haiku"}, + emptyTokens: []string{""}, + emptyFamily: "opus", + genConfig: anthropicGenerationConfig, + } + + OpenAI = &Provider{ + Name: "openai", + AgentName: "codex", + CatalogPrefix: "openai", + PricingPrefix: "openai", + EnvVars: []string{"OPENAI_API_KEY"}, + modes: map[RuntimeMode]ModeCapabilities{ + ModeAPI: {Backend: BackendOpenAI, Streaming: true, MediaTypes: []string{"image/*"}}, + ModeCLI: {Backend: BackendCodexCLI, Streaming: true, Resume: true, MediaTypes: []string{"image/*"}}, + ModeAgent: {Backend: BackendCodexAgent, Streaming: true, Resume: true, Interrupt: true, MediaTypes: []string{"image/*"}}, + ModeCmux: {Backend: BackendCodexCmux, Streaming: true, Resume: true, Keyless: true}, + }, + // A bare "codex" is the CLI, not the API — the asymmetry with "claude" + // (which stays on the API) is long-standing user-visible behaviour. + // "grok-" is served through the codex CLI. + modeTokens: sortModeTokens([]modeToken{ + {prefix: "codex-agent", mode: ModeAgent}, + {prefix: "codex", mode: ModeCLI}, + }), + claimPrefixes: []string{"gpt-", "o1", "o3", "o4"}, + identityTrim: []string{"codex-agent-", "codex-"}, + families: []string{"gpt"}, + emptyTokens: []string{"", "codex"}, + // A family name, not a model id: emptyFamily is matched against + // KnownModel.Family, so an id here matches no row and the token falls + // through unresolved ("api:codex" stayed the literal "codex"). Which gpt + // model a bare "codex" lands on is decided by priority in the catalog — + // gpt-5.6-sol carries priority 1 — not by naming it here. + emptyFamily: "gpt", + genConfig: openaiGenerationConfig, + } + + Google = &Provider{ + Name: "google", + AgentName: "gemini", + // googleai for catalog/menu/genkit ids, google for pricing keys. These + // are independent on purpose; see the Provider doc. + CatalogPrefix: "googleai", + PricingPrefix: "google", + EnvVars: []string{"GEMINI_API_KEY", "GOOGLE_API_KEY"}, + modes: map[RuntimeMode]ModeCapabilities{ + ModeAPI: {Backend: BackendGemini, Streaming: true, MediaTypes: []string{"image/*", "audio/*", "video/*", "application/pdf"}}, + ModeCLI: {Backend: BackendGeminiCLI, Streaming: true}, + }, + modeTokens: sortModeTokens([]modeToken{ + {prefix: "gemini-cli", mode: ModeCLI}, + }), + claimPrefixes: []string{"gemini", "models/gemini"}, + identityTrim: []string{"gemini-cli-"}, + families: []string{"gemini"}, + genConfig: googleGenerationConfig, + } + + DeepSeek = &Provider{ + Name: "deepseek", + AgentName: "deepseek", + CatalogPrefix: "deepseek", + PricingPrefix: "deepseek", + EnvVars: []string{"DEEPSEEK_API_KEY"}, + modes: map[RuntimeMode]ModeCapabilities{ + // DeepSeek selects reasoning by model id (deepseek-reasoner vs + // deepseek-chat) and ships no attachment support. + ModeAPI: {Backend: BackendDeepSeek, Streaming: true}, + }, + claimPrefixes: []string{"deepseek"}, + families: []string{"deepseek"}, + } +) + +// Providers returns the provider families in canonical claim order. Parse walks +// this slice and the FIRST provider to claim a token wins, so the order is a +// contract, not an accident. The claim prefixes are disjoint today; the order +// also fixes AllBackends ordering. +func Providers() []*Provider { + return []*Provider{Anthropic, OpenAI, Google, DeepSeek} +} + +// ProviderByName resolves a provider by Name, CatalogPrefix, or PricingPrefix, +// so both "google" and "googleai" find Google. +func ProviderByName(name string) (*Provider, bool) { + name = strings.ToLower(strings.TrimSpace(name)) + for _, p := range Providers() { + if name == p.Name || name == p.CatalogPrefix || name == p.PricingPrefix { + return p, true + } + } + return nil, false +} + +// ProviderFor returns the descriptor that owns a backend, and the mode that +// backend represents. It replaces Backend.Provider, Backend.Kind, Backend.Family, +// registryProviderForBackend, and modelSourceBackend — one reverse lookup over +// the same table the forward direction uses. +func ProviderFor(b Backend) (*Provider, RuntimeMode, bool) { + for _, p := range Providers() { + for _, mode := range p.Modes() { + if p.modes[mode].Backend == b { + return p, mode, true + } + } + } + return nil, "", false +} diff --git a/pkg/api/registry/supported_models.go b/pkg/api/registry/supported_models.go new file mode 100644 index 0000000..85786bd --- /dev/null +++ b/pkg/api/registry/supported_models.go @@ -0,0 +1,96 @@ +package registry + +import ( + "context" + "sort" + "sync" +) + +// liveModelHooks are the optional live model-listing functions, keyed by +// provider name. pkg/ai installs them from its provider init(), mirroring the +// existing RegisterProvider/factories seam: the HTTP calls and key resolution +// they need live above this package, but the merged view belongs here. +var ( + liveModelMu sync.RWMutex + liveModelHooks = map[string]func(context.Context) ([]KnownModel, error){} +) + +// RegisterLiveModels installs a live model-listing hook for a provider. Passing +// nil clears it. Registration is global and read by SupportedModels. +func RegisterLiveModels(provider string, fn func(context.Context) ([]KnownModel, error)) { + liveModelMu.Lock() + defer liveModelMu.Unlock() + if fn == nil { + delete(liveModelHooks, provider) + return + } + liveModelHooks[provider] = fn +} + +func (p *Provider) liveModels(ctx context.Context) ([]KnownModel, error) { + liveModelMu.RLock() + fn := liveModelHooks[p.Name] + liveModelMu.RUnlock() + if fn == nil { + return nil, nil + } + return fn(ctx) +} + +// SupportedModels returns the models this provider offers, as resolved Models +// carrying their backend, mode, and capabilities. +// +// The catalog snapshot is the floor: a live listing (when a hook is registered +// and the call succeeds) adds models the snapshot has not caught up with and +// refreshes what it knows, but a live failure degrades to the snapshot rather +// than emptying the list. Results are ordered newest-first. +// +// Models are reported on the provider's API mode by default; SupportedModelsFor +// answers for a specific mode. +func (p *Provider) SupportedModels(ctx context.Context) []Model { + return p.SupportedModelsFor(ctx, ModeAPI) +} + +// SupportedModelsFor returns the models this provider offers on one mode. +func (p *Provider) SupportedModelsFor(ctx context.Context, mode RuntimeMode) []Model { + backend, err := p.BackendFor(mode) + if err != nil { + return nil + } + + merged := map[string]KnownModel{} + order := make([]string, 0) + add := func(m KnownModel) { + if !p.availableFor(m, mode) { + return + } + if _, seen := merged[m.ID]; !seen { + order = append(order, m.ID) + } + merged[m.ID] = m + } + for _, m := range p.Models() { + if m.Preferred { + add(m) + } + } + // A live listing refines the snapshot; it never replaces it, so a provider + // outage cannot empty the catalog. + if live, err := p.liveModels(ctx); err == nil { + for _, m := range live { + if m.Provider == "" { + m.Provider = p.Name + } + add(m) + } + } + + out := make([]Model, 0, len(order)) + for _, id := range order { + out = append(out, Model{Name: id, Backend: backend}.Capabilities()) + } + sort.SliceStable(out, func(i, j int) bool { + return merged[out[i].Name].ReleaseDate > merged[out[j].Name].ReleaseDate + }) + return out +} diff --git a/pkg/api/runtime_config.go b/pkg/api/runtime_config.go index 3ff3aaa..9e04946 100644 --- a/pkg/api/runtime_config.go +++ b/pkg/api/runtime_config.go @@ -28,6 +28,14 @@ type PermissionDecision struct { UpdatedInput map[string]any } +// SchemaRepairConfig controls the optional second pass used when structured +// output fails local JSON-schema validation. Empty means use the parent +// provider/model and captain's embedded repair prompt. +type SchemaRepairConfig struct { + Model Model // optional override; empty means the parent model/backend + Prompt string // optional .prompt file path; empty means embedded default +} + // Config is the provider construction/runtime config. Model (name/backend/temp/ // effort) and Budget (cost ceiling, max tokens) come from the serializable spec // types; the rest are transport/runtime concerns that never belong in Spec. It is @@ -44,6 +52,7 @@ type Config struct { MaxConcurrent int SessionID string ProjectName string + SchemaRepair SchemaRepairConfig // CanUseTool, when set, brokers tool permissions over the stream-json control // protocol: the streaming provider asks this callback before a tool that needs @@ -52,4 +61,10 @@ type Config struct { // others ignore it. A nil callback keeps the auto-approve (bypass) behaviour. // It is never serialized (the agent process never sees the Go closure). CanUseTool PermissionFunc `json:"-"` + + // Tools are caller-supplied tools exposed to the model and executed + // in-process. Only tool-capable providers (see ToolCapableProvider — today + // the genkit API backends) honour them; other providers, which bring their + // own tool ecosystems, ignore the field. Never serialized (Go closures). + Tools []ToolDefinition `json:"-"` } diff --git a/pkg/api/runtime_event.go b/pkg/api/runtime_event.go index a97aa61..f54c49a 100644 --- a/pkg/api/runtime_event.go +++ b/pkg/api/runtime_event.go @@ -7,14 +7,22 @@ import ( // Response is the result of a buffered (non-streaming) provider execution. type Response struct { - Text string - StructuredData any - Model string - Backend Backend - Usage Usage - Duration time.Duration - CacheHit bool - Raw any + Text string + StructuredData any + TerminalOutcome *TerminalOutcome + ToolApproval *ToolApprovalState + Model string + Backend Backend + Usage Usage + // CostUSD is the response's reported cost: the provider's authoritative value + // when it supplies one (claude-cli total_cost_usd, claude-agent cost_usd), + // otherwise the provider's list-price estimate. 0 means no cost was reported + // (the buffered path used to drop it — see finding D4). Consumers should + // prefer this over recomputing from tokens. + CostUSD float64 + Duration time.Duration + CacheHit bool + Raw any // Workspace is the run's working-dir runtime state (cwd, git details, changed // files, commits, plan). Populated by the agent runner + worktree hook. @@ -64,6 +72,7 @@ type Event struct { // is raw JSON because the streaming contract does not know the caller's Go // type — the buffered Execute path unmarshals it into Request.Prompt.Schema. StructuredData json.RawMessage + ToolApproval *ToolApprovalState // Raw carries the backend-native event (e.g. claude.HistoryEntry for the // claude_cli stream) so renderers can use the rich pretty-printers in diff --git a/pkg/api/runtime_provider.go b/pkg/api/runtime_provider.go index 1700429..3210501 100644 --- a/pkg/api/runtime_provider.go +++ b/pkg/api/runtime_provider.go @@ -15,3 +15,42 @@ type StreamingProvider interface { Provider ExecuteStream(ctx context.Context, req Spec) (<-chan Event, error) } + +type ProviderUnwrapper interface { + Unwrap() Provider +} + +type InterruptibleProvider interface { + Interrupt(context.Context) error +} + +type SteerableProvider interface { + Steer(context.Context, Spec) error +} + +type CloseableProvider interface { + Close() error +} + +// ToolCapableProvider is implemented by providers that can expose and execute +// the caller-supplied Config.Tools (rather than only the backend's built-in +// tools). Resolve it with ProviderAs[ToolCapableProvider] to check support +// before relying on Config.Tools being honoured. +type ToolCapableProvider interface { + SupportsCallerTools() bool +} + +func ProviderAs[T any](provider Provider) (T, bool) { + for provider != nil { + if capability, ok := any(provider).(T); ok { + return capability, true + } + wrapper, ok := provider.(ProviderUnwrapper) + if !ok { + break + } + provider = wrapper.Unwrap() + } + var zero T + return zero, false +} diff --git a/pkg/api/runtime_provider_ginkgo_test.go b/pkg/api/runtime_provider_ginkgo_test.go new file mode 100644 index 0000000..a17adac --- /dev/null +++ b/pkg/api/runtime_provider_ginkgo_test.go @@ -0,0 +1,78 @@ +package api_test + +import ( + "context" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/credentials" +) + +type runtimeProviderStub struct{} + +func (runtimeProviderStub) Execute(context.Context, api.Spec) (*api.Response, error) { + return &api.Response{}, nil +} +func (runtimeProviderStub) GetModel() string { return "test-model" } +func (runtimeProviderStub) GetBackend() api.Backend { return api.BackendClaudeAgent } +func (runtimeProviderStub) Interrupt(context.Context) error { return nil } + +type runtimeProviderWrapper struct{ inner api.Provider } + +func (w runtimeProviderWrapper) Execute(ctx context.Context, req api.Spec) (*api.Response, error) { + return w.inner.Execute(ctx, req) +} +func (w runtimeProviderWrapper) GetModel() string { return w.inner.GetModel() } +func (w runtimeProviderWrapper) GetBackend() api.Backend { return w.inner.GetBackend() } +func (w runtimeProviderWrapper) Unwrap() api.Provider { return w.inner } + +var _ = Describe("ProviderAs", func() { + It("finds a provider capability through nested wrappers", func() { + provider := runtimeProviderWrapper{inner: runtimeProviderWrapper{inner: runtimeProviderStub{}}} + + interruptible, ok := api.ProviderAs[api.InterruptibleProvider](provider) + + Expect(ok).To(BeTrue()) + Expect(interruptible.Interrupt(context.Background())).To(Succeed()) + }) + + It("returns false for nil and missing capabilities", func() { + var provider api.Provider + + _, ok := api.ProviderAs[api.InterruptibleProvider](provider) + Expect(ok).To(BeFalse()) + + _, ok = api.ProviderAs[api.CloseableProvider](runtimeProviderStub{}) + Expect(ok).To(BeFalse()) + }) +}) + +var _ = Describe("API key resolution", func() { + BeforeEach(func() { + credentials.SetPathForTesting(filepath.Join(GinkgoT().TempDir(), "vault")) + DeferCleanup(func() { credentials.SetPathForTesting("") }) + GinkgoT().Setenv("OPENAI_API_KEY", "environment-token") + }) + + It("prefers the Captain vault over the environment", func() { + vault, err := credentials.DefaultVault() + Expect(err).NotTo(HaveOccurred()) + Expect(vault.Set("openai", "vault-token")).To(Succeed()) + + resolved, err := api.ResolveAPIKey(api.BackendOpenAI) + Expect(err).NotTo(HaveOccurred()) + Expect(resolved.Token).To(Equal("vault-token")) + Expect(resolved.Source).To(Equal(credentials.SourceVault)) + }) + + It("falls back to the environment", func() { + resolved, err := api.ResolveAPIKey(api.BackendOpenAI) + Expect(err).NotTo(HaveOccurred()) + Expect(resolved.Token).To(Equal("environment-token")) + Expect(resolved.Source).To(Equal(credentials.SourceEnvironment)) + Expect(resolved.Detail).To(Equal("OPENAI_API_KEY")) + }) +}) diff --git a/pkg/api/runtime_registry.go b/pkg/api/runtime_registry.go index 25e0daa..46c616b 100644 --- a/pkg/api/runtime_registry.go +++ b/pkg/api/runtime_registry.go @@ -3,8 +3,16 @@ package api import ( "fmt" "os" + + "github.com/flanksource/captain/pkg/credentials" ) +type ResolvedAPIKey struct { + Token string + Source string + Detail string +} + // ProviderFactory constructs a Provider for a backend from a Config. type ProviderFactory func(cfg Config) (Provider, error) @@ -44,12 +52,38 @@ func NewProvider(cfg Config) (Provider, error) { } if cfg.APIKey == "" { - cfg.APIKey = GetAPIKeyFromEnv(backend) + resolved, err := ResolveAPIKey(backend) + if err != nil { + return nil, err + } + cfg.APIKey = resolved.Token } return factory(cfg) } +// ResolveAPIKey resolves a direct provider credential from Captain's vault, +// then from the provider's supported environment variables. +func ResolveAPIKey(backend Backend) (ResolvedAPIKey, error) { + if backend.Kind() != "api" { + for _, envVar := range AuthEnvVars(backend) { + if token := os.Getenv(envVar); token != "" { + return ResolvedAPIKey{Token: token, Source: credentials.SourceEnvironment, Detail: envVar}, nil + } + } + return ResolvedAPIKey{}, nil + } + vault, err := credentials.DefaultVault() + if err != nil { + return ResolvedAPIKey{}, err + } + resolved, err := vault.Resolve(string(backend), AuthEnvVars(backend), os.Getenv) + if err != nil { + return ResolvedAPIKey{}, err + } + return ResolvedAPIKey(resolved), nil +} + // GetAPIKeyFromEnv returns the first non-empty value among a backend's auth // environment variables. func GetAPIKeyFromEnv(backend Backend) string { diff --git a/pkg/api/schema_test.go b/pkg/api/schema_test.go index f0dc450..eef39ec 100644 --- a/pkg/api/schema_test.go +++ b/pkg/api/schema_test.go @@ -18,7 +18,7 @@ func TestSchemaJSON(t *testing.T) { t.Fatalf("schema is not valid JSON: %v", err) } s := string(data) - for _, want := range []string{"xhigh", "Permissions", "Setup", "Budget", `"maxTokens"`} { + for _, want := range []string{"xhigh", "max", "ultra", "Permissions", "Setup", "Budget", `"maxTokens"`} { if !strings.Contains(s, want) { t.Errorf("schema missing %q\n%s", want, s) } diff --git a/pkg/api/serve.go b/pkg/api/serve.go new file mode 100644 index 0000000..e4fb55d --- /dev/null +++ b/pkg/api/serve.go @@ -0,0 +1,38 @@ +package api + +import ( + "os" + "strings" +) + +// ServeURLEnv names the environment variable captain serve exports with its +// listen address, so hook receivers and captain-launched agents find the right +// instance even off the default port. +const ServeURLEnv = "CAPTAIN_SERVER_URL" + +// DefaultServeURL is captain serve's default listen address. +const DefaultServeURL = "http://localhost:9020" + +// MonitorHooksEnv disables captain's session-monitoring hook injection when +// set to "off" — for fixtures and CI runs that need deterministic agent argv. +const MonitorHooksEnv = "CAPTAIN_MONITOR_HOOKS" + +// ServeBaseURL resolves the running captain serve instance: $CAPTAIN_SERVER_URL +// when set (exported by serve itself), else the default port on localhost. +func ServeBaseURL() string { + if url := strings.TrimSpace(os.Getenv(ServeURLEnv)); url != "" { + return url + } + return DefaultServeURL +} + +// MonitorHooksEnabled reports whether captain injects its session-monitoring +// hooks into an agent it launches. The hooks are captain infrastructure, not +// user hooks — Memory.SkipHooks does not suppress them. Bare runs and +// CAPTAIN_MONITOR_HOOKS=off opt out. +func MonitorHooksEnabled(spec Spec) bool { + if os.Getenv(MonitorHooksEnv) == "off" { + return false + } + return !spec.Memory.Bare && !spec.Permissions.HasPreset(PresetBare) +} diff --git a/pkg/api/serve_test.go b/pkg/api/serve_test.go new file mode 100644 index 0000000..138a916 --- /dev/null +++ b/pkg/api/serve_test.go @@ -0,0 +1,25 @@ +package api + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestServeBaseURL(t *testing.T) { + t.Setenv(ServeURLEnv, "") + assert.Equal(t, DefaultServeURL, ServeBaseURL()) + t.Setenv(ServeURLEnv, "http://localhost:9999") + assert.Equal(t, "http://localhost:9999", ServeBaseURL()) +} + +func TestMonitorHooksEnabled(t *testing.T) { + t.Setenv(MonitorHooksEnv, "") + assert.True(t, MonitorHooksEnabled(Spec{}), "default requests carry monitoring hooks") + assert.True(t, MonitorHooksEnabled(Spec{Memory: Memory{SkipHooks: true}}), + "SkipHooks governs user hooks, not captain's monitoring hooks") + assert.False(t, MonitorHooksEnabled(Spec{Memory: Memory{Bare: true}}), "bare runs opt out") + + t.Setenv(MonitorHooksEnv, "off") + assert.False(t, MonitorHooksEnabled(Spec{}), "CAPTAIN_MONITOR_HOOKS=off opts out") +} diff --git a/pkg/api/spec.go b/pkg/api/spec.go index de16722..8e8c19d 100644 --- a/pkg/api/spec.go +++ b/pkg/api/spec.go @@ -1,7 +1,9 @@ package api import ( + "encoding/json" "fmt" + "reflect" "strings" "github.com/flanksource/commons-db/shell" @@ -17,11 +19,16 @@ import ( // domain object. type Spec struct { Model `json:",inline" yaml:",inline"` - Prompt Prompt `json:"prompt" yaml:"prompt"` - Budget Budget `json:"budget,omitempty" yaml:"budget,omitempty"` - Memory Memory `json:"memory,omitempty" yaml:"memory,omitempty"` - Permissions Permissions `json:"permissions,omitempty" yaml:"permissions,omitempty"` - Setup *shell.Setup `json:"setup,omitempty" yaml:"setup,omitempty"` + Prompt Prompt `json:"prompt" yaml:"prompt"` + Messages []Message `json:"messages,omitempty" yaml:"messages,omitempty" pretty:"-"` + Budget Budget `json:"budget,omitempty" yaml:"budget,omitempty"` + Memory Memory `json:"memory,omitempty" yaml:"memory,omitempty"` + Permissions Permissions `json:"permissions,omitempty" yaml:"permissions,omitempty"` + // ToolPreferences is the serializable per-turn tool/group selection policy. + // Executable tool handlers remain in Config.Tools. + ToolPreferences ToolPreferences `json:"toolPreferences,omitempty" yaml:"toolPreferences,omitempty" pretty:"-"` + ToolApproval *ToolApprovalResume `json:"toolApproval,omitempty" yaml:"toolApproval,omitempty" pretty:"-"` + Setup *shell.Setup `json:"setup,omitempty" yaml:"setup,omitempty"` // Workflow declares the generate→verify loop (verification + finalize) around // the run. Absent = single generation, no verification. @@ -36,14 +43,134 @@ type Spec struct { CLIArgs map[string]any `json:"cliArgs,omitempty" yaml:"cliArgs,omitempty"` } +type specMarshal struct { + Model `json:",inline" yaml:",inline"` + Prompt *Prompt `json:"prompt,omitempty" yaml:"prompt,omitempty"` + Messages []Message `json:"messages,omitempty" yaml:"messages,omitempty"` + Budget *Budget `json:"budget,omitempty" yaml:"budget,omitempty"` + Memory *Memory `json:"memory,omitempty" yaml:"memory,omitempty"` + Permissions *Permissions `json:"permissions,omitempty" yaml:"permissions,omitempty"` + Preferences *ToolPreferences `json:"toolPreferences,omitempty" yaml:"toolPreferences,omitempty"` + Approval *ToolApprovalResume `json:"toolApproval,omitempty" yaml:"toolApproval,omitempty"` + Setup *shell.Setup `json:"setup,omitempty" yaml:"setup,omitempty"` + Workflow *Workflow `json:"workflow,omitempty" yaml:"workflow,omitempty"` + SessionID string `json:"sessionId,omitempty" yaml:"sessionId,omitempty"` + CLIArgs map[string]any `json:"cliArgs,omitempty" yaml:"cliArgs,omitempty"` +} + +func isEmpty(value reflect.Value) bool { + if !value.IsValid() { + return true + } + for value.Kind() == reflect.Interface || value.Kind() == reflect.Pointer { + if value.IsNil() { + return true + } + value = value.Elem() + } + if value.CanInterface() { + switch typed := value.Interface().(type) { + case Tools: + return len(typed.Policies()) == 0 + case MCP: + return !typed.Disabled && len(typed.Servers) == 0 && len(typed.Modes) == 0 + } + } + switch value.Kind() { + case reflect.Array: + for i := range value.Len() { + if !isEmpty(value.Index(i)) { + return false + } + } + return true + case reflect.Map, reflect.Slice, reflect.String: + return value.Len() == 0 + case reflect.Struct: + exported := 0 + for i := range value.NumField() { + field := value.Type().Field(i) + if !field.IsExported() || field.Tag.Get("json") == "-" || field.Tag.Get("yaml") == "-" { + continue + } + exported++ + if !isEmpty(value.Field(i)) { + return false + } + } + return exported > 0 || value.IsZero() + default: + return value.IsZero() + } +} + +func omitEmptyValue[T any](value T) *T { + if isEmpty(reflect.ValueOf(value)) { + return nil + } + return &value +} + +func omitEmptyPointer[T any](value *T) *T { + if isEmpty(reflect.ValueOf(value)) { + return nil + } + return value +} + +func (s Spec) marshalValue() specMarshal { + return specMarshal{ + Model: s.Model, + Prompt: omitEmptyValue(s.Prompt), + Messages: s.Messages, + Budget: omitEmptyValue(s.Budget), + Memory: omitEmptyValue(s.Memory), + Permissions: omitEmptyValue(s.Permissions), + Preferences: omitEmptyValue(s.ToolPreferences), + Approval: omitEmptyPointer(s.ToolApproval), + Setup: omitEmptyPointer(s.Setup), + Workflow: omitEmptyPointer(s.Workflow), + SessionID: s.SessionID, + CLIArgs: s.CLIArgs, + } +} + +func (s Spec) MarshalJSON() ([]byte, error) { + return json.Marshal(s.marshalValue()) +} + +func (s Spec) MarshalYAML() (any, error) { + return s.marshalValue(), nil +} + // Validate runs each component's validation, failing loud on the first error. func (s Spec) Validate() error { if err := s.Model.Validate(); err != nil { return fmt.Errorf("model: %w", err) } - // A verify-only spec (no body, workflow.verify present) legitimately has an - // empty prompt; only its strictness setting is checked. - if s.IsVerifyOnly() { + if s.ToolApproval != nil { + if s.hasPromptBody() || len(s.Messages) > 0 { + return fmt.Errorf("tool approval resume state, prompt body, and messages are mutually exclusive request modes") + } + if err := s.ToolApproval.Validate(); err != nil { + return fmt.Errorf("tool approval: %w", err) + } + if err := s.Prompt.SchemaStrictness.Validate(); err != nil { + return fmt.Errorf("prompt: %w", err) + } + } else if len(s.Messages) > 0 { + if err := s.ValidateRequestMode(); err != nil { + return err + } + if err := ValidateMessages(s.Messages); err != nil { + return fmt.Errorf("messages: %w", err) + } + if err := s.Prompt.SchemaStrictness.Validate(); err != nil { + return fmt.Errorf("prompt: %w", err) + } + // A verify-only spec (no body, workflow.verify present) legitimately has an + // empty prompt; only its strictness setting is checked. + } else if s.IsVerifyOnly() { if err := s.Prompt.SchemaStrictness.Validate(); err != nil { return fmt.Errorf("prompt: %w", err) } @@ -56,6 +183,9 @@ func (s Spec) Validate() error { if err := s.Permissions.Validate(); err != nil { return fmt.Errorf("permissions: %w", err) } + if err := s.ToolPreferences.Validate(); err != nil { + return err + } if err := s.Workflow.Validate(); err != nil { return fmt.Errorf("workflow: %w", err) } @@ -66,7 +196,23 @@ func (s Spec) Validate() error { // verification — a verify-only run that skips generation and verifies the // current state (e.g. scoring already-committed work). func (s Spec) IsVerifyOnly() bool { - return s.Prompt.User == "" && s.Workflow != nil && s.Workflow.Verify != nil + return s.ToolApproval == nil && len(s.Messages) == 0 && s.Prompt.User == "" && len(s.Prompt.Attachments) == 0 && s.Workflow != nil && s.Workflow.Verify != nil +} + +func (s Spec) hasPromptBody() bool { + return s.Prompt.User != "" || s.Prompt.System != "" || s.Prompt.AppendSystem != "" || len(s.Prompt.Attachments) > 0 +} + +// ValidateRequestMode rejects mixing canonical conversation history with the +// single-turn prompt body. +func (s Spec) ValidateRequestMode() error { + if s.ToolApproval != nil && (len(s.Messages) > 0 || s.hasPromptBody()) { + return fmt.Errorf("tool approval resume state, prompt body, and messages are mutually exclusive request modes") + } + if len(s.Messages) > 0 && s.hasPromptBody() { + return fmt.Errorf("prompt body and messages are mutually exclusive request modes") + } + return nil } func (s Spec) Cwd() string { diff --git a/pkg/api/spec_marshal_ginkgo_test.go b/pkg/api/spec_marshal_ginkgo_test.go new file mode 100644 index 0000000..ec0b511 --- /dev/null +++ b/pkg/api/spec_marshal_ginkgo_test.go @@ -0,0 +1,43 @@ +package api + +import ( + "encoding/json" + + "github.com/flanksource/commons-db/shell" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" +) + +var _ = Describe("Spec serialization", func() { + DescribeTable("omits empty strings and sections", + func(marshal func(any) ([]byte, error), decode func([]byte, any) error) { + encoded, err := marshal(Spec{ + Memory: Memory{Skills: []string{}}, + Permissions: Permissions{Tools: Tools{Modes: map[string]ToolMode{}}}, + Setup: &shell.Setup{}, + Workflow: &Workflow{}, + }) + Expect(err).NotTo(HaveOccurred()) + + var decoded map[string]any + Expect(decode(encoded, &decoded)).To(Succeed()) + Expect(decoded).To(BeEmpty()) + }, + Entry("JSON", json.Marshal, json.Unmarshal), + Entry("YAML", yaml.Marshal, yaml.Unmarshal), + ) + + DescribeTable("preserves non-empty values while omitting empty neighbors", + func(marshal func(any) ([]byte, error), decode func([]byte, any) error) { + encoded, err := marshal(Spec{Model: Model{Name: "opus"}}) + Expect(err).NotTo(HaveOccurred()) + + var decoded map[string]any + Expect(decode(encoded, &decoded)).To(Succeed()) + Expect(decoded).To(Equal(map[string]any{"model": "opus"})) + }, + Entry("JSON", json.Marshal, json.Unmarshal), + Entry("YAML", yaml.Marshal, yaml.Unmarshal), + ) +}) diff --git a/pkg/api/spec_merge.go b/pkg/api/spec_merge.go new file mode 100644 index 0000000..28bf8a8 --- /dev/null +++ b/pkg/api/spec_merge.go @@ -0,0 +1,135 @@ +package api + +// Merge returns a copy of s with override's set (non-zero) fields taking +// precedence. A zero-valued field in override is treated as "unset" and keeps +// s's value, so a base spec can supply defaults that an operation-specific spec +// selectively overrides: +// +// resolved := base.Merge(operation) +// +// Scalar fields (model name, effort, budget cost, prompt user, …) merge +// individually. Slices, maps, and pointers (Fallbacks, Metadata, Setup, +// Workflow, CLIArgs, Permissions sub-values) replace wholesale when set in +// override rather than deep-merging — an override that lists tools means exactly +// those tools. Boolean toggles (NoCache, Skip*) follow zero=unset: an override +// can turn a flag on but not off, since false is indistinguishable from absent. +func (s Spec) Merge(override Spec) Spec { + s.Model = s.Model.Merge(override.Model) + s.Prompt = s.Prompt.merge(override.Prompt) + if len(override.Messages) > 0 { + s.Messages = override.Messages + } + s.Budget = s.Budget.merge(override.Budget) + s.Memory = s.Memory.merge(override.Memory) + s.Permissions = s.Permissions.merge(override.Permissions) + if len(override.ToolPreferences) > 0 { + s.ToolPreferences = override.ToolPreferences + } + if override.ToolApproval != nil { + s.ToolApproval = override.ToolApproval + } + if override.Setup != nil { + s.Setup = override.Setup + } + if override.Workflow != nil { + s.Workflow = override.Workflow + } + if override.SessionID != "" { + s.SessionID = override.SessionID + } + if len(override.CLIArgs) > 0 { + s.CLIArgs = override.CLIArgs + } + return s +} + +func (p Prompt) merge(o Prompt) Prompt { + if o.User != "" { + p.User = o.User + } + if o.System != "" { + p.System = o.System + } + if o.AppendSystem != "" { + p.AppendSystem = o.AppendSystem + } + if o.Source != "" { + p.Source = o.Source + } + if o.Schema != nil { + p.Schema = o.Schema + } + if len(o.SchemaJSON) > 0 { + p.SchemaJSON = o.SchemaJSON + } + if o.SchemaStrictness != "" { + p.SchemaStrictness = o.SchemaStrictness + } + if len(o.Metadata) > 0 { + p.Metadata = o.Metadata + } + return p +} + +func (b Budget) merge(o Budget) Budget { + if o.Cost != 0 { + b.Cost = o.Cost + } + if o.MaxTokens != 0 { + b.MaxTokens = o.MaxTokens + } + if o.MaxTurns != 0 { + b.MaxTurns = o.MaxTurns + } + if o.Timeout != "" { + b.Timeout = o.Timeout + } + return b +} + +func (m Memory) merge(o Memory) Memory { + if len(o.Skills) > 0 { + m.Skills = o.Skills + } + if o.SkipProject { + m.SkipProject = true + } + if o.SkipUser { + m.SkipUser = true + } + if o.SkipSkills { + m.SkipSkills = true + } + if o.SkipHooks { + m.SkipHooks = true + } + if o.SkipMemory { + m.SkipMemory = true + } + if o.Bare { + m.Bare = true + } + return m +} + +func (p Permissions) merge(o Permissions) Permissions { + if o.Mode != "" { + p.Mode = o.Mode + } + if len(o.Presets) > 0 { + p.Presets = o.Presets + } + if len(o.Tools.Allow) > 0 || len(o.Tools.Deny) > 0 || len(o.Tools.Modes) > 0 { + p.Tools = o.Tools + } + if o.MCP.Disabled || len(o.MCP.Servers) > 0 || len(o.MCP.Modes) > 0 { + p.MCP = o.MCP + } + if len(o.Plugins) > 0 { + p.Plugins = o.Plugins + } + if len(o.Skills) > 0 { + p.Skills = o.Skills + } + return p +} diff --git a/pkg/api/spec_merge_test.go b/pkg/api/spec_merge_test.go new file mode 100644 index 0000000..2555c03 --- /dev/null +++ b/pkg/api/spec_merge_test.go @@ -0,0 +1,153 @@ +package api + +import ( + "reflect" + "testing" +) + +func TestSpec_Merge(t *testing.T) { + cases := []struct { + name string + base Spec + override Spec + check func(t *testing.T, got Spec) + }{ + { + name: "empty override preserves base", + base: sampleSpec(), + override: Spec{}, + check: func(t *testing.T, got Spec) { + if !reflect.DeepEqual(got, sampleSpec()) { + t.Fatalf("empty override mutated base:\n got=%+v", got) + } + }, + }, + { + name: "override model name wins, base budget/effort kept", + base: Spec{Model: Model{Name: "base-model", Effort: EffortMedium}, Budget: Budget{Cost: 5}}, + override: Spec{Model: Model{Name: "op-model"}}, + check: func(t *testing.T, got Spec) { + if got.Model.Name != "op-model" { + t.Errorf("Name = %q, want op-model", got.Model.Name) + } + if got.Model.Effort != EffortMedium { + t.Errorf("Effort = %q, want medium (from base)", got.Model.Effort) + } + if got.Budget.Cost != 5 { + t.Errorf("Budget.Cost = %v, want 5 (from base)", got.Budget.Cost) + } + }, + }, + { + name: "empty override model keeps base model", + base: Spec{Model: Model{Name: "base-model", Effort: EffortHigh}}, + override: Spec{Budget: Budget{Cost: 3}}, + check: func(t *testing.T, got Spec) { + if got.Model.Name != "base-model" || got.Model.Effort != EffortHigh { + t.Errorf("model = %+v, want base preserved", got.Model) + } + if got.Budget.Cost != 3 { + t.Errorf("Budget.Cost = %v, want 3 (override)", got.Budget.Cost) + } + }, + }, + { + name: "budget merges field-wise", + base: Spec{Model: Model{Name: "m"}, Budget: Budget{Cost: 5, MaxTokens: 1000, MaxTurns: 10}}, + override: Spec{Budget: Budget{Cost: 9}}, + check: func(t *testing.T, got Spec) { + if got.Budget.Cost != 9 { + t.Errorf("Cost = %v, want 9", got.Budget.Cost) + } + if got.Budget.MaxTokens != 1000 || got.Budget.MaxTurns != 10 { + t.Errorf("budget lost base fields: %+v", got.Budget) + } + }, + }, + { + name: "prompt user overrides, base system kept", + base: Spec{Model: Model{Name: "m"}, Prompt: Prompt{User: "base body", System: "be precise"}}, + override: Spec{Prompt: Prompt{User: "op body"}}, + check: func(t *testing.T, got Spec) { + if got.Prompt.User != "op body" { + t.Errorf("User = %q, want op body", got.Prompt.User) + } + if got.Prompt.System != "be precise" { + t.Errorf("System = %q, want base system kept", got.Prompt.System) + } + }, + }, + { + name: "fallbacks replace wholesale", + base: Spec{Model: Model{Name: "m", Fallbacks: []Model{{Name: "a"}, {Name: "b"}}}}, + override: Spec{Model: Model{Fallbacks: []Model{{Name: "c"}}}}, + check: func(t *testing.T, got Spec) { + if len(got.Model.Fallbacks) != 1 || got.Model.Fallbacks[0].Name != "c" { + t.Errorf("Fallbacks = %+v, want [c]", got.Model.Fallbacks) + } + }, + }, + { + name: "temperature pointer: nil override keeps base", + base: Spec{Model: Model{Name: "m", Temperature: floatPtr(0.7)}}, + override: Spec{Model: Model{Name: "m"}}, + check: func(t *testing.T, got Spec) { + if got.Model.Temperature == nil || *got.Model.Temperature != 0.7 { + t.Errorf("Temperature = %v, want 0.7 kept", got.Model.Temperature) + } + }, + }, + { + name: "temperature pointer: explicit 0.0 override wins", + base: Spec{Model: Model{Name: "m", Temperature: floatPtr(0.7)}}, + override: Spec{Model: Model{Temperature: floatPtr(0)}}, + check: func(t *testing.T, got Spec) { + if got.Model.Temperature == nil || *got.Model.Temperature != 0 { + t.Errorf("Temperature = %v, want explicit 0.0", got.Model.Temperature) + } + }, + }, + { + name: "bool noCache: base true survives override false", + base: Spec{Model: Model{Name: "m", NoCache: true}}, + override: Spec{Model: Model{Name: "m", NoCache: false}}, + check: func(t *testing.T, got Spec) { + if !got.Model.NoCache { + t.Errorf("NoCache = false, want base true preserved (false=unset)") + } + }, + }, + { + name: "setup pointer replaced when set in override", + base: Spec{Model: Model{Name: "m"}}, + override: Spec{Model: Model{Name: "m"}, SessionID: "sess-op"}, + check: func(t *testing.T, got Spec) { + if got.SessionID != "sess-op" { + t.Errorf("SessionID = %q, want sess-op", got.SessionID) + } + }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := tc.base.Merge(tc.override) + tc.check(t, got) + }) + } +} + +// TestSpec_Merge_DoesNotMutateInputs guards that Merge is pure — neither the base +// nor the override is modified (slice/map aliasing aside, the top-level structs +// must be untouched so a shared base spec can be merged repeatedly). +func TestSpec_Merge_DoesNotMutateInputs(t *testing.T) { + base := Spec{Model: Model{Name: "base", Effort: EffortLow}, Budget: Budget{Cost: 1}} + override := Spec{Model: Model{Name: "op"}, Budget: Budget{Cost: 2}} + baseCopy, overrideCopy := base, override + _ = base.Merge(override) + if !reflect.DeepEqual(base, baseCopy) { + t.Errorf("Merge mutated base: %+v != %+v", base, baseCopy) + } + if !reflect.DeepEqual(override, overrideCopy) { + t.Errorf("Merge mutated override: %+v != %+v", override, overrideCopy) + } +} diff --git a/pkg/api/spec_test.go b/pkg/api/spec_test.go index d6caaac..e2f92b6 100644 --- a/pkg/api/spec_test.go +++ b/pkg/api/spec_test.go @@ -158,3 +158,28 @@ func TestSpec_Validate(t *testing.T) { }) } } + +// TestSpec_FallbacksCompact pins the end-to-end config shape: a Spec whose model +// has compact-string fallbacks, plus the primary parsed via Expand. The grammar +// itself is covered in pkg/api/registry; this pins that Spec decoding reaches it. +func TestSpec_FallbacksCompact(t *testing.T) { + var spec Spec + src := "model: opus\neffort: high\nfallbacks:\n - agent:sonnet:medium\n - api:gpt-5.5\n" + if err := yaml.Unmarshal([]byte(src), &spec); err != nil { + t.Fatal(err) + } + if spec.Model.Name != "opus" || spec.Model.Effort != EffortHigh { + t.Fatalf("primary = %+v", spec.Model) + } + if len(spec.Model.Fallbacks) != 2 { + t.Fatalf("fallbacks = %+v", spec.Model.Fallbacks) + } + if spec.Model.Fallbacks[0].Name != "sonnet" || spec.Model.Fallbacks[0].Backend != BackendClaudeAgent { + t.Errorf("fb0 = %+v", spec.Model.Fallbacks[0]) + } + // Candidates flattens primary + fallbacks in order. + cands := spec.Model.Candidates() + if len(cands) != 3 || cands[0].Name != "opus" || cands[1].Name != "sonnet" || cands[2].Name != "gpt-5.5" { + t.Errorf("candidates = %+v", cands) + } +} diff --git a/pkg/api/terminal_outcome.go b/pkg/api/terminal_outcome.go new file mode 100644 index 0000000..bdf6e5b --- /dev/null +++ b/pkg/api/terminal_outcome.go @@ -0,0 +1,63 @@ +package api + +import ( + "fmt" + "strings" +) + +// TerminalOutcomeKind identifies a native agent terminal result. +type TerminalOutcomeKind string + +const ( + TerminalOutcomePlan TerminalOutcomeKind = "plan" + TerminalOutcomeQuestions TerminalOutcomeKind = "questions" +) + +// TerminalPlan is the plan returned by a native planning tool. +type TerminalPlan struct { + Content string `json:"content"` + Path string `json:"path,omitempty"` +} + +// TerminalQuestion is one question returned by a native ask-user tool. +type TerminalQuestion struct { + Text string `json:"text"` + Context string `json:"context,omitempty"` + Options []string `json:"options,omitempty"` +} + +// TerminalOutcome carries native plan or question completion independently of +// schema-constrained StructuredData. +type TerminalOutcome struct { + Kind TerminalOutcomeKind `json:"kind"` + Plan *TerminalPlan `json:"plan,omitempty"` + Questions []TerminalQuestion `json:"questions,omitempty"` +} + +// Validate rejects incomplete or mixed terminal payloads. +func (o TerminalOutcome) Validate() error { + switch o.Kind { + case TerminalOutcomePlan: + if o.Plan == nil || strings.TrimSpace(o.Plan.Content) == "" { + return fmt.Errorf("terminal plan content is required") + } + if len(o.Questions) > 0 { + return fmt.Errorf("plan outcome must not carry questions") + } + case TerminalOutcomeQuestions: + if o.Plan != nil { + return fmt.Errorf("questions outcome must not carry a plan") + } + if len(o.Questions) == 0 { + return fmt.Errorf("terminal questions are required") + } + for i, question := range o.Questions { + if strings.TrimSpace(question.Text) == "" { + return fmt.Errorf("terminal question %d text is required", i+1) + } + } + default: + return fmt.Errorf("unknown terminal outcome kind %q", o.Kind) + } + return nil +} diff --git a/pkg/api/terminal_outcome_test.go b/pkg/api/terminal_outcome_test.go new file mode 100644 index 0000000..5c616ee --- /dev/null +++ b/pkg/api/terminal_outcome_test.go @@ -0,0 +1,55 @@ +package api + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestTerminalOutcomeValidate(t *testing.T) { + tests := []struct { + name string + outcome TerminalOutcome + wantErr string + }{ + { + name: "plan", + outcome: TerminalOutcome{ + Kind: TerminalOutcomePlan, + Plan: &TerminalPlan{Content: "1. Inspect\n2. Change"}, + }, + }, + { + name: "questions", + outcome: TerminalOutcome{ + Kind: TerminalOutcomeQuestions, + Questions: []TerminalQuestion{{Text: "Which database?", Options: []string{"PostgreSQL", "SQLite"}}}, + }, + }, + { + name: "plan requires content", + outcome: TerminalOutcome{Kind: TerminalOutcomePlan, Plan: &TerminalPlan{}}, + wantErr: "plan content is required", + }, + { + name: "questions reject plan payload", + outcome: TerminalOutcome{ + Kind: TerminalOutcomeQuestions, + Plan: &TerminalPlan{Content: "not a question"}, + Questions: []TerminalQuestion{{Text: "Continue?"}}, + }, + wantErr: "questions outcome must not carry a plan", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.outcome.Validate() + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.ErrorContains(t, err, tt.wantErr) + }) + } +} diff --git a/pkg/api/tool_approval.go b/pkg/api/tool_approval.go new file mode 100644 index 0000000..9b9a383 --- /dev/null +++ b/pkg/api/tool_approval.go @@ -0,0 +1,252 @@ +package api + +import ( + "encoding/json" + "fmt" + "reflect" + "strings" +) + +// ToolApprovalAction selects how a suspended tool call is resolved. +type ToolApprovalAction string + +const ( + ToolApprovalApprove ToolApprovalAction = "approve" + ToolApprovalDeny ToolApprovalAction = "deny" + ToolApprovalRespond ToolApprovalAction = "respond" +) + +// ToolApprovalRequest is the serializable identity and input of a tool call +// that may require a later approval decision. +type ToolApprovalRequest struct { + ToolCallID string `json:"toolCallId" yaml:"toolCallId"` + Tool string `json:"tool" yaml:"tool"` + Input json.RawMessage `json:"input,omitempty" yaml:"input,omitempty"` +} + +// ToolApprovalCall records one tool request from the interrupted model turn. +// Result is set when the tool completed before a sibling call suspended the +// turn; such calls must not be executed again when the turn resumes. +type ToolApprovalCall struct { + Request ToolApprovalRequest `json:"request" yaml:"request"` + Result *ToolResult `json:"result,omitempty" yaml:"result,omitempty"` +} + +// ToolApprovalState is the durable state returned when a model turn suspends. +// Messages is the complete provider-neutral conversation ending with the +// assistant tool requests; Calls records which requests are pending or done. +type ToolApprovalState struct { + Messages []Message `json:"messages" yaml:"messages"` + Calls []ToolApprovalCall `json:"calls" yaml:"calls"` +} + +// ToolApprovalDecision resolves one pending call. Approve may replace Input; +// Deny may carry a Message; Respond supplies an already-computed Result. +type ToolApprovalDecision struct { + ToolCallID string `json:"toolCallId" yaml:"toolCallId"` + Tool string `json:"tool" yaml:"tool"` + Action ToolApprovalAction `json:"action" yaml:"action"` + Input json.RawMessage `json:"input,omitempty" yaml:"input,omitempty"` + Message string `json:"message,omitempty" yaml:"message,omitempty"` + Result *ToolResult `json:"result,omitempty" yaml:"result,omitempty"` +} + +// ToolApprovalResume carries durable suspension state and exactly one decision +// for every pending call into a later request. +type ToolApprovalResume struct { + State ToolApprovalState `json:"state" yaml:"state"` + Decisions []ToolApprovalDecision `json:"decisions" yaml:"decisions"` +} + +func (r ToolApprovalRequest) Validate() error { + if strings.TrimSpace(r.ToolCallID) == "" { + return fmt.Errorf("tool call ID is required") + } + if strings.TrimSpace(r.Tool) == "" { + return fmt.Errorf("tool name is required for call %q", r.ToolCallID) + } + if len(r.Input) > 0 && !json.Valid(r.Input) { + return fmt.Errorf("tool call %q input must be valid JSON", r.ToolCallID) + } + if err := validateApprovalInput(r.Input); err != nil { + return fmt.Errorf("tool call %q input: %w", r.ToolCallID, err) + } + return nil +} + +// Pending returns the unresolved calls in their original model-request order. +func (s ToolApprovalState) Pending() []ToolApprovalRequest { + pending := make([]ToolApprovalRequest, 0, len(s.Calls)) + for _, call := range s.Calls { + if call.Result == nil { + pending = append(pending, call.Request) + } + } + return pending +} + +func (s ToolApprovalState) Validate() error { + if err := ValidateMessages(s.Messages); err != nil { + return fmt.Errorf("approval messages: %w", err) + } + if len(s.Calls) == 0 { + return fmt.Errorf("approval state must contain at least one tool call") + } + requests, err := approvalMessageRequests(s.Messages) + if err != nil { + return err + } + seen := make(map[string]bool, len(s.Calls)) + pending := 0 + for i, call := range s.Calls { + if err := call.Request.Validate(); err != nil { + return fmt.Errorf("approval call %d: %w", i+1, err) + } + id := call.Request.ToolCallID + if seen[id] { + return fmt.Errorf("duplicate approval call %q", id) + } + seen[id] = true + messageRequest, ok := requests[id] + if !ok { + return fmt.Errorf("approval call %q is absent from the final assistant message", id) + } + if messageRequest.Name != call.Request.Tool { + return fmt.Errorf("approval call %q tool %q does not match message tool %q", id, call.Request.Tool, messageRequest.Name) + } + if !equalJSON(messageRequest.Input, call.Request.Input) { + return fmt.Errorf("approval call %q input does not match the final assistant message", id) + } + if call.Result == nil { + pending++ + continue + } + if err := validateApprovalResult(id, call.Result); err != nil { + return err + } + } + if len(seen) != len(requests) { + return fmt.Errorf("approval state must account for every tool request in the final assistant message") + } + if pending == 0 { + return fmt.Errorf("approval state has no pending tool calls") + } + return nil +} + +func (r ToolApprovalResume) Validate() error { + if err := r.State.Validate(); err != nil { + return err + } + calls := make(map[string]ToolApprovalCall, len(r.State.Calls)) + for _, call := range r.State.Calls { + calls[call.Request.ToolCallID] = call + } + seen := make(map[string]bool, len(r.Decisions)) + for i, decision := range r.Decisions { + if seen[decision.ToolCallID] { + return fmt.Errorf("duplicate decision for tool call %q", decision.ToolCallID) + } + seen[decision.ToolCallID] = true + call, ok := calls[decision.ToolCallID] + if !ok { + return fmt.Errorf("decision references unknown tool call %q", decision.ToolCallID) + } + if call.Result != nil { + return fmt.Errorf("tool call %q is already resolved", decision.ToolCallID) + } + if decision.Tool != call.Request.Tool { + return fmt.Errorf("decision tool %q does not match pending tool %q", decision.Tool, call.Request.Tool) + } + if err := decision.validatePayload(); err != nil { + return fmt.Errorf("decision %d for tool call %q: %w", i+1, decision.ToolCallID, err) + } + } + for _, call := range r.State.Calls { + if call.Result == nil && !seen[call.Request.ToolCallID] { + return fmt.Errorf("missing decision for pending tool call %q", call.Request.ToolCallID) + } + } + return nil +} + +func (d ToolApprovalDecision) validatePayload() error { + switch d.Action { + case ToolApprovalApprove: + if d.Message != "" || d.Result != nil { + return fmt.Errorf("approve decision can only replace tool input") + } + if len(d.Input) > 0 && !json.Valid(d.Input) { + return fmt.Errorf("approve decision input must be valid JSON") + } + if err := validateApprovalInput(d.Input); err != nil { + return fmt.Errorf("approve decision input: %w", err) + } + case ToolApprovalDeny: + if len(d.Input) > 0 || d.Result != nil { + return fmt.Errorf("deny decision can only carry a message") + } + case ToolApprovalRespond: + if len(d.Input) > 0 { + return fmt.Errorf("respond decision cannot replace tool input") + } + if d.Message != "" { + return fmt.Errorf("respond decision cannot carry a denial message") + } + if d.Result == nil { + return fmt.Errorf("respond decision requires a tool result") + } + if err := validateApprovalResult(d.ToolCallID, d.Result); err != nil { + return err + } + default: + return fmt.Errorf("invalid approval action %q (valid: approve, deny, respond)", d.Action) + } + return nil +} + +func approvalMessageRequests(messages []Message) (map[string]ToolRequest, error) { + last := messages[len(messages)-1] + if last.Role != RoleAssistant { + return nil, fmt.Errorf("approval messages must end with an assistant tool request") + } + requests := make(map[string]ToolRequest) + for _, part := range last.Parts { + if part.Type == PartToolRequest && part.ToolRequest != nil { + requests[part.ToolRequest.ToolCallID] = *part.ToolRequest + } + } + if len(requests) == 0 { + return nil, fmt.Errorf("approval messages must end with an assistant tool request") + } + return requests, nil +} + +func validateApprovalResult(callID string, result *ToolResult) error { + if result.ToolCallID != callID { + return fmt.Errorf("tool result call ID %q does not match approval call %q", result.ToolCallID, callID) + } + return (Part{Type: PartToolResult, ToolResult: result}).Validate() +} + +func equalJSON(left, right json.RawMessage) bool { + if len(left) == 0 || len(right) == 0 { + return len(left) == len(right) + } + var a, b any + return json.Unmarshal(left, &a) == nil && json.Unmarshal(right, &b) == nil && reflect.DeepEqual(a, b) +} + +func validateApprovalInput(input json.RawMessage) error { + if len(input) == 0 { + return nil + } + var value map[string]any + if err := json.Unmarshal(input, &value); err != nil { + return fmt.Errorf("must be a JSON object") + } + if value == nil { + return fmt.Errorf("must be a JSON object") + } + return nil +} diff --git a/pkg/api/tool_approval_ginkgo_test.go b/pkg/api/tool_approval_ginkgo_test.go new file mode 100644 index 0000000..36474e5 --- /dev/null +++ b/pkg/api/tool_approval_ginkgo_test.go @@ -0,0 +1,111 @@ +package api_test + +import ( + "encoding/json" + + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Resumable tool approval", func() { + validState := func() api.ToolApprovalState { + return api.ToolApprovalState{ + Messages: []api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Update and inspect the invoice."}}}, + {Role: api.RoleAssistant, Parts: []api.Part{ + {Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ToolCallID: "call-update", Name: "invoice_update", Input: json.RawMessage(`{"amount":10}`)}}, + {Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ToolCallID: "call-read", Name: "invoice_get", Input: json.RawMessage(`{"id":"inv-1"}`)}}, + }}, + }, + Calls: []api.ToolApprovalCall{ + {Request: api.ToolApprovalRequest{ToolCallID: "call-update", Tool: "invoice_update", Input: json.RawMessage(`{"amount":10}`)}}, + { + Request: api.ToolApprovalRequest{ToolCallID: "call-read", Tool: "invoice_get", Input: json.RawMessage(`{"id":"inv-1"}`)}, + Result: &api.ToolResult{ToolCallID: "call-read", Output: json.RawMessage(`{"amount":10}`)}, + }, + }, + } + } + + It("round-trips pending and already-completed calls", func() { + state := validState() + data, err := json.Marshal(state) + Expect(err).NotTo(HaveOccurred()) + + var decoded api.ToolApprovalState + Expect(json.Unmarshal(data, &decoded)).To(Succeed()) + Expect(decoded.Validate()).To(Succeed()) + Expect(decoded.Calls).To(Equal(state.Calls)) + Expect(decoded.Pending()).To(Equal([]api.ToolApprovalRequest{state.Calls[0].Request})) + }) + + It("accepts one approve decision with edited input for each pending call", func() { + resume := api.ToolApprovalResume{ + State: validState(), + Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-update", + Tool: "invoice_update", + Action: api.ToolApprovalApprove, + Input: json.RawMessage(`{"amount":12}`), + }}, + } + Expect(resume.Validate()).To(Succeed()) + + spec := api.Spec{Model: api.Model{Name: "gpt-5", Backend: api.BackendOpenAI}, ToolApproval: &resume} + Expect(spec.Validate()).To(Succeed()) + }) + + It("compares tool inputs structurally regardless of object key order", func() { + state := validState() + state.Messages[1].Parts[0].ToolRequest.Input = json.RawMessage(`{"currency":"USD","amount":10}`) + state.Calls[0].Request.Input = json.RawMessage(`{"amount":10,"currency":"USD"}`) + Expect(state.Validate()).To(Succeed()) + + state.Calls[0].Request.Input = json.RawMessage(`{"amount":11,"currency":"USD"}`) + Expect(state.Validate()).To(MatchError(ContainSubstring("input does not match"))) + }) + + DescribeTable("rejects invalid continuation decisions", + func(mutate func(*api.ToolApprovalResume), want string) { + resume := api.ToolApprovalResume{ + State: validState(), + Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-update", Tool: "invoice_update", Action: api.ToolApprovalApprove, + }}, + } + mutate(&resume) + Expect(resume.Validate()).To(MatchError(ContainSubstring(want))) + }, + Entry("missing", func(resume *api.ToolApprovalResume) { resume.Decisions = nil }, `missing decision for pending tool call "call-update"`), + Entry("duplicate", func(resume *api.ToolApprovalResume) { resume.Decisions = append(resume.Decisions, resume.Decisions[0]) }, `duplicate decision for tool call "call-update"`), + Entry("unknown call", func(resume *api.ToolApprovalResume) { resume.Decisions[0].ToolCallID = "call-other" }, `decision references unknown tool call "call-other"`), + Entry("mismatched tool", func(resume *api.ToolApprovalResume) { resume.Decisions[0].Tool = "invoice_delete" }, `decision tool "invoice_delete" does not match pending tool "invoice_update"`), + Entry("completed replay", func(resume *api.ToolApprovalResume) { + resume.Decisions[0] = api.ToolApprovalDecision{ToolCallID: "call-read", Tool: "invoice_get", Action: api.ToolApprovalApprove} + }, `tool call "call-read" is already resolved`), + ) + + It("requires deny and respond decisions to carry only their own payload", func() { + state := validState() + deny := api.ToolApprovalResume{State: state, Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-update", Tool: "invoice_update", Action: api.ToolApprovalDeny, Message: "not now", + }}} + Expect(deny.Validate()).To(Succeed()) + + respond := api.ToolApprovalResume{State: state, Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-update", Tool: "invoice_update", Action: api.ToolApprovalRespond, + Result: &api.ToolResult{ToolCallID: "call-update", Output: json.RawMessage(`{"queued":true}`)}, + }}} + Expect(respond.Validate()).To(Succeed()) + + respond.Decisions[0].Input = json.RawMessage(`{"amount":12}`) + Expect(respond.Validate()).To(MatchError(ContainSubstring("respond decision cannot replace tool input"))) + + approve := api.ToolApprovalResume{State: state, Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-update", Tool: "invoice_update", Action: api.ToolApprovalApprove, Input: json.RawMessage(`[12]`), + }}} + Expect(approve.Validate()).To(MatchError(ContainSubstring("approve decision input: must be a JSON object"))) + }) +}) diff --git a/pkg/api/tool_preferences_ginkgo_test.go b/pkg/api/tool_preferences_ginkgo_test.go new file mode 100644 index 0000000..67805b8 --- /dev/null +++ b/pkg/api/tool_preferences_ginkgo_test.go @@ -0,0 +1,72 @@ +package api_test + +import ( + "encoding/json" + + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Tool preferences", func() { + It("survives a Spec JSON round trip", func() { + in := api.Spec{ + Model: api.Model{Name: "claude-sonnet-4-6", Backend: api.BackendAnthropic}, + Prompt: api.Prompt{User: "inspect invoices"}, + ToolPreferences: api.ToolPreferences{"billing": api.ToolModeAsk, "invoice_delete": api.ToolModeOff}, + } + + encoded, err := json.Marshal(in) + Expect(err).NotTo(HaveOccurred()) + var wire map[string]any + Expect(json.Unmarshal(encoded, &wire)).To(Succeed()) + Expect(wire["toolPreferences"]).To(Equal(map[string]any{ + "billing": "ask", + "invoice_delete": "off", + })) + + var out api.Spec + Expect(json.Unmarshal(encoded, &out)).To(Succeed()) + Expect(out.ToolPreferences).To(Equal(in.ToolPreferences)) + }) + + It("replaces base preferences when an override supplies per-turn preferences", func() { + base := api.Spec{ToolPreferences: api.ToolPreferences{ + "billing": api.ToolModeAsk, + "search": api.ToolModeOff, + }} + override := api.Spec{ToolPreferences: api.ToolPreferences{ + "billing": api.ToolModeOn, + }} + + Expect(base.Merge(override).ToolPreferences).To(Equal(api.ToolPreferences{ + "billing": api.ToolModeOn, + })) + Expect(base.Merge(api.Spec{}).ToolPreferences).To(Equal(base.ToolPreferences)) + }) + + It("rejects an unknown preference before provider execution", func() { + spec := api.Spec{ + Model: api.Model{Name: "claude-sonnet-4-6", Backend: api.BackendAnthropic}, + Prompt: api.Prompt{User: "inspect invoices"}, + ToolPreferences: api.ToolPreferences{"billing": "sometimes"}, + } + + Expect(spec.Validate()).To(MatchError(ContainSubstring(`invalid tool preference "sometimes" for "billing"`))) + }) + + It("rejects removed enabled and disabled labels", func() { + for _, mode := range []api.ToolMode{"enabled", "disabled"} { + spec := api.Spec{ + Model: api.Model{Name: "claude-sonnet-4-6", Backend: api.BackendAnthropic}, + Prompt: api.Prompt{User: "inspect invoices"}, + ToolPreferences: api.ToolPreferences{"billing": mode}, + } + Expect(spec.Validate()).To(MatchError(ContainSubstring(`invalid tool preference`))) + + permissions := api.Permissions{Tools: api.Tools{Modes: map[string]api.ToolMode{"billing": mode}}} + Expect(permissions.Validate()).To(MatchError(ContainSubstring(`invalid tool mode`))) + } + }) +}) diff --git a/pkg/api/tooldef.go b/pkg/api/tooldef.go new file mode 100644 index 0000000..717a9a0 --- /dev/null +++ b/pkg/api/tooldef.go @@ -0,0 +1,72 @@ +package api + +import ( + "context" + "fmt" +) + +// ToolPreferences selects the effective per-turn mode for a tool name or group. +// An exact tool-name entry takes precedence over its group entry. +type ToolPreferences map[string]ToolMode + +// Validate rejects unknown modes before a provider request is assembled. +func (p ToolPreferences) Validate() error { + if _, exists := p[""]; exists { + return fmt.Errorf("tool preference key cannot be empty") + } + for _, key := range sortedKeys(p) { + mode := p[key] + if _, ok := NormalizeToolMode(mode); !ok { + return fmt.Errorf("invalid tool preference %q for %q (valid: on, ask, off, auto)", mode, key) + } + } + return nil +} + +// ToolHandler runs a caller-supplied tool. input is the model's tool-call +// arguments (the decoded JSON object); the returned value is marshaled back to +// the model as the tool result. Returning an error surfaces it to the model as +// a failed tool call. +type ToolHandler func(ctx context.Context, input map[string]any) (any, error) + +// ToolDefinition is a caller-supplied tool that a tool-capable provider (see +// ToolCapableProvider) exposes to the model and executes in-process. It carries +// a Go handler, so — like CanUseTool — it is a runtime concern that lives on +// Config, never on the serializable Spec, and is never marshaled. +type ToolDefinition struct { + // Name is the tool id the model calls (provider-safe: letters, digits, _-). + Name string + // Description tells the model when/how to use the tool. + Description string + // InputSchema is the JSON Schema (decoded to a map) for the tool arguments. + // Nil means a no-argument tool. + InputSchema map[string]any + // Group is the preference key shared by related tools. A tool-name preference + // overrides a group preference. + Group string + // Parent and Icon retain presentation metadata for catalogs without affecting + // provider execution. + Parent string + Icon string + // Strict opts this tool into provider strict-schema enforcement. Safety hints + // prioritize tools when a provider caps strict-tool definitions. + Strict *bool + ReadOnlyHint *bool + DestructiveHint *bool + IdempotentHint *bool + // Handler executes the tool in-process. Required. + Handler ToolHandler `json:"-"` + // DefaultPermission controls exposure: off omits the tool, ask routes calls + // through Config.CanUseTool, on auto-runs, and auto defers to runtime policy. + DefaultPermission ToolMode + // Annotations carries opaque caller metadata (e.g. the originating CLI + // verb/method/path) for policies that want the raw values; providers ignore it. + Annotations map[string]string `json:",omitempty"` +} + +// NeedsApproval reports whether a call to this tool must go through +// Config.CanUseTool before running. +func (t ToolDefinition) NeedsApproval() bool { + mode, ok := NormalizeToolMode(t.DefaultPermission) + return ok && mode == ToolModeAsk +} diff --git a/pkg/api/workflow.go b/pkg/api/workflow.go index 8739462..2013611 100644 --- a/pkg/api/workflow.go +++ b/pkg/api/workflow.go @@ -1,25 +1,25 @@ package api -import ( - "encoding/json" - "fmt" -) +import "fmt" // Workflow declares the generate→verify loop around a run as hook // declarations: an optional verification stage (commands run after each -// generation, whose failure feedback drives a re-run), an optional postRun -// stage (commit the result), and an optional typed output schema. It is the -// serializable form of pkg/ai/agent's Verify/PostRun/Output hooks, and mirrors -// clicky-ui's AISpecRuntimeLocalWorkflow so the SpecRuntimeEditor "Verify" -// section round-trips. +// generation, whose failure feedback drives a re-run), and an optional postRun +// stage (commit the result). It is the serializable form of pkg/ai/agent's +// Verify/PostRun hooks, and mirrors clicky-ui's AISpecRuntimeLocalWorkflow so +// the SpecRuntimeEditor "Verify" section round-trips. // // A spec with a Verify but an empty Prompt.User runs verify-only (generation is // skipped); a spec with no Verify runs generate-only (today's behaviour). type Workflow struct { Verify *Verify `json:"verify,omitempty" yaml:"verify,omitempty"` PostRun *PostRun `json:"postRun,omitempty" yaml:"postRun,omitempty"` - // Output declares the workflow's typed final-result schema. - Output *Output `json:"output,omitempty" yaml:"output,omitempty"` + + // AutoVerifyWithoutFixture is the explicit policy opt-in for hosts that + // project a successful generate-only run into a durable verified state. A + // false value keeps the durable work item open when no verification fixture + // ran; success by itself is not treated as proof of correctness. + AutoVerifyWithoutFixture bool `json:"autoVerifyWithoutFixture,omitempty" yaml:"autoVerifyWithoutFixture,omitempty"` } // Verify is the loop's definition-of-done: it runs after each generation and @@ -50,12 +50,6 @@ type PostRun struct { KeepWorktree bool `json:"keepWorktree,omitempty" yaml:"keepWorktree,omitempty"` } -// Output declares the workflow's typed final-result schema (a JSON Schema -// document) so callers/editors know the shape of the run's structured result. -type Output struct { - SchemaJSON json.RawMessage `json:"schemaJSON,omitempty" yaml:"schemaJSON,omitempty"` -} - // Validate checks the workflow's enum-typed fields. func (w *Workflow) Validate() error { if w == nil || w.Verify == nil { diff --git a/pkg/api/workflow_test.go b/pkg/api/workflow_test.go index 7af0bb8..ea9ca67 100644 --- a/pkg/api/workflow_test.go +++ b/pkg/api/workflow_test.go @@ -17,8 +17,8 @@ func TestWorkflowSpecRoundTrip(t *testing.T) { Scope: VerifyScopeChanged, MaxIterations: 3, }, - PostRun: &PostRun{Commit: true, CommitMessage: "apply"}, - Output: &Output{SchemaJSON: json.RawMessage(`{"type":"object"}`)}, + PostRun: &PostRun{Commit: true, CommitMessage: "apply"}, + AutoVerifyWithoutFixture: true, }, } @@ -43,8 +43,8 @@ func TestWorkflowSpecRoundTrip(t *testing.T) { if got.Workflow.PostRun == nil || !got.Workflow.PostRun.Commit { t.Errorf("postRun not preserved: %+v", got.Workflow.PostRun) } - if got.Workflow.Output == nil || string(got.Workflow.Output.SchemaJSON) != string(spec.Workflow.Output.SchemaJSON) { - t.Errorf("output schema not preserved: %+v", got.Workflow.Output) + if !got.Workflow.AutoVerifyWithoutFixture { + t.Errorf("autoVerifyWithoutFixture not preserved: %+v", got.Workflow) } } @@ -63,11 +63,32 @@ func TestSpecSchemaIncludesWorkflow(t *testing.T) { if err != nil { t.Fatalf("schema: %v", err) } - for _, want := range []string{"workflow", "Workflow", "Verify", "commands", "maxIterations"} { + for _, want := range []string{"workflow", "Workflow", "Verify", "commands", "maxIterations", "autoVerifyWithoutFixture"} { if !strings.Contains(string(data), want) { t.Errorf("reflected spec schema missing %q", want) } } + var reflected map[string]any + if err := json.Unmarshal(data, &reflected); err != nil { + t.Fatalf("decode schema: %v", err) + } + definitions, ok := reflected["$defs"].(map[string]any) + if !ok { + t.Fatalf("reflected spec schema has no $defs object") + } + workflow, ok := definitions["Workflow"].(map[string]any) + if !ok { + t.Fatalf("reflected spec schema has no Workflow definition") + } + properties, ok := workflow["properties"].(map[string]any) + if !ok { + t.Fatalf("reflected Workflow schema has no properties object") + } + for _, gone := range []string{"output", "Output"} { + if _, exists := properties[gone]; exists { + t.Errorf("reflected Workflow schema still contains removed field %q", gone) + } + } } func TestVerifyScopeValidate(t *testing.T) { diff --git a/pkg/attachments/attachments_suite_test.go b/pkg/attachments/attachments_suite_test.go new file mode 100644 index 0000000..5821f2f --- /dev/null +++ b/pkg/attachments/attachments_suite_test.go @@ -0,0 +1,13 @@ +package attachments_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAttachments(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Attachments Suite") +} diff --git a/pkg/attachments/gc.go b/pkg/attachments/gc.go new file mode 100644 index 0000000..1bc049c --- /dev/null +++ b/pkg/attachments/gc.go @@ -0,0 +1,68 @@ +package attachments + +import ( + "fmt" + "io/fs" + "os" + "path/filepath" + "sort" + "strings" + "time" + + "github.com/flanksource/captain/pkg/api" +) + +type GCResult struct { + RemovedIDs []string `json:"removedIds" pretty:"label=Attachments"` + RemovedBytes int64 `json:"removedBytes" pretty:"label=Bytes"` + DryRun bool `json:"dryRun" pretty:"label=Dry Run"` + RetentionDays int `json:"retentionDays" pretty:"label=Retention Days"` +} + +func (s *Store) GC(referenced map[string]struct{}, retention time.Duration, dryRun bool) (GCResult, error) { + if retention <= 0 { + return GCResult{}, fmt.Errorf("attachment retention must be positive") + } + result := GCResult{DryRun: dryRun, RetentionDays: int(retention.Hours() / 24)} + cutoff := time.Now().Add(-retention) + root := filepath.Join(s.directory, "sha256") + err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + if os.IsNotExist(walkErr) { + return nil + } + return walkErr + } + if entry.IsDir() || strings.HasPrefix(entry.Name(), ".attachment-") { + return nil + } + id := api.AttachmentIDPrefix + entry.Name() + if _, ok := referenced[id]; ok { + return nil + } + ref := api.AttachmentRef{ID: id} + if ref.Validate() != nil { + return nil + } + info, err := entry.Info() + if err != nil { + return err + } + if !info.ModTime().Before(cutoff) { + return nil + } + result.RemovedIDs = append(result.RemovedIDs, id) + result.RemovedBytes += info.Size() + if !dryRun { + if err := os.Remove(path); err != nil { + return fmt.Errorf("remove attachment %s: %w", id, err) + } + } + return nil + }) + if err != nil && !os.IsNotExist(err) { + return GCResult{}, fmt.Errorf("scan attachment store: %w", err) + } + sort.Strings(result.RemovedIDs) + return result, nil +} diff --git a/pkg/attachments/limits.go b/pkg/attachments/limits.go new file mode 100644 index 0000000..ec7c2c0 --- /dev/null +++ b/pkg/attachments/limits.go @@ -0,0 +1,48 @@ +package attachments + +const ( + DefaultMaxFileBytes int64 = 20 << 20 + DefaultMaxRequestBytes int64 = 50 << 20 + DefaultMaxFiles = 10 +) + +type Limits struct { + MaxFileBytes int64 + MaxRequestBytes int64 + MaxFiles int +} + +func DefaultLimits() Limits { + return Limits{ + MaxFileBytes: DefaultMaxFileBytes, + MaxRequestBytes: DefaultMaxRequestBytes, + MaxFiles: DefaultMaxFiles, + } +} + +func (l Limits) withDefaults() Limits { + defaults := DefaultLimits() + if l.MaxFileBytes == 0 { + l.MaxFileBytes = defaults.MaxFileBytes + } + if l.MaxRequestBytes == 0 { + l.MaxRequestBytes = defaults.MaxRequestBytes + } + if l.MaxFiles == 0 { + l.MaxFiles = defaults.MaxFiles + } + return l +} + +func (l Limits) validate() error { + if l.MaxFileBytes < 1 { + return invalidLimit("max file bytes", l.MaxFileBytes) + } + if l.MaxRequestBytes < 1 { + return invalidLimit("max request bytes", l.MaxRequestBytes) + } + if l.MaxFiles < 1 { + return invalidLimit("max files", l.MaxFiles) + } + return nil +} diff --git a/pkg/attachments/resolve.go b/pkg/attachments/resolve.go new file mode 100644 index 0000000..0640dcf --- /dev/null +++ b/pkg/attachments/resolve.go @@ -0,0 +1,127 @@ +package attachments + +import ( + "context" + "fmt" + "mime" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + + "github.com/flanksource/captain/pkg/api" +) + +func (s *Store) Resolve(ctx context.Context, refs []api.AttachmentRef, baseDir string) ([]api.AttachmentRef, error) { + if len(refs) > s.limits.MaxFiles { + return nil, fmt.Errorf("attachment request has %d files and exceeds %d file limit", len(refs), s.limits.MaxFiles) + } + resolved := make([]api.AttachmentRef, 0, len(refs)) + var total int64 + for i, ref := range refs { + if err := ref.Validate(); err != nil { + return nil, fmt.Errorf("attachment %d: %w", i+1, err) + } + prepared, err := s.resolveOne(ctx, ref, baseDir) + if err != nil { + return nil, fmt.Errorf("attachment %d: %w", i+1, err) + } + total += prepared.Size + if total > s.limits.MaxRequestBytes { + return nil, fmt.Errorf("attachments total %d bytes exceeds %d byte request limit", total, s.limits.MaxRequestBytes) + } + resolved = append(resolved, prepared) + } + return resolved, nil +} + +func (s *Store) resolveOne(ctx context.Context, ref api.AttachmentRef, baseDir string) (api.AttachmentRef, error) { + content, filename, err := s.readSource(ctx, ref, baseDir) + if err != nil { + return api.AttachmentRef{}, err + } + mediaType := detectMediaType(content, filename) + if ref.MediaType != "" && canonicalMediaType(ref.MediaType) != mediaType { + return api.AttachmentRef{}, fmt.Errorf("declared media type %s does not match detected %s", ref.MediaType, mediaType) + } + id, path, err := s.persist(content) + if err != nil { + return api.AttachmentRef{}, err + } + digest := strings.TrimPrefix(id, api.AttachmentIDPrefix) + if ref.SHA256 != "" && !strings.EqualFold(ref.SHA256, digest) { + return api.AttachmentRef{}, fmt.Errorf("declared sha256 %s does not match content %s", ref.SHA256, digest) + } + if ref.Filename != "" { + filename = ref.Filename + } + return api.AttachmentRef{ + ID: id, + Filename: filename, + MediaType: mediaType, + Size: int64(len(content)), + SHA256: digest, + }.WithPreparedContent(api.AttachmentContent{Bytes: content, Path: path}), nil +} + +func (s *Store) readSource(ctx context.Context, ref api.AttachmentRef, baseDir string) ([]byte, string, error) { + switch { + case ref.ID != "": + file, err := s.Open(ref.ID) + if err != nil { + return nil, "", err + } + defer file.Close() + content, err := readLimited(file, s.limits.MaxFileBytes) + return content, ref.Filename, err + case ref.Path != "": + path := ref.Path + if !filepath.IsAbs(path) { + path = filepath.Join(baseDir, path) + } + file, err := os.Open(path) + if err != nil { + return nil, "", fmt.Errorf("open attachment path %s: %w", path, err) + } + defer file.Close() + content, err := readLimited(file, s.limits.MaxFileBytes) + return content, filepath.Base(path), err + case ref.URL != "": + request, err := http.NewRequestWithContext(ctx, http.MethodGet, ref.URL, nil) + if err != nil { + return nil, "", fmt.Errorf("create attachment request: %w", err) + } + response, err := s.httpClient.Do(request) + if err != nil { + return nil, "", fmt.Errorf("download attachment: %w", err) + } + defer response.Body.Close() + if response.StatusCode < 200 || response.StatusCode >= 300 { + return nil, "", fmt.Errorf("download attachment: unexpected HTTP status %s", response.Status) + } + content, err := readLimited(response.Body, s.limits.MaxFileBytes) + parsed, _ := url.Parse(ref.URL) + return content, filepath.Base(parsed.Path), err + default: + return nil, "", fmt.Errorf("attachment source is required") + } +} + +func detectMediaType(content []byte, filename string) string { + mediaType := canonicalMediaType(http.DetectContentType(content)) + if mediaType == "application/octet-stream" { + if extensionType := canonicalMediaType(mime.TypeByExtension(filepath.Ext(filename))); extensionType != "" { + return extensionType + } + } + return mediaType +} + +func canonicalMediaType(value string) string { + mediaType, _, err := mime.ParseMediaType(value) + if err != nil { + return strings.ToLower(strings.TrimSpace(value)) + } + return strings.ToLower(mediaType) +} diff --git a/pkg/attachments/store.go b/pkg/attachments/store.go new file mode 100644 index 0000000..0d8d76a --- /dev/null +++ b/pkg/attachments/store.go @@ -0,0 +1,171 @@ +package attachments + +import ( + "crypto/sha256" + "errors" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" + + "github.com/flanksource/captain/pkg/api" +) + +type StoreOptions struct { + Directory string + Limits Limits + HTTPClient *http.Client +} + +type Store struct { + directory string + limits Limits + httpClient *http.Client +} + +func NewStore(opts StoreOptions) (*Store, error) { + if strings.TrimSpace(opts.Directory) == "" { + return nil, errors.New("attachment store directory is required") + } + limits := opts.Limits.withDefaults() + if err := limits.validate(); err != nil { + return nil, err + } + directory, err := filepath.Abs(opts.Directory) + if err != nil { + return nil, fmt.Errorf("resolve attachment store directory: %w", err) + } + if err := ensurePrivateDirectory(directory); err != nil { + return nil, err + } + client := opts.HTTPClient + if client == nil { + client = &http.Client{ + Timeout: 30 * time.Second, + CheckRedirect: func(_ *http.Request, via []*http.Request) error { + if len(via) >= 5 { + return fmt.Errorf("attachment download exceeded 5 redirects") + } + return nil + }, + } + } + return &Store{directory: directory, limits: limits, httpClient: client}, nil +} + +func (s *Store) Limits() Limits { return s.limits } + +func (s *Store) Directory() string { return s.directory } + +func (s *Store) Put(reader io.Reader, filename, declaredMediaType string) (api.AttachmentRef, error) { + content, err := readLimited(reader, s.limits.MaxFileBytes) + if err != nil { + return api.AttachmentRef{}, err + } + mediaType := detectMediaType(content, filename) + if declaredMediaType != "" && canonicalMediaType(declaredMediaType) != "application/octet-stream" && canonicalMediaType(declaredMediaType) != mediaType { + return api.AttachmentRef{}, fmt.Errorf("declared media type %s does not match detected %s", declaredMediaType, mediaType) + } + id, path, err := s.persist(content) + if err != nil { + return api.AttachmentRef{}, err + } + digest := strings.TrimPrefix(id, api.AttachmentIDPrefix) + return api.AttachmentRef{ + ID: id, Filename: filename, MediaType: mediaType, + Size: int64(len(content)), SHA256: digest, + }.WithPreparedContent(api.AttachmentContent{Bytes: content, Path: path}), nil +} + +func (s *Store) Path(id string) string { + digest := strings.TrimPrefix(id, api.AttachmentIDPrefix) + if len(digest) < 2 { + return "" + } + return filepath.Join(s.directory, "sha256", digest[:2], digest) +} + +func (s *Store) Open(id string) (*os.File, error) { + ref := api.AttachmentRef{ID: id} + if err := ref.Validate(); err != nil { + return nil, err + } + root, err := os.OpenRoot(s.directory) + if err != nil { + return nil, fmt.Errorf("open attachment store: %w", err) + } + defer root.Close() + digest := strings.TrimPrefix(id, api.AttachmentIDPrefix) + file, err := root.Open(filepath.Join("sha256", digest[:2], digest)) + if err != nil { + return nil, fmt.Errorf("open attachment %s: %w", id, err) + } + return file, nil +} + +func (s *Store) persist(content []byte) (string, string, error) { + digest := fmt.Sprintf("%x", sha256.Sum256(content)) + id := api.AttachmentIDPrefix + digest + path := s.Path(id) + if info, err := os.Stat(path); err == nil { + if !info.Mode().IsRegular() { + return "", "", fmt.Errorf("attachment path %s is not a regular file", path) + } + return id, path, nil + } else if !errors.Is(err, os.ErrNotExist) { + return "", "", fmt.Errorf("inspect attachment path %s: %w", path, err) + } + dir := filepath.Dir(path) + if err := ensurePrivateDirectory(dir); err != nil { + return "", "", err + } + tmp, err := os.CreateTemp(dir, ".attachment-*") + if err != nil { + return "", "", fmt.Errorf("create attachment temp file: %w", err) + } + tmpName := tmp.Name() + defer os.Remove(tmpName) + if err := tmp.Chmod(0o600); err != nil { + _ = tmp.Close() + return "", "", fmt.Errorf("set attachment permissions: %w", err) + } + if _, err := tmp.Write(content); err != nil { + _ = tmp.Close() + return "", "", fmt.Errorf("write attachment: %w", err) + } + if err := tmp.Close(); err != nil { + return "", "", fmt.Errorf("close attachment: %w", err) + } + if err := os.Rename(tmpName, path); err != nil { + return "", "", fmt.Errorf("publish attachment: %w", err) + } + return id, path, nil +} + +func ensurePrivateDirectory(path string) error { + if err := os.MkdirAll(path, 0o700); err != nil { + return fmt.Errorf("create attachment directory %s: %w", path, err) + } + if err := os.Chmod(path, 0o700); err != nil { + return fmt.Errorf("set attachment directory permissions %s: %w", path, err) + } + return nil +} + +func readLimited(reader io.Reader, limit int64) ([]byte, error) { + content, err := io.ReadAll(io.LimitReader(reader, limit+1)) + if err != nil { + return nil, err + } + if int64(len(content)) > limit { + return nil, fmt.Errorf("attachment exceeds %d byte file limit", limit) + } + return content, nil +} + +func invalidLimit(name string, value any) error { + return fmt.Errorf("attachment %s must be positive, got %v", name, value) +} diff --git a/pkg/attachments/store_ginkgo_test.go b/pkg/attachments/store_ginkgo_test.go new file mode 100644 index 0000000..376b7bc --- /dev/null +++ b/pkg/attachments/store_ginkgo_test.go @@ -0,0 +1,122 @@ +package attachments_test + +import ( + "context" + "crypto/sha256" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/attachments" +) + +var _ = Describe("Store", func() { + It("resolves a local file into an immutable content-addressed attachment", func() { + root := GinkgoT().TempDir() + input := filepath.Join(root, "diagram.png") + content := append([]byte("\x89PNG\r\n\x1a\n"), make([]byte, 512)...) + Expect(os.WriteFile(input, content, 0o600)).To(Succeed()) + + store, err := attachments.NewStore(attachments.StoreOptions{Directory: filepath.Join(root, ".captain", "attachments")}) + Expect(err).NotTo(HaveOccurred()) + resolved, err := store.Resolve(context.Background(), []api.AttachmentRef{{Path: input}}, "") + Expect(err).NotTo(HaveOccurred()) + Expect(resolved).To(HaveLen(1)) + + digest := fmt.Sprintf("%x", sha256.Sum256(content)) + Expect(resolved[0].ID).To(Equal(api.AttachmentIDPrefix + digest)) + Expect(resolved[0].Filename).To(Equal("diagram.png")) + Expect(resolved[0].MediaType).To(Equal("image/png")) + Expect(resolved[0].Size).To(Equal(int64(len(content)))) + Expect(resolved[0].SHA256).To(Equal(digest)) + prepared, ok := resolved[0].PreparedContent() + Expect(ok).To(BeTrue()) + Expect(prepared.Bytes).To(Equal(content)) + Expect(prepared.Path).To(Equal(store.Path(resolved[0].ID))) + + info, err := os.Stat(store.Path(resolved[0].ID)) + Expect(err).NotTo(HaveOccurred()) + Expect(info.Mode().Perm()).To(Equal(os.FileMode(0o600))) + }) + + It("rejects declared media types that disagree with detected content", func() { + root := GinkgoT().TempDir() + input := filepath.Join(root, "invoice.pdf") + Expect(os.WriteFile(input, []byte("plain text, not a PDF"), 0o600)).To(Succeed()) + store, err := attachments.NewStore(attachments.StoreOptions{Directory: filepath.Join(root, "store")}) + Expect(err).NotTo(HaveOccurred()) + + _, err = store.Resolve(context.Background(), []api.AttachmentRef{{Path: input, MediaType: "application/pdf"}}, "") + Expect(err).To(MatchError(ContainSubstring("declared media type application/pdf does not match detected text/plain"))) + }) + + It("enforces file, request, and count limits before execution", func() { + root := GinkgoT().TempDir() + input := filepath.Join(root, "large.txt") + Expect(os.WriteFile(input, []byte("12345"), 0o600)).To(Succeed()) + store, err := attachments.NewStore(attachments.StoreOptions{ + Directory: filepath.Join(root, "store"), + Limits: attachments.Limits{ + MaxFileBytes: 4, + MaxRequestBytes: 8, + MaxFiles: 1, + }, + }) + Expect(err).NotTo(HaveOccurred()) + + _, err = store.Resolve(context.Background(), []api.AttachmentRef{{Path: input}}, "") + Expect(err).To(MatchError(ContainSubstring("exceeds 4 byte file limit"))) + _, err = store.Resolve(context.Background(), []api.AttachmentRef{{Path: input}, {Path: input}}, "") + Expect(err).To(MatchError(ContainSubstring("exceeds 1 file limit"))) + }) + + It("does not follow attachment symlinks outside the store", func() { + root := GinkgoT().TempDir() + store, err := attachments.NewStore(attachments.StoreOptions{Directory: filepath.Join(root, "store")}) + Expect(err).NotTo(HaveOccurred()) + + outside := filepath.Join(root, "outside.txt") + Expect(os.WriteFile(outside, []byte("secret"), 0o600)).To(Succeed()) + id := api.AttachmentIDPrefix + strings.Repeat("a", sha256.Size*2) + path := store.Path(id) + Expect(os.MkdirAll(filepath.Dir(path), 0o700)).To(Succeed()) + Expect(os.Symlink(outside, path)).To(Succeed()) + + file, err := store.Open(id) + if file != nil { + DeferCleanup(file.Close) + } + Expect(err).To(HaveOccurred()) + }) + + It("garbage-collects only old unreferenced blobs and supports dry-run", func() { + root := GinkgoT().TempDir() + store, err := attachments.NewStore(attachments.StoreOptions{Directory: filepath.Join(root, "store")}) + Expect(err).NotTo(HaveOccurred()) + kept, err := store.Put(strings.NewReader("keep"), "keep.txt", "text/plain") + Expect(err).NotTo(HaveOccurred()) + removed, err := store.Put(strings.NewReader("remove"), "remove.txt", "text/plain") + Expect(err).NotTo(HaveOccurred()) + old := time.Now().Add(-31 * 24 * time.Hour) + Expect(os.Chtimes(store.Path(kept.ID), old, old)).To(Succeed()) + Expect(os.Chtimes(store.Path(removed.ID), old, old)).To(Succeed()) + + dryRun, err := store.GC(map[string]struct{}{kept.ID: {}}, 30*24*time.Hour, true) + Expect(err).NotTo(HaveOccurred()) + Expect(dryRun.RemovedIDs).To(Equal([]string{removed.ID})) + _, err = os.Stat(store.Path(removed.ID)) + Expect(err).NotTo(HaveOccurred()) + + result, err := store.GC(map[string]struct{}{kept.ID: {}}, 30*24*time.Hour, false) + Expect(err).NotTo(HaveOccurred()) + Expect(result.RemovedIDs).To(Equal([]string{removed.ID})) + Expect(store.Path(kept.ID)).To(BeAnExistingFile()) + Expect(store.Path(removed.ID)).NotTo(BeAnExistingFile()) + }) +}) diff --git a/pkg/bash/category_config.yaml b/pkg/bash/category_config.yaml index 845bd8d..4abd804 100644 --- a/pkg/bash/category_config.yaml +++ b/pkg/bash/category_config.yaml @@ -306,6 +306,10 @@ categories: tools: - Task - TodoWrite + - TaskCreate + - TaskUpdate + - TaskGet + - TaskList - EnterPlanMode commands: - task @@ -323,11 +327,36 @@ categories: tools: - Edit - Write + - MultiEdit - NotebookEdit + # Codex file mutations. apply_patch is the wire name; CodexPatchApply is + # the normalized name history renders. Without these, every Codex file + # write classifies as `other`. + - apply_patch + - CodexPatchApply commands: - mkdir + - mv + - cp + - touch + - install + - patch + - tee patterns: - "sed.* -i" + # perl bundles its flags, so -i appears as -i, -pi, -ni or -i.bak. Matching + # a literal " -i" misses every bundled form. + - "^perl\\s+-[a-zA-Z.]*i" + - "^mv\\b" + - "^cp\\b" + - "^touch\\b" + - "^patch\\b" + - "^tee\\b" + # Deliberately no generic `> file` redirect rule. RE2 has no lookahead, so + # `> /dev/null` cannot be excluded cleanly, and `edit` also sources + # captain_artifacts — a false positive there invents a file-modification + # record. Rehydration from the session file recovers anything a miss drops, + # so precision beats recall here. clarify: tools: diff --git a/pkg/bash/category_test.go b/pkg/bash/category_test.go index 2e47745..044096c 100644 --- a/pkg/bash/category_test.go +++ b/pkg/bash/category_test.go @@ -238,6 +238,19 @@ func TestClassifyTool(t *testing.T) { {"WebFetch", CategoryExplore}, {"WebSearch", CategoryExplore}, {"UnknownTool", CategoryOther}, + + // File mutations that previously fell through to `other`, so curated + // transcript storage would have dropped them along with the tool noise. + {"MultiEdit", CategoryEdit}, + {"apply_patch", CategoryEdit}, + {"CodexPatchApply", CategoryEdit}, + + // Agent task state. TodoWrite was already covered; the Task* family was + // not, so task transitions classified as `other`. + {"TaskCreate", CategoryPlan}, + {"TaskUpdate", CategoryPlan}, + {"TaskGet", CategoryPlan}, + {"TaskList", CategoryPlan}, } for _, tt := range tests { @@ -250,6 +263,36 @@ func TestClassifyTool(t *testing.T) { } } +// Shell-level file mutations must stay distinct from read-only inspection of the +// same binary: `sed -i` edits, bare `sed` explores. +func TestClassifyBashFileMutations(t *testing.T) { + classifier := NewCategoryClassifier(DefaultCategoryConfig()) + + tests := []struct { + command string + expected Category + }{ + {"sed -i 's/a/b/' file.go", CategoryEdit}, + {"sed -n '1,5p' file.go", CategoryExplore}, + {"perl -pi -e 's/a/b/' file.go", CategoryEdit}, + {"mv old.go new.go", CategoryEdit}, + {"cp src.go dst.go", CategoryEdit}, + {"touch new.go", CategoryEdit}, + {"mkdir -p pkg/foo", CategoryEdit}, + {"cat file.go", CategoryExplore}, + {"rg pattern .", CategoryExplore}, + } + + for _, tt := range tests { + t.Run(tt.command, func(t *testing.T) { + got := classifier.ClassifyBash(tt.command) + if got != tt.expected { + t.Errorf("ClassifyBash(%q) = %q, want %q", tt.command, got, tt.expected) + } + }) + } +} + func TestClassifyToolWithPath(t *testing.T) { classifier := NewCategoryClassifier(DefaultCategoryConfig()) diff --git a/pkg/captainconfig/config.go b/pkg/captainconfig/config.go index fb57720..d5278ef 100644 --- a/pkg/captainconfig/config.go +++ b/pkg/captainconfig/config.go @@ -11,16 +11,52 @@ import ( "io/fs" "os" "path/filepath" + "strings" + "github.com/flanksource/captain/pkg/api/registry" + "golang.org/x/sys/unix" "gopkg.in/yaml.v3" ) type Config struct { - AI AIDefaults `yaml:"ai"` - Prompts PromptDefaults `yaml:"prompts"` + AI AIDefaults `yaml:"ai"` + Prompts PromptDefaults `yaml:"prompts"` + Attachments AttachmentDefaults `yaml:"attachments"` +} + +type AttachmentDefaults struct { + Directory string `yaml:"directory,omitempty"` + MaxFileBytes int64 `yaml:"maxFileBytes,omitempty"` + MaxRequestBytes int64 `yaml:"maxRequestBytes,omitempty"` + MaxFiles int `yaml:"maxFiles,omitempty"` + Retention string `yaml:"retention,omitempty"` +} + +func (a AttachmentDefaults) WithDefaults() AttachmentDefaults { + if a.Directory == "" { + a.Directory = ".captain/attachments" + } + if a.MaxFileBytes == 0 { + a.MaxFileBytes = 20 << 20 + } + if a.MaxRequestBytes == 0 { + a.MaxRequestBytes = 50 << 20 + } + if a.MaxFiles == 0 { + a.MaxFiles = 10 + } + if a.Retention == "" { + a.Retention = "30d" + } + return a } type AIDefaults struct { + DefaultProvider string `yaml:"defaultProvider,omitempty"` + Providers map[string]ProviderDefaults `yaml:"providers,omitempty"` + + // Legacy global selection fields are read so existing configurations can be + // projected into the provider map. Save omits them after migration. Backend string `yaml:"backend,omitempty"` Model string `yaml:"model,omitempty"` ReasoningEffort string `yaml:"reasoningEffort,omitempty"` @@ -37,8 +73,63 @@ type AIDefaults struct { NoMemory bool `yaml:"noMemory,omitempty"` } +type ProviderDefaults struct { + Agent string `yaml:"agent,omitempty" json:"agent"` + Model string `yaml:"model,omitempty" json:"model"` + ReasoningEffort string `yaml:"reasoningEffort,omitempty" json:"effort"` +} + +func (a AIDefaults) ActiveProvider() string { + if provider := registry.Backend(strings.TrimSpace(a.DefaultProvider)); provider != "" && provider.Provider() == provider { + return string(provider) + } + if provider := registry.Backend(strings.TrimSpace(a.Backend)).Provider(); provider != "" { + return string(provider) + } + if backend, err := registry.InferBackend(strings.TrimSpace(a.Model)); err == nil { + return string(backend.Provider()) + } + return string(registry.AnthropicProvider) +} + +func (a AIDefaults) legacyProvider() string { + if provider := registry.Backend(strings.TrimSpace(a.Backend)).Provider(); provider != "" { + return string(provider) + } + if backend, err := registry.InferBackend(strings.TrimSpace(a.Model)); err == nil { + return string(backend.Provider()) + } + return strings.TrimSpace(a.DefaultProvider) +} + +func (a AIDefaults) Provider(provider string) ProviderDefaults { + provider = strings.TrimSpace(provider) + defaults := a.Providers[provider] + if provider != a.legacyProvider() { + return defaults + } + legacyAgent := registry.Backend(strings.TrimSpace(a.Backend)) + if defaults.Agent == "" && legacyAgent.Provider() == registry.Backend(provider) { + defaults.Agent = string(legacyAgent) + } + if defaults.Model == "" { + defaults.Model = strings.TrimSpace(a.Model) + } + if defaults.ReasoningEffort == "" { + defaults.ReasoningEffort = strings.TrimSpace(a.ReasoningEffort) + } + return defaults +} + type PromptDefaults struct { - Dirs []string `yaml:"dirs,omitempty"` + Dirs []string `yaml:"dirs,omitempty"` + SchemaRepair SchemaRepairDefaults `yaml:"schemaRepair,omitempty"` +} + +type SchemaRepairDefaults struct { + Model string `yaml:"model,omitempty"` + Backend string `yaml:"backend,omitempty"` + Prompt string `yaml:"prompt,omitempty"` } // pathOverride lets tests redirect Path() to a temp directory without touching @@ -71,6 +162,10 @@ func Load() (Config, bool, error) { if err != nil { return Config{}, false, err } + return load(path) +} + +func load(path string) (Config, bool, error) { data, err := os.ReadFile(path) if err != nil { if errors.Is(err, fs.ErrNotExist) { @@ -93,6 +188,69 @@ func Save(cfg Config) error { if err != nil { return err } + return withLock(path, func() error { return write(path, normalize(cfg)) }) +} + +// Update serializes a read-modify-write operation so concurrent CLI and web +// configuration changes cannot overwrite unrelated settings. +func Update(update func(*Config) error) error { + if update == nil { + return fmt.Errorf("config update function is required") + } + path, err := Path() + if err != nil { + return err + } + return withLock(path, func() error { + cfg, _, err := load(path) + if err != nil { + return err + } + if err := update(&cfg); err != nil { + return err + } + return write(path, normalize(cfg)) + }) +} + +func normalize(cfg Config) Config { + a := &cfg.AI + hasLegacy := a.Backend != "" || a.Model != "" || a.ReasoningEffort != "" + if hasLegacy { + provider := a.legacyProvider() + if provider == "" { + provider = a.ActiveProvider() + } + if a.Providers == nil { + a.Providers = map[string]ProviderDefaults{} + } + a.Providers[provider] = a.Provider(provider) + if a.DefaultProvider == "" { + a.DefaultProvider = provider + } + } + a.Backend, a.Model, a.ReasoningEffort = "", "", "" + return cfg +} + +func withLock(path string, action func() error) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("ensure %s: %w", dir, err) + } + lock, err := os.OpenFile(path+".lock", os.O_CREATE|os.O_RDWR, 0o644) + if err != nil { + return fmt.Errorf("open config lock: %w", err) + } + defer func() { _ = lock.Close() }() + if err := unix.Flock(int(lock.Fd()), unix.LOCK_EX); err != nil { + return fmt.Errorf("lock config: %w", err) + } + defer func() { _ = unix.Flock(int(lock.Fd()), unix.LOCK_UN) }() + return action() +} + +func write(path string, cfg Config) error { data, err := yaml.Marshal(cfg) if err != nil { return fmt.Errorf("marshal config: %w", err) diff --git a/pkg/captainconfig/config_test.go b/pkg/captainconfig/config_test.go index e3d23fd..3b3dba4 100644 --- a/pkg/captainconfig/config_test.go +++ b/pkg/captainconfig/config_test.go @@ -4,6 +4,7 @@ import ( "os" "path/filepath" "reflect" + "strings" "testing" ) @@ -34,17 +35,26 @@ func TestSaveLoad_RoundTrip(t *testing.T) { path := withTempPath(t) want := Config{ AI: AIDefaults{ - Backend: "anthropic", - Model: "claude-sonnet-4-6", - ReasoningEffort: "medium", - BudgetUSD: 2.5, - MaxTokens: 8192, - Temperature: 0.2, - Timeout: "180s", - NoCache: true, - NoMCP: true, - NoHooks: false, - NoMemory: true, + DefaultProvider: "anthropic", + Providers: map[string]ProviderDefaults{ + "anthropic": {Agent: "claude-agent", Model: "claude-sonnet-4-6", ReasoningEffort: "medium"}, + }, + BudgetUSD: 2.5, + MaxTokens: 8192, + Temperature: 0.2, + Timeout: "180s", + NoCache: true, + NoMCP: true, + NoHooks: false, + NoMemory: true, + }, + Prompts: PromptDefaults{ + Dirs: []string{"/repo/prompts"}, + SchemaRepair: SchemaRepairDefaults{ + Model: "gpt-5", + Backend: "openai", + Prompt: "/repo/prompts/json-repair.prompt", + }, }, } if err := Save(want); err != nil { @@ -71,6 +81,81 @@ func TestSaveLoad_RoundTrip(t *testing.T) { } } +func TestLegacyDefaultsProjectIntoActiveProviderAndMigrateOnSave(t *testing.T) { + path := withTempPath(t) + legacy := []byte("ai:\n backend: codex-agent\n model: gpt-5.6-sol\n reasoningEffort: high\n maxTokens: 4096\n") + if err := os.WriteFile(path, legacy, 0o644); err != nil { + t.Fatalf("seed legacy config: %v", err) + } + cfg, _, err := Load() + if err != nil { + t.Fatalf("Load() legacy config: %v", err) + } + if got := cfg.AI.ActiveProvider(); got != "openai" { + t.Fatalf("ActiveProvider() = %q, want openai", got) + } + if got := cfg.AI.Provider("openai"); !reflect.DeepEqual(got, ProviderDefaults{ + Agent: "codex-agent", Model: "gpt-5.6-sol", ReasoningEffort: "high", + }) { + t.Fatalf("Provider(openai) = %+v", got) + } + if err := Save(cfg); err != nil { + t.Fatalf("Save() migrated config: %v", err) + } + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read migrated config: %v", err) + } + text := string(data) + if strings.Contains(text, "backend:") || strings.Contains(text, "\n model:") || !strings.Contains(text, "defaultProvider: openai") { + t.Fatalf("legacy fields were not migrated:\n%s", text) + } +} + +func TestLegacyDefaultsStayWithTheirProviderWhenDefaultProviderChanges(t *testing.T) { + cfg := Config{AI: AIDefaults{ + DefaultProvider: "gemini", + Backend: "codex-agent", + Model: "gpt-5.6-sol", + ReasoningEffort: "high", + }} + + normalized := normalize(cfg) + if got := normalized.AI.Providers["openai"]; got.Agent != "codex-agent" || got.Model != "gpt-5.6-sol" || got.ReasoningEffort != "high" { + t.Fatalf("openai legacy defaults = %+v", got) + } + if got := normalized.AI.Providers["gemini"]; got != (ProviderDefaults{}) { + t.Fatalf("gemini inherited openai legacy defaults: %+v", got) + } + if normalized.AI.DefaultProvider != "gemini" { + t.Fatalf("default provider = %q", normalized.AI.DefaultProvider) + } +} + +func TestUpdatePreservesUnrelatedConfiguration(t *testing.T) { + withTempPath(t) + wantPrompt := PromptDefaults{Dirs: []string{"/repo/prompts"}} + if err := Save(Config{Prompts: wantPrompt}); err != nil { + t.Fatalf("seed config: %v", err) + } + if err := Update(func(cfg *Config) error { + cfg.AI.DefaultProvider = "gemini" + cfg.AI.Providers = map[string]ProviderDefaults{ + "gemini": {Agent: "gemini-cli", Model: "gemini-3.5-flash"}, + } + return nil + }); err != nil { + t.Fatalf("Update(): %v", err) + } + got, _, err := Load() + if err != nil { + t.Fatalf("Load(): %v", err) + } + if !reflect.DeepEqual(got.Prompts, wantPrompt) || got.AI.DefaultProvider != "gemini" { + t.Fatalf("updated config = %+v", got) + } +} + func TestLoad_MalformedYAMLReturnsError(t *testing.T) { path := withTempPath(t) if err := os.WriteFile(path, []byte("ai: [not, a, mapping"), 0o644); err != nil { @@ -93,7 +178,7 @@ func TestSave_AtomicLeavesNoTempFile(t *testing.T) { } for _, e := range entries { name := e.Name() - if name == ".captain.yaml" { + if name == ".captain.yaml" || name == ".captain.yaml.lock" { continue } t.Errorf("found stray file alongside config: %s", name) diff --git a/pkg/claude/cost.go b/pkg/claude/cost.go index 0d95899..fc976b4 100644 --- a/pkg/claude/cost.go +++ b/pkg/claude/cost.go @@ -1,20 +1,14 @@ package claude import ( - "github.com/segmentio/encoding/json" "sort" - "strings" -) -type ModelFamily string + "github.com/segmentio/encoding/json" -const ( - ModelFamilyOpus4 ModelFamily = "opus-4" - ModelFamilySonnet4 ModelFamily = "sonnet-4" - ModelFamilyHaiku4 ModelFamily = "haiku-4" - ModelFamilyUnknown ModelFamily = "unknown" + "github.com/flanksource/captain/pkg/api/registry" ) +// ModelPricing is a model's list price in USD per million tokens. type ModelPricing struct { InputPerMTok float64 OutputPerMTok float64 @@ -22,57 +16,35 @@ type ModelPricing struct { CacheReadPerMTok float64 } -// PricingTable maps model families to their per-million-token pricing in USD. -// Source: https://docs.anthropic.com/en/docs/about-claude/models -var PricingTable = map[ModelFamily]ModelPricing{ - ModelFamilyOpus4: { - InputPerMTok: 15.0, - OutputPerMTok: 75.0, - CacheWritePerMTok: 18.75, - CacheReadPerMTok: 1.50, - }, - ModelFamilySonnet4: { - InputPerMTok: 3.0, - OutputPerMTok: 15.0, - CacheWritePerMTok: 3.75, - CacheReadPerMTok: 0.30, - }, - ModelFamilyHaiku4: { - InputPerMTok: 0.80, - OutputPerMTok: 4.0, - CacheWritePerMTok: 1.0, - CacheReadPerMTok: 0.08, - }, -} - -func ClassifyModel(model string) ModelFamily { - m := strings.ToLower(model) - switch { - case strings.Contains(m, "opus"): - return ModelFamilyOpus4 - case strings.Contains(m, "sonnet"): - return ModelFamilySonnet4 - case strings.Contains(m, "haiku"): - return ModelFamilyHaiku4 - default: - return ModelFamilyUnknown +// PricingFor reads a model's list price from the generated catalog, which +// carries models.dev's per-model rates. ok is false when the catalog snapshot +// prices no such model; callers must render that as "unknown", never as another +// model's rate. This replaced a hand-written opus/sonnet/haiku family table that +// classified on substring alone and so kept billing Opus 4.5 and newer at the +// retired 4.1 rate of $15/$75 rather than $5/$25. +func PricingFor(model string) (ModelPricing, bool) { + cost, ok := registry.CostFor(model) + if !ok { + return ModelPricing{}, false } + return ModelPricing{ + InputPerMTok: cost.Input, + OutputPerMTok: cost.Output, + CacheWritePerMTok: cost.CacheWrite, + CacheReadPerMTok: cost.CacheRead, + }, true } -// PricingFor returns the per-million-token pricing for a model, falling back to -// Sonnet rates for unrecognized model families (matching CalculateCost). -func PricingFor(model string) ModelPricing { - if pricing, ok := PricingTable[ClassifyModel(model)]; ok { - return pricing - } - return PricingTable[ModelFamilySonnet4] -} - +// CalculateCost totals a usage record in USD, returning 0 for a model the +// catalog does not price. func CalculateCost(usage *Usage, model string) float64 { if usage == nil { return 0 } - pricing := PricingFor(model) + pricing, ok := PricingFor(model) + if !ok { + return 0 + } return float64(usage.InputTokens)*pricing.InputPerMTok/1e6 + float64(usage.OutputTokens)*pricing.OutputPerMTok/1e6 + float64(usage.CacheCreationInputTokens)*pricing.CacheWritePerMTok/1e6 + diff --git a/pkg/claude/cost_test.go b/pkg/claude/cost_test.go index 1494c71..dd229c4 100644 --- a/pkg/claude/cost_test.go +++ b/pkg/claude/cost_test.go @@ -7,23 +7,29 @@ import ( "github.com/stretchr/testify/assert" ) -func TestClassifyModel(t *testing.T) { +// TestPricingFor covers the catalog lookup that replaced the opus/sonnet/haiku +// family table. Rates are the vendors' published per-MTok list prices; the Opus +// case is the regression guard, since the old table billed every "opus" id at +// the retired 4.1 rate of $15/$75. +func TestPricingFor(t *testing.T) { tests := []struct { - model string - expected ModelFamily + model string + want ModelPricing + priced bool }{ - {"claude-opus-4-6", ModelFamilyOpus4}, - {"claude-opus-4-5-20251101", ModelFamilyOpus4}, - {"claude-sonnet-4-6", ModelFamilySonnet4}, - {"claude-sonnet-4-5-20241022", ModelFamilySonnet4}, - {"claude-haiku-4-5-20251001", ModelFamilyHaiku4}, - {"", ModelFamilyUnknown}, - {"gpt-4o", ModelFamilyUnknown}, + {"claude-opus-4-6", ModelPricing{5, 25, 6.25, 0.5}, true}, + {"claude-opus-4-5-20251101", ModelPricing{5, 25, 6.25, 0.5}, true}, + {"claude-sonnet-4-6", ModelPricing{3, 15, 3.75, 0.3}, true}, + {"claude-haiku-4-5-20251001", ModelPricing{1, 5, 1.25, 0.1}, true}, + {"", ModelPricing{}, false}, + {"gpt-4o", ModelPricing{}, false}, } for _, tt := range tests { t.Run(tt.model, func(t *testing.T) { - assert.Equal(t, tt.expected, ClassifyModel(tt.model)) + got, ok := PricingFor(tt.model) + assert.Equal(t, tt.priced, ok) + assert.Equal(t, tt.want, got) }) } } @@ -48,7 +54,7 @@ func TestCalculateCost(t *testing.T) { OutputTokens: 1_000_000, }, model: "claude-opus-4-6", - expected: 15.0 + 75.0, // $90 + expected: 5.0 + 25.0, // $30 at the Opus 4.5+ rate, not the retired $90 }, { name: "sonnet with cache", @@ -69,16 +75,16 @@ func TestCalculateCost(t *testing.T) { OutputTokens: 50_000, }, model: "claude-haiku-4-5-20251001", - expected: 0.08 + 0.20, // $0.28 + expected: 0.10 + 0.25, // $0.35 at $1/$5 }, { - name: "unknown model falls back to sonnet pricing", + name: "unpriced model reports no cost rather than a guess", usage: &Usage{ InputTokens: 1_000_000, OutputTokens: 1_000_000, }, model: "unknown-model", - expected: 3.0 + 15.0, // $18 + expected: 0, }, } diff --git a/pkg/claude/history.go b/pkg/claude/history.go index d473fae..5b6229e 100644 --- a/pkg/claude/history.go +++ b/pkg/claude/history.go @@ -24,8 +24,22 @@ type HistoryEntry struct { // PlanFilePath is set on synthetic entries surfaced from plan_mode / // plan_mode_exit attachments; it points at the session's plan file even when // the transcript carries no ExitPlanMode tool call or plan-file write. - PlanFilePath string `json:"-"` - RawLine json.RawMessage `json:"-"` + PlanFilePath string `json:"-"` + Event *TranscriptEvent `json:"-"` + RawLine json.RawMessage `json:"-"` + // Line is the 1-based JSONL line number the entry was read from, so + // downstream consumers can seek back into the transcript file. + Line int `json:"-"` +} + +// TranscriptEvent is a non-message, non-tool transcript line. It lets callers +// preserve session/turn metadata without pretending discovery or budget records +// are model messages. +type TranscriptEvent struct { + Type string `json:"type"` + Scope string `json:"scope,omitempty"` + Subtype string `json:"subtype,omitempty"` + Data map[string]any `json:"data,omitempty"` } // Message represents a conversation message @@ -107,6 +121,7 @@ type CacheCreation struct { // ServerToolUse tracks server-side tool usage type ServerToolUse struct { WebSearchRequests int `json:"web_search_requests,omitempty"` + WebFetchRequests int `json:"web_fetch_requests,omitempty"` } // IsUserMessage returns true if this is a user message diff --git a/pkg/claude/hooks.go b/pkg/claude/hooks.go index 70f13ef..3e9dd5d 100644 --- a/pkg/claude/hooks.go +++ b/pkg/claude/hooks.go @@ -30,6 +30,13 @@ type HookInput struct { TranscriptPath string `json:"transcript_path,omitempty"` StopHookReason string `json:"stop_hook_reason,omitempty"` Prompt string `json:"prompt,omitempty"` + CWD string `json:"cwd,omitempty"` + HookEventName string `json:"hook_event_name,omitempty"` + // Source is set on SessionStart: startup|resume|clear|compact. + Source string `json:"source,omitempty"` + // Reason is set on SessionEnd: clear|resume|logout|prompt_input_exit|other. + Reason string `json:"reason,omitempty"` + LastAssistantMessage string `json:"last_assistant_message,omitempty"` } // HookOutput is the JSON returned by hooks diff --git a/pkg/claude/hooks_test.go b/pkg/claude/hooks_test.go index eec263f..1ba43c2 100644 --- a/pkg/claude/hooks_test.go +++ b/pkg/claude/hooks_test.go @@ -2,6 +2,7 @@ package claude import ( "encoding/json" + "reflect" "testing" ) @@ -107,6 +108,110 @@ func TestHookOutput_Marshal(t *testing.T) { } } +func TestHookInput_Unmarshal_LifecycleEvents(t *testing.T) { + cases := []struct { + name string + payload string + want HookInput + }{ + { + name: "SessionStart", + payload: `{ + "session_id": "abc123", + "transcript_path": "/Users/x/.claude/projects/-repo/abc123.jsonl", + "cwd": "/repo", + "hook_event_name": "SessionStart", + "source": "startup" + }`, + want: HookInput{ + SessionID: "abc123", + TranscriptPath: "/Users/x/.claude/projects/-repo/abc123.jsonl", + CWD: "/repo", + HookEventName: "SessionStart", + Source: "startup", + }, + }, + { + name: "SessionEnd", + payload: `{ + "session_id": "abc123", + "transcript_path": "/Users/x/.claude/projects/-repo/abc123.jsonl", + "cwd": "/repo", + "hook_event_name": "SessionEnd", + "reason": "prompt_input_exit" + }`, + want: HookInput{ + SessionID: "abc123", + TranscriptPath: "/Users/x/.claude/projects/-repo/abc123.jsonl", + CWD: "/repo", + HookEventName: "SessionEnd", + Reason: "prompt_input_exit", + }, + }, + { + name: "Stop", + payload: `{ + "session_id": "abc123", + "transcript_path": "/Users/x/.claude/projects/-repo/abc123.jsonl", + "cwd": "/repo", + "hook_event_name": "Stop", + "last_assistant_message": "done" + }`, + want: HookInput{ + SessionID: "abc123", + TranscriptPath: "/Users/x/.claude/projects/-repo/abc123.jsonl", + CWD: "/repo", + HookEventName: "Stop", + LastAssistantMessage: "done", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var got HookInput + if err := json.Unmarshal([]byte(tc.payload), &got); err != nil { + t.Fatalf("unmarshal failed: %v", err) + } + if !reflect.DeepEqual(got, tc.want) { + t.Errorf("got %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestHooksConfig_Marshal_SessionLifecycle(t *testing.T) { + config := HooksConfig{Hooks: map[HookEventType][]HookMatcher{ + HookEventSessionStart: {{Hooks: []Hook{{Type: HookTypeCommand, Command: "captain hook monitor notify --provider claude", Timeout: 10}}}}, + HookEventSessionEnd: {{Hooks: []Hook{{Type: HookTypeCommand, Command: "captain hook monitor notify --provider claude", Timeout: 10}}}}, + }} + + data, err := json.Marshal(config) + if err != nil { + t.Fatalf("marshal failed: %v", err) + } + + var parsed map[string]map[string][]map[string]any + if err := json.Unmarshal(data, &parsed); err != nil { + t.Fatalf("unmarshal failed: %v", err) + } + + for _, event := range []string{"SessionStart", "SessionEnd"} { + matchers := parsed["hooks"][event] + if len(matchers) != 1 { + t.Fatalf("%s: expected 1 matcher, got %d", event, len(matchers)) + } + if _, exists := matchers[0]["matcher"]; exists { + t.Errorf("%s: matcher key should be omitted for lifecycle events", event) + } + hooks := matchers[0]["hooks"].([]any) + hook := hooks[0].(map[string]any) + if hook["type"] != "command" || hook["timeout"] != float64(10) { + t.Errorf("%s: unexpected hook entry: %+v", event, hook) + } + } +} + func TestHookOutput_Marshal_OmitsNil(t *testing.T) { output := HookOutput{Continue: true} diff --git a/pkg/claude/reader.go b/pkg/claude/reader.go index addceb3..ef99b8d 100644 --- a/pkg/claude/reader.go +++ b/pkg/claude/reader.go @@ -99,6 +99,7 @@ func readJSONL(r io.Reader, fallbackToHistoryEntry bool, opts ReadOptions) ([]Hi if opts.KeepRaw { entry.RawLine = append(json.RawMessage(nil), line...) } + entry.Line = lineNo entries = append(entries, entry) continue } @@ -110,6 +111,7 @@ func readJSONL(r io.Reader, fallbackToHistoryEntry bool, opts ReadOptions) ([]Hi if opts.KeepRaw { entry.RawLine = append(json.RawMessage(nil), line...) } + entry.Line = lineNo entries = append(entries, entry) } } @@ -125,6 +127,11 @@ type streamJSONLine struct { Subtype string `json:"subtype,omitempty"` Message json.RawMessage `json:"message,omitempty"` + // Content is the top-level content on system/local_command lines. Kept raw + // because other line types (e.g. queue-operation) carry a non-string content + // object; contentString decodes it only when it is a JSON string. + Content json.RawMessage `json:"content,omitempty"` + // Stream-json fields SessionIDSnake string `json:"session_id,omitempty"` TimestampSnake string `json:"timestamp,omitempty"` @@ -136,7 +143,7 @@ type streamJSONLine struct { CWD string `json:"cwd,omitempty"` GitBranch string `json:"gitBranch,omitempty"` Slug string `json:"slug,omitempty"` - Error string `json:"error,omitempty"` + Error json.RawMessage `json:"error,omitempty"` Attachment json.RawMessage `json:"attachment,omitempty"` } @@ -151,21 +158,34 @@ func (sj streamJSONLine) timestamp() string { return sj.TimestampSnake } +// contentString returns the top-level content when it is a JSON string, and "" +// when it is absent or a non-string object. Only system/local_command lines +// carry a string content wrapper this reader needs to inspect. +func (sj streamJSONLine) contentString() string { + if len(sj.Content) == 0 || sj.Content[0] != '"' { + return "" + } + var s string + if err := json.Unmarshal(sj.Content, &s); err != nil { + return "" + } + return s +} + // knownSessionStorageTypes are line types that appear in on-disk session // files for state tracking but carry no useful row-level information. // Listed explicitly so they don't pollute the unhandled-types diagnostic. var knownSessionStorageTypes = map[string]bool{ "file-history-snapshot": true, - "last-prompt": true, + "file-history-delta": true, // incremental checkpoint bookkeeping for the snapshot above "permission-mode": true, "agent-name": true, // Operational/streaming state with no unique row-level content — the real // content surfaces via the actual user/assistant messages. Listed so they // don't pollute the unhandled-types diagnostic. - "mode": true, // active mode marker (e.g. {"mode":"normal"}) - "bridge-session": true, // cloud bridge-session linkage - "progress": true, // intermediate streaming progress, superseded by the final message - "queue-operation": true, // message-queue bookkeeping; dequeued content appears as a real message + "mode": true, // active mode marker (e.g. {"mode":"normal"}) + "bridge-session": true, // cloud bridge-session linkage + "progress": true, // intermediate streaming progress, superseded by the final message } // planAttachment is the plan-mode attachment Claude Code writes when entering @@ -200,6 +220,187 @@ func attachmentEntry(sj streamJSONLine) []HistoryEntry { }) } +func metadataEventEntry(sj streamJSONLine, eventType, scope string, data map[string]any) []HistoryEntry { + if eventType == "" { + return nil + } + return single(HistoryEntry{ + SessionID: sj.sessionID(), + UUID: sj.UUID, + Timestamp: sj.timestamp(), + CWD: sj.CWD, + GitBranch: sj.GitBranch, + Slug: sj.Slug, + Event: &TranscriptEvent{ + Type: eventType, + Scope: scope, + Data: data, + }, + }) +} + +func rawObject(raw []byte) map[string]any { + var out map[string]any + if err := json.Unmarshal(raw, &out); err != nil { + return nil + } + delete(out, "message") + return out +} + +func attachmentEventEntry(sj streamJSONLine, raw []byte) []HistoryEntry { + var attachment map[string]any + if err := json.Unmarshal(sj.Attachment, &attachment); err != nil { + return nil + } + typ, _ := attachment["type"].(string) + switch typ { + case "deferred_tools_delta", "agent_listing_delta", "skill_listing": + return metadataEventEntry(sj, typ, "session", attachment) + case "budget_usd": + return metadataEventEntry(sj, typ, "turn", attachment) + case "goal_status": + // Session-scoped goal directive. Kept as a map so future goal fields + // (beyond condition/met/sentinel/reason) round-trip without parser churn. + return metadataEventEntry(sj, "goal_status", "session", attachment) + default: + return attachmentEntry(sj) + } +} + +// claudeCommandRecord is the parsed form of a Claude slash-command wrapper +// (//). Claude writes it as a +// text-only user message (modern) or a system/local_command line (older). +type claudeCommandRecord struct { + Name string + Message string + Args string +} + +// claudeCommandOutputRecord is the parsed form of a or +// wrapper carrying a slash command's captured output. +type claudeCommandOutputRecord struct { + Stream string + Content string +} + +// cutWrapperTag extracts the inner content of the section that s +// must begin with, returning the content and the remainder after the closing +// tag. last selects strings.LastIndex for the closing tag so a trailing +// free-form section (args/output) may itself contain the literal closing tag. +func cutWrapperTag(s, tag string, last bool) (inner, rest string, err error) { + open, closeTag := "<"+tag+">", "" + if !strings.HasPrefix(s, open) { + return "", "", fmt.Errorf("expected <%s>", tag) + } + body := s[len(open):] + idx := strings.Index(body, closeTag) + if last { + idx = strings.LastIndex(body, closeTag) + } + if idx < 0 { + return "", "", fmt.Errorf("unterminated <%s>", tag) + } + return body[:idx], body[idx+len(closeTag):], nil +} + +// parseClaudeCommandRecord recognizes a complete slash-command wrapper. It +// returns matched=false for text that does not open with (so +// ordinary prose is left as chat) and matched=true with an error for a +// recognized-but-malformed wrapper (surfaced as a ParseError row). Empty +// command arguments are valid. +func parseClaudeCommandRecord(text string) (claudeCommandRecord, bool, error) { + trimmed := strings.TrimSpace(text) + if !strings.HasPrefix(trimmed, "") { + return claudeCommandRecord{}, false, nil + } + name, rest, err := cutWrapperTag(trimmed, "command-name", false) + if err != nil { + return claudeCommandRecord{}, true, err + } + message, rest, err := cutWrapperTag(strings.TrimSpace(rest), "command-message", false) + if err != nil { + return claudeCommandRecord{}, true, err + } + args, rest, err := cutWrapperTag(strings.TrimSpace(rest), "command-args", true) + if err != nil { + return claudeCommandRecord{}, true, err + } + if strings.TrimSpace(rest) != "" { + return claudeCommandRecord{}, true, fmt.Errorf("unexpected content after ") + } + return claudeCommandRecord{Name: name, Message: message, Args: args}, true, nil +} + +// parseClaudeCommandOutputRecord recognizes a complete local-command output +// wrapper. Its matched/error contract mirrors parseClaudeCommandRecord. Empty +// output is valid. +func parseClaudeCommandOutputRecord(text string) (claudeCommandOutputRecord, bool, error) { + trimmed := strings.TrimSpace(text) + for _, s := range []struct{ tag, stream string }{ + {"local-command-stdout", "stdout"}, + {"local-command-stderr", "stderr"}, + } { + if !strings.HasPrefix(trimmed, "<"+s.tag+">") { + continue + } + content, rest, err := cutWrapperTag(trimmed, s.tag, true) + if err != nil { + return claudeCommandOutputRecord{}, true, err + } + if strings.TrimSpace(rest) != "" { + return claudeCommandOutputRecord{}, true, fmt.Errorf("unexpected content after ", s.tag) + } + return claudeCommandOutputRecord{Stream: s.stream, Content: content}, true, nil + } + return claudeCommandOutputRecord{}, false, nil +} + +func claudeCommandEntry(sj streamJSONLine, command claudeCommandRecord) HistoryEntry { + return metadataEventEntry(sj, "claude_command", "turn", map[string]any{ + "command_name": command.Name, + "command_message": command.Message, + "command_args": command.Args, + })[0] +} + +func claudeCommandOutputEntry(sj streamJSONLine, output claudeCommandOutputRecord) HistoryEntry { + return metadataEventEntry(sj, "claude_command_output", "turn", map[string]any{ + "stream": output.Stream, + "content": output.Content, + })[0] +} + +// classifyClaudeCommandText converts a command wrapper or its output into +// structured event rows. It returns ok=false when text is not a recognized +// wrapper so the caller keeps its normal handling. A recognized-but-malformed +// wrapper yields a ParseError row (ok=true) rather than reverting to raw text. +func classifyClaudeCommandText(sj streamJSONLine, text string, raw []byte, lineNo int) ([]HistoryEntry, bool) { + if command, matched, err := parseClaudeCommandRecord(text); matched { + if err != nil { + return []HistoryEntry{parseErrorEntry(lineNo, raw, err, sj.timestamp())}, true + } + return single(claudeCommandEntry(sj, command)), true + } + if output, matched, err := parseClaudeCommandOutputRecord(text); matched { + if err != nil { + return []HistoryEntry{parseErrorEntry(lineNo, raw, err, sj.timestamp())}, true + } + return single(claudeCommandOutputEntry(sj, output)), true + } + return nil, false +} + +// singleTextContent returns the text of a message that is exactly one text +// block, so command-wrapper classification only fires on text-only user turns +// and never on mixed tool_use/tool_result content. +func singleTextContent(msg Message) (string, bool) { + if len(msg.Content) != 1 || msg.Content[0].Type != ContentTypeText { + return "", false + } + return msg.Content[0].Text, true +} + // dispatchEvent routes a typed line to zero, one, or more HistoryEntry rows. // A single line can emit multiple entries — e.g. an assistant message with a // top-level "error" field yields both the regular assistant entry and an @@ -215,6 +416,13 @@ func dispatchEvent(sj streamJSONLine, raw []byte, lineNo int) []HistoryEntry { if err := json.Unmarshal(sj.Message, &msg); err != nil { return []HistoryEntry{parseErrorEntry(lineNo, raw, err, sj.timestamp())} } + if sj.Type == "user" { + if text, ok := singleTextContent(msg); ok { + if entries, ok := classifyClaudeCommandText(sj, text, raw, lineNo); ok { + return entries + } + } + } out := []HistoryEntry{{ SessionID: sj.sessionID(), UUID: sj.UUID, @@ -230,6 +438,9 @@ func dispatchEvent(sj streamJSONLine, raw []byte, lineNo int) []HistoryEntry { } return out + case "queue-operation": + return metadataEventEntry(sj, "queue-operation", "turn", rawObject(raw)) + case "system": switch sj.Subtype { case "init": @@ -263,11 +474,19 @@ func dispatchEvent(sj streamJSONLine, raw []byte, lineNo int) []HistoryEntry { "content", "compactMetadata", "level", })) case "local_command": - return single(syntheticEntry(sj, "LocalCommand", raw, []string{"content", "level"})) + if entries, ok := classifyClaudeCommandText(sj, sj.contentString(), raw, lineNo); ok { + return entries + } + return single(syntheticEntry(sj, "LocalCommand", raw, []string{"content", "level", "cwd"})) case "scheduled_task_fire": return single(syntheticEntry(sj, "ScheduledTaskFire", raw, []string{"content"})) case "informational": return single(syntheticEntry(sj, "Informational", raw, []string{"content", "level"})) + case "api_error": + return single(syntheticEntry(sj, "ApiError", raw, []string{ + "error", "level", "retryInMs", "retryAttempt", "maxRetries", + "api_error_status", + })) } case "result": @@ -287,8 +506,20 @@ func dispatchEvent(sj streamJSONLine, raw []byte, lineNo int) []HistoryEntry { "prNumber", "prUrl", "prRepository", })) + case "worktree-state": + return single(syntheticEntry(sj, "WorktreeState", raw, []string{"worktreeSession"})) + + case "relocated": + return single(syntheticEntry(sj, "Relocated", raw, []string{"relocatedCwd"})) + + case "started": + return single(syntheticEntry(sj, "Started", raw, []string{"cwd"})) + case "attachment": - return attachmentEntry(sj) + return attachmentEventEntry(sj, raw) + + case "last-prompt": + return metadataEventEntry(sj, "last-prompt", "session", rawObject(raw)) } if knownSessionStorageTypes[sj.Type] { @@ -310,7 +541,7 @@ func single(e HistoryEntry) []HistoryEntry { return []HistoryEntry{e} } // carries a top-level "error" field (e.g. invalid_request, 4xx/5xx from the // model API). The error would otherwise be invisible in history output. func apiErrorFromAssistantLine(sj streamJSONLine, raw []byte) (HistoryEntry, bool) { - if sj.Error == "" { + if len(sj.Error) == 0 || string(sj.Error) == "null" { return HistoryEntry{}, false } return syntheticEntry(sj, "ApiError", raw, []string{ diff --git a/pkg/claude/reader_command_test.go b/pkg/claude/reader_command_test.go new file mode 100644 index 0000000..375c737 --- /dev/null +++ b/pkg/claude/reader_command_test.go @@ -0,0 +1,213 @@ +package claude + +import ( + "fmt" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/claude/tools" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/segmentio/encoding/json" +) + +func TestClaudeCommandParsing(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Claude Command Parsing Suite") +} + +// jsonString encodes s as a JSON string literal for embedding in a JSONL line. +func jsonString(s string) string { + b, _ := json.Marshal(s) + return string(b) +} + +// userLine builds a modern text-only user record whose message content is text. +func userLine(uuid, text string) string { + return fmt.Sprintf( + `{"type":"user","uuid":%q,"sessionId":"s1","timestamp":"2026-07-14T12:16:09.113Z","cwd":"/repo","gitBranch":"main","message":{"role":"user","content":%s}}`, + uuid, jsonString(text)) +} + +// systemLocalCommandLine builds the older system/local_command record shape. +func systemLocalCommandLine(uuid, content string) string { + return fmt.Sprintf( + `{"type":"system","subtype":"local_command","uuid":%q,"sessionId":"s1","timestamp":"2026-07-14T12:16:09.113Z","content":%s,"level":"info"}`, + uuid, jsonString(content)) +} + +func goalStatusLine(uuid, condition string, met, sentinel bool) string { + return fmt.Sprintf( + `{"type":"attachment","uuid":%q,"sessionId":"s1","timestamp":"2026-07-14T12:16:09.113Z","cwd":"/repo","attachment":{"type":"goal_status","met":%t,"sentinel":%t,"condition":%s}}`, + uuid, met, sentinel, jsonString(condition)) +} + +// commandWrapper renders the slash-command envelope with the realistic newline +// and indentation Claude writes between sections. +func commandWrapper(name, message, args string) string { + return fmt.Sprintf( + "%s\n %s\n %s", + name, message, args) +} + +func stdoutWrapper(content string) string { + return "" + content + "" +} + +func stderrWrapper(content string) string { + return "" + content + "" +} + +func readEntries(lines ...string) []HistoryEntry { + entries, err := ReadHistory(strings.NewReader(strings.Join(lines, "\n"))) + Expect(err).NotTo(HaveOccurred()) + return entries +} + +var _ = Describe("Claude command / goal parsing in the shared reader", func() { + const goalCondition = "push and monitor the docker build on PR #32, if any dependencies need updates open PR's for them and pin to the git SHA until they merge into main and create a release" + + Describe("the exact reported /goal three-record shape", func() { + var entries []HistoryEntry + + BeforeEach(func() { + entries = readEntries( + goalStatusLine("uuid-goal", goalCondition, false, true), + userLine("uuid-cmd", commandWrapper("/goal", "goal", goalCondition)), + userLine("uuid-out", stdoutWrapper("Goal set: "+goalCondition)), + ) + }) + + It("surfaces goal_status as a session-scoped event preserving its fields", func() { + Expect(entries).To(HaveLen(3)) + ev := entries[0].Event + Expect(ev).NotTo(BeNil()) + Expect(ev.Type).To(Equal("goal_status")) + Expect(ev.Scope).To(Equal("session")) + Expect(ev.Data["condition"]).To(Equal(goalCondition)) + Expect(ev.Data["met"]).To(Equal(false)) + Expect(ev.Data["sentinel"]).To(Equal(true)) + }) + + It("parses the /goal record into a claude_command event with split fields", func() { + ev := entries[1].Event + Expect(ev).NotTo(BeNil()) + Expect(ev.Type).To(Equal("claude_command")) + Expect(ev.Scope).To(Equal("turn")) + Expect(ev.Data["command_name"]).To(Equal("/goal")) + Expect(ev.Data["command_message"]).To(Equal("goal")) + Expect(ev.Data["command_args"]).To(Equal(goalCondition)) + }) + + It("parses the stdout record into a wrapper-free claude_command_output event", func() { + ev := entries[2].Event + Expect(ev).NotTo(BeNil()) + Expect(ev.Type).To(Equal("claude_command_output")) + Expect(ev.Data["stream"]).To(Equal("stdout")) + Expect(ev.Data["content"]).To(Equal("Goal set: " + goalCondition)) + }) + + It("emits no raw User text carrying the wrapper tags", func() { + for _, e := range entries { + Expect(e.Message.Role).To(BeEmpty()) + Expect(e.Message.GetTextContent()).NotTo(ContainSubstring("")) + Expect(e.Message.GetTextContent()).NotTo(ContainSubstring("")) + } + }) + + It("projects the three records as non-operational event activity", func() { + uses := ExtractToolUses(entries) + names := make([]string, 0, len(uses)) + for _, u := range uses { + names = append(names, u.Tool) + Expect(tools.IsEventToolName(u.Tool)).To(BeTrue(), "expected %q to be an event tool", u.Tool) + } + Expect(names).To(Equal([]string{"GoalStatus", "ClaudeCommand", "ClaudeCommand"})) + }) + }) + + DescribeTable("command wrappers across command names and argument shapes", + func(line string, wantName, wantMessage, wantArgs string) { + entries := readEntries(line) + Expect(entries).To(HaveLen(1)) + ev := entries[0].Event + Expect(ev).NotTo(BeNil()) + Expect(ev.Type).To(Equal("claude_command")) + Expect(ev.Scope).To(Equal("turn")) + Expect(ev.Data["command_name"]).To(Equal(wantName)) + Expect(ev.Data["command_message"]).To(Equal(wantMessage)) + Expect(ev.Data["command_args"]).To(Equal(wantArgs)) + }, + Entry("/plan with empty arguments (modern user shape)", + userLine("u1", commandWrapper("/plan", "plan", "")), "/plan", "plan", ""), + Entry("/clear with empty arguments (modern user shape)", + userLine("u2", commandWrapper("/clear", "clear", "")), "/clear", "clear", ""), + Entry("/effort with non-empty arguments (modern user shape)", + userLine("u3", commandWrapper("/effort", "effort", "high")), "/effort", "effort", "high"), + Entry("/usage via the older system/local_command shape", + systemLocalCommandLine("u4", commandWrapper("/usage", "usage", "")), "/usage", "usage", ""), + Entry("multiline arguments preserved verbatim", + userLine("u5", commandWrapper("/goal", "goal", "line1\nline2")), "/goal", "goal", "line1\nline2"), + ) + + DescribeTable("output wrappers across streams and shapes", + func(line, wantStream, wantContent string) { + entries := readEntries(line) + Expect(entries).To(HaveLen(1)) + ev := entries[0].Event + Expect(ev).NotTo(BeNil()) + Expect(ev.Type).To(Equal("claude_command_output")) + Expect(ev.Scope).To(Equal("turn")) + Expect(ev.Data["stream"]).To(Equal(wantStream)) + Expect(ev.Data["content"]).To(Equal(wantContent)) + }, + Entry("empty stdout (modern user shape)", + userLine("o1", stdoutWrapper("")), "stdout", ""), + Entry("empty stdout (older system/local_command shape)", + systemLocalCommandLine("o2", stdoutWrapper("")), "stdout", ""), + Entry("non-empty stdout (older system/local_command shape)", + systemLocalCommandLine("o3", stdoutWrapper("MCP dialog dismissed")), "stdout", "MCP dialog dismissed"), + Entry("stderr stream", + userLine("o4", stderrWrapper("boom")), "stderr", "boom"), + Entry("multiline output with ANSI escapes preserved", + userLine("o5", stdoutWrapper("\x1b[31mred\x1b[0m\nsecond line")), "stdout", "\x1b[31mred\x1b[0m\nsecond line"), + ) + + Describe("false positives and malformed wrappers", func() { + It("leaves ordinary prose that merely mentions a tag as a User message", func() { + text := "Please review how is parsed in the reader." + entries := readEntries(userLine("p1", text)) + Expect(entries).To(HaveLen(1)) + Expect(entries[0].Event).To(BeNil()) + Expect(entries[0].Message.Role).To(Equal(MessageRoleUser)) + Expect(entries[0].Message.GetTextContent()).To(Equal(text)) + }) + + It("does not classify mixed content that is not a single text block", func() { + line := fmt.Sprintf( + `{"type":"user","uuid":"m1","sessionId":"s1","timestamp":"2026-07-14T12:16:09.113Z","message":{"role":"user","content":[{"type":"text","text":%s},{"type":"tool_result","tool_use_id":"t1","content":"ok"}]}}`, + jsonString(commandWrapper("/goal", "goal", "do things"))) + entries := readEntries(line) + Expect(entries).To(HaveLen(1)) + Expect(entries[0].Event).To(BeNil()) + Expect(entries[0].Message.Role).To(Equal(MessageRoleUser)) + }) + + It("surfaces a recognized-but-incomplete wrapper as a ParseError without stopping the read", func() { + partial := "/goal" // missing message + args sections + entries := readEntries( + userLine("bad", partial), + userLine("good", "a normal follow-up prompt"), + ) + Expect(entries).To(HaveLen(2)) + + parseErr := entries[0].Message.GetToolUses() + Expect(parseErr).To(HaveLen(1)) + Expect(parseErr[0].Name).To(Equal("ParseError")) + + Expect(entries[1].Event).To(BeNil()) + Expect(entries[1].Message.Role).To(Equal(MessageRoleUser)) + Expect(entries[1].Message.GetTextContent()).To(Equal("a normal follow-up prompt")) + }) + }) +}) diff --git a/pkg/claude/reader_test.go b/pkg/claude/reader_test.go index e3f1ab9..99c157e 100644 --- a/pkg/claude/reader_test.go +++ b/pkg/claude/reader_test.go @@ -237,14 +237,57 @@ func TestReadStreamJSON_UnhandledTypes(t *testing.T) { if err != nil { t.Fatalf("ReadStreamJSON(state) failed: %v", err) } - if len(stateEntries) != 0 { - t.Errorf("operational state types should produce no rows, got %d", len(stateEntries)) + if len(stateEntries) != 1 { + t.Fatalf("queue-operation should produce one turn metadata event, got %d", len(stateEntries)) + } + if stateEntries[0].Event == nil || stateEntries[0].Event.Type != "queue-operation" || stateEntries[0].Event.Scope != "turn" { + t.Fatalf("queue-operation event = %+v, want turn metadata event", stateEntries[0].Event) } if got := SnapshotUnhandledStreamTypes(); len(got) != 0 { t.Errorf("operational state types should not be reported as unhandled, got %v", got) } } +func TestReadHistory_MetadataEvents(t *testing.T) { + jsonl := `{"type":"attachment","sessionId":"s","uuid":"tools","timestamp":"2026-07-05T10:00:00Z","attachment":{"type":"deferred_tools_delta","addedNames":["Read","Bash"],"pendingMcpServers":["github"]}} +{"type":"attachment","sessionId":"s","uuid":"agents","timestamp":"2026-07-05T10:00:01Z","attachment":{"type":"agent_listing_delta","addedTypes":["general-purpose"]}} +{"type":"attachment","sessionId":"s","uuid":"skills","timestamp":"2026-07-05T10:00:02Z","attachment":{"type":"skill_listing","names":["gavel-runner"]}} +{"type":"queue-operation","sessionId":"s","uuid":"queue","timestamp":"2026-07-05T10:00:03Z","operation":"enqueue","content":{"type":"message"}} +{"type":"attachment","sessionId":"s","uuid":"budget","timestamp":"2026-07-05T10:00:04Z","attachment":{"type":"budget_usd","used":1.25,"total":5,"remaining":3.75}} +{"type":"last-prompt","sessionId":"s","uuid":"prompt","timestamp":"2026-07-05T10:00:05Z","content":"fix it"}` + + entries, err := ReadHistory(strings.NewReader(jsonl)) + if err != nil { + t.Fatalf("ReadHistory failed: %v", err) + } + if len(entries) != 6 { + t.Fatalf("entries = %d, want 6", len(entries)) + } + want := []struct { + typ string + scope string + uuid string + }{ + {"deferred_tools_delta", "session", "tools"}, + {"agent_listing_delta", "session", "agents"}, + {"skill_listing", "session", "skills"}, + {"queue-operation", "turn", "queue"}, + {"budget_usd", "turn", "budget"}, + {"last-prompt", "session", "prompt"}, + } + for i, w := range want { + if entries[i].Event == nil { + t.Fatalf("entry[%d] missing event", i) + } + if entries[i].Event.Type != w.typ || entries[i].Event.Scope != w.scope || entries[i].UUID != w.uuid { + t.Errorf("entry[%d] event = %+v uuid=%q, want %s/%s/%s", i, entries[i].Event, entries[i].UUID, w.typ, w.scope, w.uuid) + } + if len(entries[i].Message.Content) != 0 { + t.Errorf("entry[%d] metadata event should not create message content: %+v", i, entries[i].Message.Content) + } + } +} + // TestReadStreamJSON_PrLinkSurfaced verifies a pr-link line becomes a PrLink // synthetic row carrying the PR fields, rather than being dropped as unhandled. func TestReadStreamJSON_PrLinkSurfaced(t *testing.T) { @@ -273,13 +316,35 @@ func TestReadStreamJSON_PrLinkSurfaced(t *testing.T) { } } +// TestReadStreamJSON_FileHistoryDeltaIgnored verifies file-history-delta — the +// incremental sibling of file-history-snapshot — is explicitly ignored rather +// than counted as an unhandled stream type (which surfaced as a spurious +// "unhandled stream types: file-history-delta=15" warning on real sessions). +func TestReadStreamJSON_FileHistoryDeltaIgnored(t *testing.T) { + ResetUnhandledStreamTypes() + input := `{"type":"file-history-delta","messageId":"m1","snapshotMessageId":"s1","trackingPath":"/Users/x/.claude/plans/p.md","backup":{"backupFileName":null,"version":1,"backupTime":"2026-07-23T17:18:00Z","realParentDir":"/Users/x/.claude/plans"},"timestamp":"2026-07-23T17:18:00Z"}` + entries, err := ReadStreamJSON(strings.NewReader(input)) + if err != nil { + t.Fatalf("ReadStreamJSON failed: %v", err) + } + if len(entries) != 0 { + t.Fatalf("file-history-delta should emit no entries, got %d: %+v", len(entries), entries) + } + if got := SnapshotUnhandledStreamTypes(); got["file-history-delta"] != 0 { + t.Errorf("file-history-delta should be a known storage type, not counted unhandled: %v", got) + } +} + // TestReadStreamJSON_ContentSystemSubtypes verifies the content-bearing system // subtypes surface as synthetic rows (carrying their content) rather than being // dropped as unhandled. func TestReadStreamJSON_ContentSystemSubtypes(t *testing.T) { ResetUnhandledStreamTypes() + // A local_command whose content is a recognized command/output wrapper now + // surfaces as a structured claude_command(_output) event; only untagged + // content falls through to the generic LocalCommand row exercised here. input := `{"type":"system","subtype":"compact_boundary","content":"Conversation compacted","uuid":"c"} -{"type":"system","subtype":"local_command","content":"ok","uuid":"l"} +{"type":"system","subtype":"local_command","content":"cleared pending input","uuid":"l"} {"type":"system","subtype":"scheduled_task_fire","content":"resuming /loop","uuid":"s"} {"type":"system","subtype":"informational","content":"Remote Control disconnected","uuid":"i"}` @@ -289,7 +354,7 @@ func TestReadStreamJSON_ContentSystemSubtypes(t *testing.T) { } wantRows := map[string]string{ "CompactBoundary": "Conversation compacted", - "LocalCommand": "ok", + "LocalCommand": "cleared pending input", "ScheduledTaskFire": "resuming /loop", "Informational": "Remote Control disconnected", } @@ -319,6 +384,58 @@ func TestReadStreamJSON_ContentSystemSubtypes(t *testing.T) { } } +func TestReadStreamJSON_SystemApiErrorSurfaced(t *testing.T) { + ResetUnhandledStreamTypes() + input := `{"type":"system","subtype":"api_error","error":{"message":"rate limited","status":429},"retryInMs":1000,"retryAttempt":1,"maxRetries":3,"uuid":"api"}` + + entries, err := ReadStreamJSON(strings.NewReader(input)) + if err != nil { + t.Fatalf("ReadStreamJSON failed: %v", err) + } + if len(entries) != 1 { + t.Fatalf("expected 1 ApiError entry, got %d", len(entries)) + } + uses := entries[0].Message.GetToolUses() + if len(uses) != 1 || uses[0].Name != "ApiError" { + t.Fatalf("expected ApiError synthetic tool, got %+v", uses) + } + var in map[string]any + if err := json.Unmarshal(uses[0].Input, &in); err != nil { + t.Fatalf("unmarshal ApiError input: %v", err) + } + if in["retryAttempt"] != float64(1) || in["maxRetries"] != float64(3) { + t.Errorf("ApiError input missing retry fields: %v", in) + } + if got := SnapshotUnhandledStreamTypes(); len(got) != 0 { + t.Errorf("system/api_error should be handled, got unhandled: %v", got) + } +} + +func TestReadStreamJSON_WorktreeLifecycleEvents(t *testing.T) { + ResetUnhandledStreamTypes() + input := `{"type":"worktree-state","worktreeSession":{"worktreeName":"feature","worktreePath":"/repo","worktreeBranch":"feature/x"},"uuid":"w"} +{"type":"relocated","relocatedCwd":"/repo/subdir","uuid":"r"} +{"type":"started","cwd":"/repo","uuid":"s"}` + + entries, err := ReadStreamJSON(strings.NewReader(input)) + if err != nil { + t.Fatalf("ReadStreamJSON failed: %v", err) + } + if len(entries) != 3 { + t.Fatalf("expected 3 lifecycle entries, got %d", len(entries)) + } + want := []string{"WorktreeState", "Relocated", "Started"} + for i, name := range want { + uses := entries[i].Message.GetToolUses() + if len(uses) != 1 || uses[0].Name != name { + t.Fatalf("entry[%d] expected %s, got %+v", i, name, uses) + } + } + if got := SnapshotUnhandledStreamTypes(); len(got) != 0 { + t.Errorf("worktree lifecycle events should be handled, got unhandled: %v", got) + } +} + func TestReadHistory_SessionFileEvents(t *testing.T) { // On-disk session files use camelCase fields and a different mix of // types from stream-json. ReadHistory must recognize the same set of diff --git a/pkg/claude/session.go b/pkg/claude/session.go index 2e8b3de..4aad22a 100644 --- a/pkg/claude/session.go +++ b/pkg/claude/session.go @@ -1,13 +1,14 @@ package claude import ( - "github.com/segmentio/encoding/json" + "bufio" "os" "path/filepath" "strings" "time" "github.com/flanksource/captain/pkg/claude/tools" + "github.com/segmentio/encoding/json" ) // GetClaudeHome returns the path to the Claude Code home directory (~/.claude) @@ -63,22 +64,28 @@ func FindProjectRoot(dir string) string { // For git worktrees (where .git is a file), Root is the worktree directory (for correct // relative paths) and MainRoot is the main repository root (for project naming). func FindProjectInfo(dir string) ProjectInfo { - if dir == "" { + if dir == "" || strings.Contains(dir, "..") { return ProjectInfo{} } current := dir for { - for _, marker := range projectMarkers { - markerPath := filepath.Join(current, marker) - info, err := os.Stat(markerPath) - if err != nil { - continue + entries, err := os.ReadDir(current) + if err == nil { + entriesByName := make(map[string]os.DirEntry) + for _, entry := range entries { + entriesByName[entry.Name()] = entry } - pi := ProjectInfo{Root: current, MarkerFile: marker} - if marker == ".git" && !info.IsDir() { - pi.MainRoot = resolveWorktreeRoot(markerPath) + for _, marker := range projectMarkers { + entry, ok := entriesByName[marker] + if !ok || entry.Type()&os.ModeSymlink != 0 { + continue + } + pi := ProjectInfo{Root: current, MarkerFile: marker} + if marker == ".git" && !entry.IsDir() { + pi.MainRoot = resolveWorktreeRoot(filepath.Join(current, marker)) + } + return pi } - return pi } parent := filepath.Dir(current) if parent == current { @@ -95,11 +102,28 @@ func FindProjectInfo(dir string) ProjectInfo { // // This function extracts the main repo root from that path. func resolveWorktreeRoot(gitFilePath string) string { - data, err := os.ReadFile(gitFilePath) + if filepath.Base(gitFilePath) != ".git" { + return "" + } + root, err := os.OpenRoot(filepath.Dir(gitFilePath)) if err != nil { return "" } - line := strings.TrimSpace(string(data)) + defer func() { _ = root.Close() }() + file, err := root.Open(".git") + if err != nil { + return "" + } + defer func() { _ = file.Close() }() + + scanner := bufio.NewScanner(file) + if !scanner.Scan() { + return "" + } + line := strings.TrimSpace(scanner.Text()) + if scanner.Scan() || scanner.Err() != nil { + return "" + } if !strings.HasPrefix(line, "gitdir: ") { return "" } diff --git a/pkg/claude/session_security_ginkgo_test.go b/pkg/claude/session_security_ginkgo_test.go new file mode 100644 index 0000000..ee9aa98 --- /dev/null +++ b/pkg/claude/session_security_ginkgo_test.go @@ -0,0 +1,54 @@ +package claude + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("worktree metadata", func() { + It("rejects parent traversal in project discovery", func() { + root := GinkgoT().TempDir() + project := filepath.Join(root, "project") + outside := filepath.Join(root, "outside") + Expect(os.MkdirAll(project, 0o755)).To(Succeed()) + Expect(os.MkdirAll(outside, 0o755)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(outside, "go.mod"), []byte("module example.com/outside\n"), 0o600)).To(Succeed()) + + info := FindProjectInfo(project + string(os.PathSeparator) + ".." + string(os.PathSeparator) + "outside") + + Expect(info.MarkerFile).To(BeEmpty()) + Expect(info.Root).To(BeEmpty()) + }) + + It("does not treat a marker symlink outside the project as a match", func() { + root := GinkgoT().TempDir() + project := filepath.Join(root, "project") + Expect(os.MkdirAll(project, 0o755)).To(Succeed()) + outside := filepath.Join(root, "outside-go.mod") + Expect(os.WriteFile(outside, []byte("module example.com/outside\n"), 0o600)).To(Succeed()) + Expect(os.Symlink(outside, filepath.Join(project, "go.mod"))).To(Succeed()) + + Expect(FindProjectInfo(project)).NotTo(SatisfyAll( + HaveField("Root", project), + HaveField("MarkerFile", "go.mod"), + )) + }) + + It("does not follow a .git symlink outside the worktree", func() { + root := GinkgoT().TempDir() + mainRepo := filepath.Join(root, "main") + gitDir := filepath.Join(mainRepo, ".git", "worktrees", "feature") + Expect(os.MkdirAll(gitDir, 0o755)).To(Succeed()) + + worktree := filepath.Join(root, "worktree") + Expect(os.MkdirAll(worktree, 0o755)).To(Succeed()) + outside := filepath.Join(root, "outside-git-file") + Expect(os.WriteFile(outside, []byte("gitdir: "+gitDir+"\n"), 0o600)).To(Succeed()) + Expect(os.Symlink(outside, filepath.Join(worktree, ".git"))).To(Succeed()) + + Expect(resolveWorktreeRoot(filepath.Join(worktree, ".git"))).To(BeEmpty()) + }) +}) diff --git a/pkg/claude/tools.go b/pkg/claude/tools.go index 24209d4..efdf162 100644 --- a/pkg/claude/tools.go +++ b/pkg/claude/tools.go @@ -59,18 +59,6 @@ type WebSearch struct { Query string `json:"query"` } -// TodoWrite represents a TodoWrite tool invocation -type TodoWrite struct { - Todos []Todo `json:"todos"` -} - -// Todo represents a single todo item -type Todo struct { - ActiveForm string `json:"activeForm,omitempty"` - Content string `json:"content,omitempty"` - Status string `json:"status,omitempty"` -} - // AskUserQuestion represents question data type Questions struct { Questions []Question `json:"questions"` diff --git a/pkg/claude/tools/agent.go b/pkg/claude/tools/agent.go index eb45da5..b02ea22 100644 --- a/pkg/claude/tools/agent.go +++ b/pkg/claude/tools/agent.go @@ -75,10 +75,10 @@ func (t *TodoWriteTool) Detail() api.Textable { text := clicky.Text("").Append("Plan", "font-bold") for _, item := range items { text = text.NewLine().Append("- ", "text-gray-500") - if item.status != "" { - text = text.Append(item.status+": ", "text-gray-500") + if item.Status != "" { + text = text.Append(item.Status+": ", "text-gray-500") } - text = text.Append(item.text, "") + text = text.Append(item.Text, "") } return &text } @@ -93,21 +93,28 @@ func (t *TodoWriteTool) Pretty() api.Text { return text } -type todoWriteItem struct { - text string - status string +// TodoItem is one entry of an agent task list, normalized across the shapes the +// providers emit: Claude's TodoWrite uses content/activeForm, Codex's +// update_plan uses step. Persisted into captain_sessions.metadata["todos"], so +// the JSON tags are a storage contract. +type TodoItem struct { + Text string `json:"text"` + Status string `json:"status,omitempty"` } -func (t *TodoWriteTool) todoItems() []todoWriteItem { +func (t *TodoWriteTool) todoItems() []TodoItem { raw := t.Input["todos"] if raw == nil { raw = t.Input["plan"] } - return todoItems(raw) + return TodoItems(raw) } -func todoItems(raw any) []todoWriteItem { - var out []todoWriteItem +// TodoItems normalizes a raw TodoWrite/update_plan payload into task items. It +// tolerates both []any and []map[string]any, and both provider key vocabularies. +// Entries without usable text are skipped rather than yielding empty items. +func TodoItems(raw any) []TodoItem { + var out []TodoItem switch todos := raw.(type) { case []any: for _, todo := range todos { @@ -125,20 +132,20 @@ func todoItems(raw any) []todoWriteItem { return out } -func todoItem(raw any) (todoWriteItem, bool) { +func todoItem(raw any) (TodoItem, bool) { m, ok := raw.(map[string]any) if !ok { - return todoWriteItem{}, false + return TodoItem{}, false } text := firstString(m, "step", "content", "activeForm", "title", "description") text = compactText(text) if text == "" { - return todoWriteItem{}, false + return TodoItem{}, false } - return todoWriteItem{ - text: truncateText(text, 160), - status: compactText(firstString(m, "status")), + return TodoItem{ + Text: truncateText(text, 160), + Status: compactText(firstString(m, "status")), }, true } @@ -151,16 +158,16 @@ func firstString(m map[string]any, keys ...string) string { return "" } -func todoPreview(items []todoWriteItem, max int) string { +func todoPreview(items []TodoItem, max int) string { if len(items) == 0 || max <= 0 { return "" } parts := make([]string, 0, max) for _, item := range items { - if item.text == "" { + if item.Text == "" { continue } - parts = append(parts, truncateText(item.text, 60)) + parts = append(parts, truncateText(item.Text, 60)) if len(parts) == max { break } diff --git a/pkg/claude/tools/assistant.go b/pkg/claude/tools/assistant.go index 4ff1c0e..7152803 100644 --- a/pkg/claude/tools/assistant.go +++ b/pkg/claude/tools/assistant.go @@ -10,7 +10,7 @@ import ( type AssistantTool struct{ BaseTool } func (t *AssistantTool) Name() string { return "Assistant" } -func (t *AssistantTool) Category() string { return "message" } +func (t *AssistantTool) Category() string { return "chat" } func (t *AssistantTool) FilePath() string { return "" } func (t *AssistantTool) ExtractPath() string { return "" } @@ -18,7 +18,7 @@ func (t *AssistantTool) Pretty() api.Text { icon := icons.Icon{Unicode: "🤖", Iconify: "mdi:robot", Style: "muted"} text := t.header(icon, "assistant", "text-blue-500 font-medium") if body := t.Str("text"); body != "" { - text = text.Append(" "+body, "text-gray-700 max-w-[tw-20ch]") + text = text.Append(" "+body, "text-muted max-w-[tw-20ch]") } return text } @@ -30,7 +30,7 @@ func (t *AssistantTool) Detail() api.Textable { return t.BaseTool.Detail() } type ReasoningTool struct{ BaseTool } func (t *ReasoningTool) Name() string { return "Reasoning" } -func (t *ReasoningTool) Category() string { return "message" } +func (t *ReasoningTool) Category() string { return "chat" } func (t *ReasoningTool) FilePath() string { return "" } func (t *ReasoningTool) ExtractPath() string { return "" } @@ -38,7 +38,7 @@ func (t *ReasoningTool) Pretty() api.Text { icon := icons.Icon{Unicode: "💭", Iconify: "mdi:thought-bubble", Style: "muted"} text := t.header(icon, "reasoning", "text-purple-500 font-medium") if body := t.Str("text"); body != "" { - text = text.Append(" "+body, "text-gray-500 italic max-w-[tw-20ch]") + text = text.Append(" "+body, "text-muted italic max-w-[tw-20ch]") } return text } diff --git a/pkg/claude/tools/claude_command_event_test.go b/pkg/claude/tools/claude_command_event_test.go new file mode 100644 index 0000000..f10cf46 --- /dev/null +++ b/pkg/claude/tools/claude_command_event_test.go @@ -0,0 +1,126 @@ +package tools + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("ClaudeCommand and GoalStatus renderers", func() { + Describe("ClaudeCommandTool", func() { + It("renders a slash-command invocation with an argument preview", func() { + tool := NewTool(BaseTool{ + RawTool: "ClaudeCommand", + Input: map[string]any{ + "event": "claude_command", + "scope": "turn", + "command_name": "/goal", + "command_message": "goal", + "command_args": "ship the docker build on PR #32", + }, + }) + Expect(tool).To(BeAssignableToTypeOf(&ClaudeCommandTool{})) + Expect(tool.Name()).To(Equal("ClaudeCommand")) + Expect(tool.Category()).To(Equal("chat")) + + pretty := tool.Pretty().String() + Expect(pretty).To(ContainSubstring("/goal")) + Expect(pretty).To(ContainSubstring("ship the docker build on PR #32")) + + detail := tool.Detail() + Expect(detail).NotTo(BeNil()) + Expect(detail.String()).To(ContainSubstring("ship the docker build on PR #32")) + Expect(detail.String()).NotTo(ContainSubstring("")) + }) + + It("renders stdout output with the stream label and preview, wrapper-free", func() { + tool := NewTool(BaseTool{ + RawTool: "ClaudeCommand", + Input: map[string]any{ + "event": "claude_command_output", + "scope": "turn", + "stream": "stdout", + "content": "Goal set: ship it", + }, + }) + Expect(tool).To(BeAssignableToTypeOf(&ClaudeCommandTool{})) + + pretty := tool.Pretty().String() + Expect(pretty).To(ContainSubstring("stdout")) + Expect(pretty).To(ContainSubstring("Goal set: ship it")) + Expect(pretty).NotTo(ContainSubstring("")) + + detail := tool.Detail() + Expect(detail).NotTo(BeNil()) + Expect(detail.String()).To(ContainSubstring("Goal set: ship it")) + Expect(detail.String()).NotTo(ContainSubstring("")) + }) + + It("renders stderr output with the stderr stream label", func() { + tool := NewTool(BaseTool{ + RawTool: "ClaudeCommand", + Input: map[string]any{ + "event": "claude_command_output", + "scope": "turn", + "stream": "stderr", + "content": "command failed", + }, + }) + pretty := tool.Pretty().String() + Expect(pretty).To(ContainSubstring("stderr")) + Expect(pretty).To(ContainSubstring("command failed")) + }) + }) + + Describe("GoalStatusTool", func() { + It("renders an active goal with a condition preview", func() { + tool := NewTool(BaseTool{ + RawTool: "GoalStatus", + Input: map[string]any{ + "event": "goal_status", + "scope": "session", + "met": false, + "sentinel": true, + "condition": "release once CI is green", + }, + }) + Expect(tool).To(BeAssignableToTypeOf(&GoalStatusTool{})) + Expect(tool.Name()).To(Equal("GoalStatus")) + Expect(tool.Category()).To(Equal("chat")) + + pretty := tool.Pretty().String() + Expect(pretty).To(ContainSubstring("goal active")) + Expect(pretty).To(ContainSubstring("release once CI is green")) + + detail := tool.Detail() + Expect(detail).NotTo(BeNil()) + Expect(detail.String()).To(ContainSubstring("release once CI is green")) + }) + + It("renders a met goal", func() { + tool := NewTool(BaseTool{ + RawTool: "GoalStatus", + Input: map[string]any{ + "event": "goal_status", + "met": true, + "condition": "release once CI is green", + }, + }) + Expect(tool.Pretty().String()).To(ContainSubstring("goal met")) + }) + + It("includes a failure reason in the detail", func() { + tool := NewTool(BaseTool{ + RawTool: "GoalStatus", + Input: map[string]any{ + "event": "goal_status", + "met": false, + "condition": "release once CI is green", + "reason": "docker build still failing", + }, + }) + detail := tool.Detail() + Expect(detail).NotTo(BeNil()) + Expect(detail.String()).To(ContainSubstring("docker build still failing")) + }) + }) +}) diff --git a/pkg/claude/tools/event.go b/pkg/claude/tools/event.go new file mode 100644 index 0000000..1a63ea3 --- /dev/null +++ b/pkg/claude/tools/event.go @@ -0,0 +1,203 @@ +package tools + +import ( + "fmt" + "strings" + + "github.com/flanksource/clicky" + "github.com/flanksource/clicky/api" + "github.com/flanksource/clicky/api/icons" +) + +// EventToolName maps raw Codex/Claude event names to concrete synthetic tools. +// The raw event name is still preserved in BaseTool.Input["event"]. +func EventToolName(eventType string) string { + switch strings.TrimSpace(eventType) { + case "token_count": + return "TokenCount" + case "task_started": + return "TaskStarted" + case "task_complete": + return "TaskComplete" + case "turn_aborted": + return "TurnAborted" + case "context_compacted": + return "ContextCompacted" + case "thread_rolled_back": + return "ThreadRolledBack" + case "item_completed": + return "ItemCompleted" + case "exec_command_end": + return "CodexExecCommand" + case "patch_apply_end": + return "CodexPatchApply" + case "mcp_tool_call_end": + return "MCPToolCall" + case "web_search_end": + return "WebSearchEvent" + case "view_image_tool_call": + return "ViewImage" + case "guardian_assessment": + return "GuardianAssessment" + case "entered_review_mode", "exited_review_mode": + return "ReviewMode" + case "collab_agent_spawn_end": + return "CollabAgentSpawn" + case "collab_agent_interaction_end": + return "CollabAgentInteraction" + case "collab_waiting_end": + return "CollabWaiting" + case "collab_close_end": + return "CollabClose" + case "error": + return "ApiError" + case "queue-operation": + return "QueueOperation" + case "deferred_tools_delta": + return "DeferredToolsDelta" + case "agent_listing_delta": + return "AgentListingDelta" + case "memory_citation": + return "MemoryCitation" + case "skill_listing": + return "SkillListing" + case "budget_usd": + return "Budget" + case "worktree-state": + return "WorktreeState" + case "relocated": + return "Relocated" + case "started": + return "Started" + case "claude_command", "claude_command_output": + return "ClaudeCommand" + case "goal_status": + return "GoalStatus" + default: + return "Event" + } +} + +// IsEventToolName reports whether name is a synthetic transcript event row. +func IsEventToolName(name string) bool { + switch name { + case "Event", "TokenCount", "TaskStarted", "TaskComplete", "TurnAborted", + "ContextCompacted", "ThreadRolledBack", "ItemCompleted", + "CodexExecCommand", "CodexPatchApply", "MCPToolCall", + "WebSearchEvent", "ViewImage", "GuardianAssessment", "ReviewMode", + "CollabAgentSpawn", "CollabAgentInteraction", "CollabWaiting", + "CollabClose", "QueueOperation", "DeferredToolsDelta", + "AgentListingDelta", "MemoryCitation", "SkillListing", "Budget", "PrLink", + "CompactBoundary", "LocalCommand", "ScheduledTaskFire", + "Informational", "WorktreeState", "Relocated", "Started", "UserShellCommand", + "ClaudeCommand", "GoalStatus": + return true + default: + return false + } +} + +type EventTool struct{ BaseTool } + +func (t *EventTool) Name() string { return "Event" } +func (t *EventTool) Category() string { return "chat" } +func (t *EventTool) FilePath() string { return "" } +func (t *EventTool) ExtractPath() string { return "" } + +func (t *EventTool) Pretty() api.Text { + name := t.Str("event") + if name == "" { + name = "event" + } + text := clicky.Text(""). + Add(icons.Info). + Append(" "+name, "text-slate-500 font-medium") + if msg := t.Str("message"); msg != "" { + text = text.Append(" "+msg, "text-gray-500 max-w-[tw-20ch]") + } + if duration := eventInt(t.Input["duration_ms"]); duration > 0 { + text = text.Append(fmt.Sprintf(" %dms", duration), "text-gray-500") + } + return text +} + +func (t *EventTool) Detail() api.Textable { return t.BaseTool.Detail() } + +type TokenCountTool struct{ BaseTool } + +func (t *TokenCountTool) Name() string { return "TokenCount" } +func (t *TokenCountTool) Category() string { return "chat" } +func (t *TokenCountTool) FilePath() string { return "" } +func (t *TokenCountTool) ExtractPath() string { return "" } +func (t *TokenCountTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *TokenCountTool) Pretty() api.Text { + text := eventText(icons.Icon{Unicode: "◌", Iconify: "mdi:counter", Style: "muted"}, "tokens", "text-slate-500 font-medium") + if total := eventInt(t.Input["total_tokens"]); total > 0 { + text = text.Append(" "+FormatCompactTokens(int(total)), "text-muted") + } + if in := eventInt(t.Input["input_tokens"]); in > 0 { + text = text.Append(fmt.Sprintf(" in=%s", FormatCompactTokens(int(in))), "text-gray-500") + } + if out := eventInt(t.Input["output_tokens"]); out > 0 { + text = text.Append(fmt.Sprintf(" out=%s", FormatCompactTokens(int(out))), "text-gray-500") + } + return text +} + +type TaskStartedTool struct{ BaseTool } + +func (t *TaskStartedTool) Name() string { return "TaskStarted" } +func (t *TaskStartedTool) Category() string { return "chat" } +func (t *TaskStartedTool) FilePath() string { return "" } +func (t *TaskStartedTool) ExtractPath() string { return "" } +func (t *TaskStartedTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *TaskStartedTool) Pretty() api.Text { + text := eventText(icons.Play, "task started", "text-green-600 font-medium") + if mode := t.Str("collaboration_mode_kind"); mode != "" { + text = text.Append(" "+mode, "text-gray-500") + } + if window := eventInt(t.Input["model_context_window"]); window > 0 { + text = text.Append(fmt.Sprintf(" ctx=%s", FormatCompactTokens(int(window))), "text-gray-500") + } + return text +} + +type TaskCompleteTool struct{ BaseTool } + +func (t *TaskCompleteTool) Name() string { return "TaskComplete" } +func (t *TaskCompleteTool) Category() string { return "chat" } +func (t *TaskCompleteTool) FilePath() string { return "" } +func (t *TaskCompleteTool) ExtractPath() string { return "" } +func (t *TaskCompleteTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *TaskCompleteTool) Pretty() api.Text { + text := eventText(icons.Check, "task complete", "text-green-600 font-medium") + if duration := eventInt(t.Input["duration_ms"]); duration > 0 { + text = text.Append(" "+formatDurationMS(float64(duration)), "text-gray-500") + } + if msg := eventString(t.Input["last_agent_message"]); msg != "" { + text = text.Append(" "+eventPreview(msg, 90), "text-gray-500 italic") + } + return text +} + +type LifecycleEventTool struct{ BaseTool } + +func (t *LifecycleEventTool) Name() string { return t.RawTool } +func (t *LifecycleEventTool) Category() string { return "chat" } +func (t *LifecycleEventTool) FilePath() string { return "" } +func (t *LifecycleEventTool) ExtractPath() string { return "" } +func (t *LifecycleEventTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *LifecycleEventTool) Pretty() api.Text { + label := strings.TrimSpace(t.Str("event")) + if label == "" { + label = strings.ToLower(t.RawTool) + } + text := eventText(icons.Info, strings.ReplaceAll(label, "_", " "), "text-slate-500 font-medium") + if reason := t.Str("reason"); reason != "" { + text = text.Append(" "+reason, "text-gray-500") + } + if turns := eventInt(t.Input["num_turns"]); turns > 0 { + text = text.Append(fmt.Sprintf(" turns=%d", turns), "text-gray-500") + } + return text +} diff --git a/pkg/claude/tools/event_helpers.go b/pkg/claude/tools/event_helpers.go new file mode 100644 index 0000000..5fd3b9a --- /dev/null +++ b/pkg/claude/tools/event_helpers.go @@ -0,0 +1,239 @@ +package tools + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/flanksource/clicky" + "github.com/flanksource/clicky/api" +) + +func eventInt(value any) int64 { + switch v := value.(type) { + case int: + return int64(v) + case int64: + return v + case float64: + return int64(v) + default: + return 0 + } +} + +func eventFloat(value any) float64 { + switch v := value.(type) { + case int: + return float64(v) + case int64: + return float64(v) + case float64: + return v + default: + return 0 + } +} + +func eventString(value any) string { + switch v := value.(type) { + case string: + return v + default: + return "" + } +} + +func eventText(icon api.Textable, label, color string) api.Text { + return clicky.Text("").Add(icon).Append(" "+label, color) +} + +func eventPreview(s string, max int) string { + s = strings.Join(strings.Fields(strings.TrimSpace(s)), " ") + if max <= 0 || len(s) <= max { + return s + } + if max <= 3 { + return s[:max] + } + return s[:max-3] + "..." +} + +func firstNonEmptyEvent(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} + +func mapLen(value any) int { + switch v := value.(type) { + case map[string]any: + return len(v) + case map[string]map[string]any: + return len(v) + default: + return 0 + } +} + +func listLen(value any) int { + switch v := value.(type) { + case []any: + return len(v) + case []string: + return len(v) + default: + return 0 + } +} + +func appendListCount(text *api.Text, label string, value any) { + if n := listLen(value); n > 0 { + *text = text.Append(fmt.Sprintf(" %s=%d", label, n), "text-gray-500") + } +} + +func statusColor(status string) string { + switch strings.ToLower(status) { + case "completed", "success", "ok", "approved", "pass", "passed": + return "text-green-600" + case "failed", "error", "denied", "blocked", "rejected", "interrupted": + return "text-red-500" + default: + return "text-gray-500" + } +} + +func codexCommandString(value any) string { + switch v := value.(type) { + case string: + return v + case []any: + parts := make([]string, 0, len(v)) + for _, item := range v { + if s, ok := item.(string); ok { + parts = append(parts, s) + } + } + return strings.Join(parts, " ") + case []string: + return strings.Join(v, " ") + default: + return "" + } +} + +func commandOutputDetail(base BaseTool) api.Textable { + if d := base.Detail(); d != nil { + return d + } + stdout := strings.TrimSpace(base.Str("stdout")) + stderr := strings.TrimSpace(base.Str("stderr")) + if stdout == "" && stderr == "" { + stdout = strings.TrimSpace(base.Str("aggregated_output")) + } + if stdout == "" && stderr == "" { + return nil + } + text := clicky.Text("") + if stdout != "" { + text = text.Append("stdout: ", "font-bold text-muted").Append(stdout, "") + } + if stderr != "" { + if stdout != "" { + text = text.NewLine() + } + text = text.Append("stderr: ", "font-bold text-red-500").Append(stderr, "") + } + return &text +} + +func invocationName(value any) (string, string) { + m, ok := value.(map[string]any) + if !ok { + return "", "" + } + return eventString(m["server"]), eventString(m["tool"]) +} + +func durationString(value any) string { + m, ok := value.(map[string]any) + if !ok { + return "" + } + if secs := eventFloat(m["secs"]); secs > 0 { + return formatDurationMS(secs * 1000) + } + if nanos := eventFloat(m["nanos"]); nanos > 0 { + return formatDurationMS(nanos / 1_000_000) + } + return "" +} + +func resultStatus(value any) string { + m, ok := value.(map[string]any) + if !ok { + return "" + } + for k := range m { + if strings.EqualFold(k, "ok") { + return "ok" + } + if strings.EqualFold(k, "err") || strings.EqualFold(k, "error") { + return "error" + } + } + return "" +} + +func actionQuery(value any) string { + m, ok := value.(map[string]any) + if !ok { + return "" + } + if q := eventString(m["query"]); q != "" { + return q + } + if qs, ok := m["queries"].([]any); ok && len(qs) > 0 { + if q, ok := qs[0].(string); ok { + return q + } + } + return "" +} + +func actionSummary(value any) string { + m, ok := value.(map[string]any) + if !ok { + return "" + } + parts := make([]string, 0, 3) + for _, key := range []string{"tool", "command", "cwd"} { + if value := eventString(m[key]); value != "" { + if key == "cwd" { + value = filepath.Base(value) + } + parts = append(parts, value) + } + } + return strings.Join(parts, " ") +} + +func findingsCount(value any) int { + m, ok := value.(map[string]any) + if !ok { + return 0 + } + return listLen(m["findings"]) +} + +func nestedString(value any, key string) string { + m, ok := value.(map[string]any) + if !ok { + return "" + } + return eventString(m[key]) +} diff --git a/pkg/claude/tools/event_renderers.go b/pkg/claude/tools/event_renderers.go new file mode 100644 index 0000000..e41230c --- /dev/null +++ b/pkg/claude/tools/event_renderers.go @@ -0,0 +1,469 @@ +package tools + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/bash" + "github.com/flanksource/clicky" + "github.com/flanksource/clicky/api" + "github.com/flanksource/clicky/api/icons" +) + +type CodexExecCommandTool struct{ BaseTool } + +func (t *CodexExecCommandTool) Name() string { return "CodexExecCommand" } +func (t *CodexExecCommandTool) Category() string { return "chat" } +func (t *CodexExecCommandTool) FilePath() string { return "" } +func (t *CodexExecCommandTool) ExtractPath() string { + if cwd := t.Str("cwd"); cwd != "" { + return cwd + } + return "" +} +func (t *CodexExecCommandTool) Detail() api.Textable { return commandOutputDetail(t.BaseTool) } +func (t *CodexExecCommandTool) Pretty() api.Text { + text := eventText(icons.Icon{Unicode: "💻", Iconify: "codicon:terminal", Style: "muted"}, "exec", "text-green-500 font-medium") + if cmd := codexCommandString(t.Input["command"]); cmd != "" { + text = text.Append(" "+eventPreview(cmd, 120), "text-muted") + } + if status := t.Str("status"); status != "" { + color := "text-green-600" + if status != "completed" && status != "success" { + color = "text-red-500" + } + text = text.Append(" "+status, color) + } + if code := eventInt(t.Input["exit_code"]); code != 0 { + text = text.Append(fmt.Sprintf(" exit=%d", code), "text-red-500") + } + return text +} + +type UserShellCommandTool struct{ BaseTool } + +func (t *UserShellCommandTool) Name() string { return "UserShellCommand" } +func (t *UserShellCommandTool) Category() string { return "chat" } +func (t *UserShellCommandTool) FilePath() string { return "" } +func (t *UserShellCommandTool) ExtractPath() string { return "" } +func (t *UserShellCommandTool) Detail() api.Textable { return commandOutputDetail(t.BaseTool) } +func (t *UserShellCommandTool) Pretty() api.Text { + text := eventText(icons.Icon{Unicode: "💻", Iconify: "codicon:terminal", Style: "muted"}, "local command", "text-slate-500 font-medium") + if cmd := t.Str("command"); cmd != "" { + text = text.Append(" "+eventPreview(cmd, 120), "text-muted") + } + if duration := eventFloat(t.Input["duration_ms"]); duration > 0 { + text = text.Append(" "+formatDurationMS(duration), "text-gray-500") + } + if code := eventInt(t.Input["exit_code"]); code != 0 { + text = text.Append(fmt.Sprintf(" exit=%d", code), "text-red-500") + } + return text +} + +type CodexPatchApplyTool struct{ BaseTool } + +func (t *CodexPatchApplyTool) Name() string { return "CodexPatchApply" } + +// Category is edit, not chat: this row is Codex applying a patch, i.e. a file +// mutation. The other event renderers here are genuinely conversational, but +// grouping a write with them hides Codex file changes from `--category edit` +// and from anything sourcing file modifications off the category. +func (t *CodexPatchApplyTool) Category() string { return string(bash.CategoryEdit) } +func (t *CodexPatchApplyTool) FilePath() string { return "" } +func (t *CodexPatchApplyTool) ExtractPath() string { return "" } +func (t *CodexPatchApplyTool) Detail() api.Textable { return commandOutputDetail(t.BaseTool) } +func (t *CodexPatchApplyTool) Pretty() api.Text { + text := eventText(icons.Icon{Unicode: "✏️", Iconify: "codicon:edit", Style: "muted"}, "patch", "text-orange-500 font-medium") + if success, ok := t.Input["success"].(bool); ok { + if success { + text = text.Append(" applied", "text-green-600") + } else { + text = text.Append(" failed", "text-red-500") + } + } + if n := mapLen(t.Input["changes"]); n > 0 { + text = text.Append(fmt.Sprintf(" files=%d", n), "text-gray-500") + } + return text +} + +type MCPToolCallTool struct{ BaseTool } + +func (t *MCPToolCallTool) Name() string { return "MCPToolCall" } +func (t *MCPToolCallTool) Category() string { return "chat" } +func (t *MCPToolCallTool) FilePath() string { return "" } +func (t *MCPToolCallTool) ExtractPath() string { return "" } +func (t *MCPToolCallTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *MCPToolCallTool) Pretty() api.Text { + server, tool := invocationName(t.Input["invocation"]) + text := eventText(icons.Package, "mcp", "text-indigo-500 font-medium") + if server != "" || tool != "" { + text = text.Append(" "+strings.Trim(server+"."+tool, "."), "text-muted") + } + if d := durationString(t.Input["duration"]); d != "" { + text = text.Append(" "+d, "text-gray-500") + } + if resultStatus := resultStatus(t.Input["result"]); resultStatus != "" { + text = text.Append(" "+resultStatus, statusColor(resultStatus)) + } + return text +} + +type WebSearchEventTool struct{ BaseTool } + +func (t *WebSearchEventTool) Name() string { return "WebSearchEvent" } +func (t *WebSearchEventTool) Category() string { return "chat" } +func (t *WebSearchEventTool) FilePath() string { return "" } +func (t *WebSearchEventTool) ExtractPath() string { return "" } +func (t *WebSearchEventTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *WebSearchEventTool) Pretty() api.Text { + text := eventText(icons.Search, "web search", "text-purple-500 font-medium") + if query := firstNonEmptyEvent(t.Str("query"), actionQuery(t.Input["action"])); query != "" { + text = text.Append(" "+eventPreview(query, 100), "text-muted") + } + return text +} + +type ViewImageTool struct{ BaseTool } + +func (t *ViewImageTool) Name() string { return "ViewImage" } +func (t *ViewImageTool) Category() string { return "chat" } +func (t *ViewImageTool) FilePath() string { return t.Str("path") } +func (t *ViewImageTool) ExtractPath() string { return t.Str("path") } +func (t *ViewImageTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *ViewImageTool) Pretty() api.Text { + text := eventText(icons.File, "image", "text-cyan-500 font-medium") + if path := t.Str("path"); path != "" { + text = text.Append(" "+ShortenPath(path), "text-muted") + } + return text +} + +type GuardianAssessmentTool struct{ BaseTool } + +func (t *GuardianAssessmentTool) Name() string { return "GuardianAssessment" } +func (t *GuardianAssessmentTool) Category() string { return "chat" } +func (t *GuardianAssessmentTool) FilePath() string { return "" } +func (t *GuardianAssessmentTool) ExtractPath() string { return "" } +func (t *GuardianAssessmentTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *GuardianAssessmentTool) Pretty() api.Text { + text := eventText(icons.Icon{Unicode: "🛡", Iconify: "mdi:shield-check", Style: "muted"}, "guardian", "text-amber-600 font-medium") + if status := t.Str("status"); status != "" { + text = text.Append(" "+status, statusColor(status)) + } + if action := actionSummary(t.Input["action"]); action != "" { + text = text.Append(" "+eventPreview(action, 100), "text-muted") + } + return text +} + +type ReviewModeTool struct{ BaseTool } + +func (t *ReviewModeTool) Name() string { return "ReviewMode" } +func (t *ReviewModeTool) Category() string { return "chat" } +func (t *ReviewModeTool) FilePath() string { return "" } +func (t *ReviewModeTool) ExtractPath() string { return "" } +func (t *ReviewModeTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *ReviewModeTool) Pretty() api.Text { + action := strings.TrimSuffix(strings.TrimPrefix(t.Str("event"), "entered_"), "_mode") + if action == "" { + action = strings.TrimSuffix(strings.TrimPrefix(t.Str("event"), "exited_"), "_mode") + } + text := eventText(icons.Info, "review "+action, "text-purple-500 font-medium") + if n := findingsCount(t.Input["review_output"]); n > 0 { + text = text.Append(fmt.Sprintf(" findings=%d", n), "text-gray-500") + } + return text +} + +type CollabEventTool struct{ BaseTool } + +func (t *CollabEventTool) Name() string { return t.RawTool } +func (t *CollabEventTool) Category() string { return "chat" } +func (t *CollabEventTool) FilePath() string { return "" } +func (t *CollabEventTool) ExtractPath() string { return "" } +func (t *CollabEventTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *CollabEventTool) Pretty() api.Text { + label := strings.TrimPrefix(strings.TrimPrefix(t.Str("event"), "collab_"), "agent_") + label = strings.TrimSuffix(label, "_end") + if label == "" { + label = strings.ToLower(t.RawTool) + } + text := eventText(icons.Icon{Unicode: "🤝", Iconify: "mdi:handshake", Style: "muted"}, "collab "+strings.ReplaceAll(label, "_", " "), "text-indigo-500 font-medium") + if nick := firstNonEmptyEvent(t.Str("nickname"), nestedString(t.Input["receiver"], "nickname")); nick != "" { + text = text.Append(" "+nick, "text-muted") + } + if status := t.Str("status"); status != "" { + text = text.Append(" "+status, statusColor(status)) + } + return text +} + +type QueueOperationTool struct{ BaseTool } + +func (t *QueueOperationTool) Name() string { return "QueueOperation" } +func (t *QueueOperationTool) Category() string { return "chat" } +func (t *QueueOperationTool) FilePath() string { return "" } +func (t *QueueOperationTool) ExtractPath() string { return "" } +func (t *QueueOperationTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *QueueOperationTool) Pretty() api.Text { + op := firstNonEmptyEvent(t.Str("operation"), "queue") + text := eventText(icons.Package, "queue "+op, "text-slate-500 font-medium") + if content := t.Str("content"); content != "" { + text = text.Append(" "+eventPreview(content, 100), "text-gray-500") + } + return text +} + +type DeferredToolsDeltaTool struct{ BaseTool } + +func (t *DeferredToolsDeltaTool) Name() string { return "DeferredToolsDelta" } +func (t *DeferredToolsDeltaTool) Category() string { return "chat" } +func (t *DeferredToolsDeltaTool) FilePath() string { return "" } +func (t *DeferredToolsDeltaTool) ExtractPath() string { return "" } +func (t *DeferredToolsDeltaTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *DeferredToolsDeltaTool) Pretty() api.Text { + text := eventText(icons.Package, "tools", "text-slate-500 font-medium") + appendListCount(&text, "added", t.Input["addedNames"]) + appendListCount(&text, "pending", t.Input["pendingMcpServers"]) + return text +} + +type AgentListingDeltaTool struct{ BaseTool } + +func (t *AgentListingDeltaTool) Name() string { return "AgentListingDelta" } +func (t *AgentListingDeltaTool) Category() string { return "chat" } +func (t *AgentListingDeltaTool) FilePath() string { return "" } +func (t *AgentListingDeltaTool) ExtractPath() string { return "" } +func (t *AgentListingDeltaTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *AgentListingDeltaTool) Pretty() api.Text { + text := eventText(icons.Icon{Unicode: "🤖", Iconify: "mdi:robot", Style: "muted"}, "agents", "text-indigo-500 font-medium") + appendListCount(&text, "added", t.Input["addedTypes"]) + return text +} + +type SkillListingTool struct{ BaseTool } + +func (t *SkillListingTool) Name() string { return "SkillListing" } +func (t *SkillListingTool) Category() string { return "chat" } +func (t *SkillListingTool) FilePath() string { return "" } +func (t *SkillListingTool) ExtractPath() string { return "" } +func (t *SkillListingTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *SkillListingTool) Pretty() api.Text { + text := eventText(icons.Info, "skills", "text-teal-500 font-medium") + if n := listLen(t.Input["names"]); n > 0 { + return text.Append(fmt.Sprintf(" count=%d", n), "text-gray-500") + } + if count := eventInt(t.Input["skillCount"]); count > 0 { + return text.Append(fmt.Sprintf(" count=%d", count), "text-gray-500") + } + return text +} + +type BudgetTool struct{ BaseTool } + +func (t *BudgetTool) Name() string { return "Budget" } +func (t *BudgetTool) Category() string { return "chat" } +func (t *BudgetTool) FilePath() string { return "" } +func (t *BudgetTool) ExtractPath() string { return "" } +func (t *BudgetTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *BudgetTool) Pretty() api.Text { + text := eventText(icons.Icon{Unicode: "$", Iconify: "mdi:cash", Style: "muted"}, "budget", "text-green-600 font-medium") + if used := eventFloat(t.Input["used"]); used > 0 { + text = text.Append(fmt.Sprintf(" used=$%.2f", used), "text-muted") + } + if total := eventFloat(t.Input["total"]); total > 0 { + text = text.Append(fmt.Sprintf(" total=$%.2f", total), "text-gray-500") + } + if remaining := eventFloat(t.Input["remaining"]); remaining > 0 { + text = text.Append(fmt.Sprintf(" remaining=$%.2f", remaining), "text-gray-500") + } + return text +} + +type PrLinkTool struct{ BaseTool } + +func (t *PrLinkTool) Name() string { return "PrLink" } +func (t *PrLinkTool) Category() string { return "chat" } +func (t *PrLinkTool) FilePath() string { return "" } +func (t *PrLinkTool) ExtractPath() string { return "" } +func (t *PrLinkTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *PrLinkTool) Pretty() api.Text { + text := eventText(icons.Git, "pr", "text-cyan-600 font-medium") + if n := eventInt(t.Input["prNumber"]); n > 0 { + text = text.Append(fmt.Sprintf(" #%d", n), "text-muted") + } + if repo := t.Str("prRepository"); repo != "" { + text = text.Append(" "+repo, "text-gray-500") + } + return text +} + +type ContentEventTool struct{ BaseTool } + +func (t *ContentEventTool) Name() string { return t.RawTool } +func (t *ContentEventTool) Category() string { return "chat" } +func (t *ContentEventTool) FilePath() string { return "" } +func (t *ContentEventTool) ExtractPath() string { return "" } +func (t *ContentEventTool) Detail() api.Textable { + if d := t.BaseTool.Detail(); d != nil { + return d + } + if c := strings.TrimSpace(t.Str("content")); c != "" { + text := clicky.Text("").Append(c, "") + return &text + } + return nil +} +func (t *ContentEventTool) Pretty() api.Text { + label := strings.TrimSpace(t.Str("event")) + if label == "" { + label = strings.ToLower(t.RawTool) + } + text := eventText(icons.Info, strings.ReplaceAll(label, "_", " "), "text-slate-500 font-medium") + if content := t.Str("content"); content != "" { + text = text.Append(" "+eventPreview(content, 100), "text-gray-500") + } + return text +} + +type WorktreeStateTool struct{ BaseTool } + +func (t *WorktreeStateTool) Name() string { return "WorktreeState" } +func (t *WorktreeStateTool) Category() string { return "chat" } +func (t *WorktreeStateTool) FilePath() string { + return nestedString(t.Input["worktreeSession"], "worktreePath") +} +func (t *WorktreeStateTool) ExtractPath() string { return t.FilePath() } +func (t *WorktreeStateTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *WorktreeStateTool) Pretty() api.Text { + text := eventText(icons.Git, "worktree", "text-green-600 font-medium") + if name := nestedString(t.Input["worktreeSession"], "worktreeName"); name != "" { + text = text.Append(" "+name, "text-muted") + } + if branch := nestedString(t.Input["worktreeSession"], "worktreeBranch"); branch != "" { + text = text.Append(" "+branch, "text-gray-500") + } + return text +} + +type RelocatedTool struct{ BaseTool } + +func (t *RelocatedTool) Name() string { return "Relocated" } +func (t *RelocatedTool) Category() string { return "chat" } +func (t *RelocatedTool) FilePath() string { return t.Str("relocatedCwd") } +func (t *RelocatedTool) ExtractPath() string { return t.Str("relocatedCwd") } +func (t *RelocatedTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *RelocatedTool) Pretty() api.Text { + text := eventText(icons.Folder, "relocated", "text-cyan-600 font-medium") + if cwd := t.Str("relocatedCwd"); cwd != "" { + text = text.Append(" "+ShortenPath(cwd), "text-muted") + } + return text +} + +type StartedTool struct{ BaseTool } + +func (t *StartedTool) Name() string { return "Started" } +func (t *StartedTool) Category() string { return "chat" } +func (t *StartedTool) FilePath() string { return "" } +func (t *StartedTool) ExtractPath() string { return "" } +func (t *StartedTool) Detail() api.Textable { return t.BaseTool.Detail() } +func (t *StartedTool) Pretty() api.Text { + return eventText(icons.Play, "started", "text-green-600 font-medium") +} + +// ClaudeCommandTool renders a Claude slash-command invocation (claude_command) +// or its captured output (claude_command_output) as a concise, non-operational +// event row. The wrapper tags are stripped by the reader; this only formats the +// preserved fields. +type ClaudeCommandTool struct{ BaseTool } + +func (t *ClaudeCommandTool) Name() string { return "ClaudeCommand" } +func (t *ClaudeCommandTool) Category() string { return "chat" } +func (t *ClaudeCommandTool) FilePath() string { return "" } +func (t *ClaudeCommandTool) ExtractPath() string { return "" } + +func (t *ClaudeCommandTool) isOutput() bool { + return t.Str("event") == "claude_command_output" || t.Str("stream") != "" +} + +func (t *ClaudeCommandTool) Pretty() api.Text { + if t.isOutput() { + stream := firstNonEmptyEvent(t.Str("stream"), "stdout") + color := "text-slate-500 font-medium" + if stream == "stderr" { + color = "text-red-500 font-medium" + } + text := eventText(icons.Terminal, stream, color) + if content := t.Str("content"); content != "" { + text = text.Append(" "+eventPreview(content, 100), "text-muted") + } + return text + } + name := firstNonEmptyEvent(t.Str("command_name"), "command") + text := eventText(icons.Terminal, name, "text-indigo-500 font-medium") + if args := t.Str("command_args"); args != "" { + text = text.Append(" "+eventPreview(args, 100), "text-muted") + } + return text +} + +func (t *ClaudeCommandTool) Detail() api.Textable { + if d := t.BaseTool.Detail(); d != nil { + return d + } + body := t.Str("command_args") + if t.isOutput() { + body = t.Str("content") + } + if strings.TrimSpace(body) == "" { + return nil + } + text := clicky.Text("").Append(body, "") + return &text +} + +// GoalStatusTool renders a session-scoped Claude goal directive as a concise, +// non-operational event row. +type GoalStatusTool struct{ BaseTool } + +func (t *GoalStatusTool) Name() string { return "GoalStatus" } +func (t *GoalStatusTool) Category() string { return "chat" } +func (t *GoalStatusTool) FilePath() string { return "" } +func (t *GoalStatusTool) ExtractPath() string { return "" } + +func (t *GoalStatusTool) Pretty() api.Text { + label, color := "goal active", "text-amber-600 font-medium" + if met, _ := t.Input["met"].(bool); met { + label, color = "goal met", "text-green-600 font-medium" + } + text := eventText(icons.Target, label, color) + if condition := t.Str("condition"); condition != "" { + text = text.Append(" "+eventPreview(condition, 100), "text-muted") + } + return text +} + +func (t *GoalStatusTool) Detail() api.Textable { + if d := t.BaseTool.Detail(); d != nil { + return d + } + condition := strings.TrimSpace(t.Str("condition")) + reason := strings.TrimSpace(t.Str("reason")) + if condition == "" && reason == "" { + return nil + } + text := clicky.Text("") + if condition != "" { + text = text.Append("condition: ", "font-bold text-muted").Append(condition, "") + } + if reason != "" { + if condition != "" { + text = text.NewLine() + } + text = text.Append("reason: ", "font-bold text-muted").Append(reason, "") + } + return &text +} diff --git a/pkg/claude/tools/muted_styles_test.go b/pkg/claude/tools/muted_styles_test.go new file mode 100644 index 0000000..4aeebc4 --- /dev/null +++ b/pkg/claude/tools/muted_styles_test.go @@ -0,0 +1,122 @@ +package tools + +import ( + "strings" + "testing" + + "github.com/flanksource/clicky/api" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestMutedStyles(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Muted Pretty Styles Suite") +} + +var _ = Describe("semantic muted pretty styles", func() { + DescribeTable("renders secondary session text with the semantic muted style", + func(render func() api.Textable, content string) { + text := render() + Expect(text).NotTo(BeNil()) + Expect(textStyleForContent(text, content)).To(ContainElement("text-muted")) + Expect(text.HTML()).NotTo(MatchRegexp(`text-gray-(600|700)`)) + Expect(text.String()).To(ContainSubstring(strings.TrimSpace(content))) + }, + Entry("assistant body", func() api.Textable { + return (&AssistantTool{BaseTool: BaseTool{Input: map[string]any{"text": "assistant payload"}}}).Pretty() + }, " assistant payload"), + Entry("reasoning body", func() api.Textable { + return (&ReasoningTool{BaseTool: BaseTool{Input: map[string]any{"text": "reasoning payload"}}}).Pretty() + }, " reasoning payload"), + Entry("session model", func() api.Textable { + return (&SystemInitTool{BaseTool: BaseTool{Input: map[string]any{"model": "model-x"}}}).Pretty() + }, " model-x"), + Entry("hook stdout label", func() api.Textable { + return (&HookResponseTool{BaseTool: BaseTool{Input: map[string]any{"stdout": "hook output"}}}).Detail() + }, "stdout: "), + Entry("parse error raw label", func() api.Textable { + return (&ParseErrorTool{BaseTool: BaseTool{Input: map[string]any{"raw": "raw payload"}}}).Detail() + }, "raw: "), + Entry("turn label", func() api.Textable { + return (&TurnDurationTool{}).Pretty() + }, " turn"), + Entry("away summary label", func() api.Textable { + return (&AwaySummaryTool{}).Pretty() + }, " away-summary"), + Entry("session title", func() api.Textable { + return (&SessionTitleTool{BaseTool: BaseTool{Input: map[string]any{"aiTitle": "session title"}}}).Pretty() + }, " session title"), + Entry("token total", func() api.Textable { + return (&TokenCountTool{BaseTool: BaseTool{Input: map[string]any{"total_tokens": 12}}}).Pretty() + }, " 12"), + Entry("Codex command", func() api.Textable { + return (&CodexExecCommandTool{BaseTool: BaseTool{Input: map[string]any{"command": "command --flag"}}}).Pretty() + }, " command --flag"), + Entry("user shell command", func() api.Textable { + return (&UserShellCommandTool{BaseTool: BaseTool{Input: map[string]any{"command": "local --flag"}}}).Pretty() + }, " local --flag"), + Entry("MCP invocation", func() api.Textable { + return (&MCPToolCallTool{BaseTool: BaseTool{Input: map[string]any{"invocation": map[string]any{"server": "server", "tool": "tool"}}}}).Pretty() + }, " server.tool"), + Entry("web search query", func() api.Textable { + return (&WebSearchEventTool{BaseTool: BaseTool{Input: map[string]any{"query": "search terms"}}}).Pretty() + }, " search terms"), + Entry("image path", func() api.Textable { + return (&ViewImageTool{BaseTool: BaseTool{Input: map[string]any{"path": "image.png"}}}).Pretty() + }, " image.png"), + Entry("guardian action", func() api.Textable { + return (&GuardianAssessmentTool{BaseTool: BaseTool{Input: map[string]any{"action": map[string]any{"tool": "tool-name"}}}}).Pretty() + }, " tool-name"), + Entry("collaboration nickname", func() api.Textable { + return (&CollabEventTool{BaseTool: BaseTool{Input: map[string]any{"nickname": "worker"}}}).Pretty() + }, " worker"), + Entry("budget used", func() api.Textable { + return (&BudgetTool{BaseTool: BaseTool{Input: map[string]any{"used": 1.25}}}).Pretty() + }, " used=$1.25"), + Entry("pull request number", func() api.Textable { + return (&PrLinkTool{BaseTool: BaseTool{Input: map[string]any{"prNumber": 42}}}).Pretty() + }, " #42"), + Entry("worktree name", func() api.Textable { + return (&WorktreeStateTool{BaseTool: BaseTool{Input: map[string]any{"worktreeSession": map[string]any{"worktreeName": "feature-name"}}}}).Pretty() + }, " feature-name"), + Entry("relocated directory", func() api.Textable { + return (&RelocatedTool{BaseTool: BaseTool{Input: map[string]any{"relocatedCwd": "repo"}}}).Pretty() + }, " repo"), + Entry("command stdout label", func() api.Textable { + return (&CodexExecCommandTool{BaseTool: BaseTool{Input: map[string]any{"stdout": "command output"}}}).Detail() + }, "stdout: "), + Entry("claude command arguments", func() api.Textable { + return (&ClaudeCommandTool{BaseTool: BaseTool{Input: map[string]any{"event": "claude_command", "command_args": "argument preview"}}}).Pretty() + }, " argument preview"), + Entry("claude command output", func() api.Textable { + return (&ClaudeCommandTool{BaseTool: BaseTool{Input: map[string]any{"event": "claude_command_output", "stream": "stdout", "content": "output preview"}}}).Pretty() + }, " output preview"), + Entry("goal condition", func() api.Textable { + return (&GoalStatusTool{BaseTool: BaseTool{Input: map[string]any{"condition": "goal condition"}}}).Pretty() + }, " goal condition"), + ) +}) + +func textStyleForContent(value api.Textable, content string) []string { + switch text := value.(type) { + case api.Text: + return findTextStyle(text, content) + case *api.Text: + return findTextStyle(*text, content) + default: + return nil + } +} + +func findTextStyle(text api.Text, content string) []string { + if text.Content == content { + return strings.Fields(text.Style) + } + for _, child := range text.Children { + if style := textStyleForContent(child, content); style != nil { + return style + } + } + return nil +} diff --git a/pkg/claude/tools/stream.go b/pkg/claude/tools/stream.go index ab9e5c3..26028c9 100644 --- a/pkg/claude/tools/stream.go +++ b/pkg/claude/tools/stream.go @@ -22,7 +22,7 @@ func (t *SystemInitTool) Pretty() api.Text { Add(icons.Icon{Unicode: "🚀", Iconify: "mdi:rocket-launch", Style: "muted"}). Append(" init", "text-purple-500 font-medium") if model := t.Str("model"); model != "" { - text = text.Append(" "+model, "text-gray-700") + text = text.Append(" "+model, "text-muted") } if cwd := t.Str("cwd"); cwd != "" { text = text.Append(" cwd="+ShortenPath(cwd), "text-gray-500") @@ -86,7 +86,7 @@ func (t *HookResponseTool) Detail() api.Textable { } text := clicky.Text("") if stdout != "" { - text = text.Append("stdout: ", "font-bold text-gray-600").Append(stdout, "") + text = text.Append("stdout: ", "font-bold text-muted").Append(stdout, "") } if stderr != "" { if stdout != "" { @@ -156,7 +156,7 @@ func (t *ApiErrorTool) Pretty() api.Text { text := clicky.Text(""). Add(icons.Icon{Unicode: "❌", Iconify: "mdi:alert-circle", Style: "muted"}). Append(" api-error", "text-red-500 font-bold") - if errStr := t.Str("error"); errStr != "" { + if errStr := errorSummary(t.Input["error"]); errStr != "" { text = text.Append(" "+errStr, "text-red-500") } if status := int(t.Float("api_error_status")); status > 0 { @@ -168,6 +168,23 @@ func (t *ApiErrorTool) Pretty() api.Text { return text } +func errorSummary(value any) string { + switch v := value.(type) { + case string: + return v + case map[string]any: + for _, key := range []string{"formatted", "message", "status"} { + if s, ok := v[key].(string); ok && s != "" { + return s + } + if n, ok := v[key].(float64); ok && n != 0 { + return fmt.Sprintf("%.0f", n) + } + } + } + return "" +} + func (t *ApiErrorTool) Detail() api.Textable { if d := t.BaseTool.Detail(); d != nil { return d @@ -198,7 +215,7 @@ func (t *ParseErrorTool) Detail() api.Textable { if raw == "" { return t.BaseTool.Detail() } - text := clicky.Text("").Append("raw: ", "font-bold text-gray-600").Append(raw, "") + text := clicky.Text("").Append("raw: ", "font-bold text-muted").Append(raw, "") return &text } @@ -244,7 +261,7 @@ func (t *TurnDurationTool) Category() string { return "system" } func (t *TurnDurationTool) Pretty() api.Text { text := clicky.Text(""). Add(icons.Icon{Unicode: "⏱", Iconify: "mdi:timer-outline", Style: "muted"}). - Append(" turn", "text-gray-700 font-medium") + Append(" turn", "text-muted font-medium") if ms := t.Float("durationMs"); ms > 0 { text = text.Append(" "+formatDurationMS(ms), "text-gray-500") } @@ -264,7 +281,7 @@ func (t *AwaySummaryTool) Category() string { return "system" } func (t *AwaySummaryTool) Pretty() api.Text { text := clicky.Text(""). Add(icons.Icon{Unicode: "💤", Iconify: "mdi:sleep", Style: "muted"}). - Append(" away-summary", "text-gray-700 font-medium") + Append(" away-summary", "text-muted font-medium") if content := strings.TrimSpace(t.Str("content")); content != "" { preview := content if len(preview) > 80 { @@ -296,7 +313,7 @@ func (t *SessionTitleTool) Pretty() api.Text { Add(icons.Icon{Unicode: "🏷", Iconify: "mdi:tag", Style: "muted"}). Append(" title", "text-purple-500 font-medium") if title := t.Str("aiTitle"); title != "" { - text = text.Append(" "+title, "text-gray-700") + text = text.Append(" "+title, "text-muted") } return text } diff --git a/pkg/claude/tools/stream_test.go b/pkg/claude/tools/stream_test.go index 57f402a..145372f 100644 --- a/pkg/claude/tools/stream_test.go +++ b/pkg/claude/tools/stream_test.go @@ -1,6 +1,7 @@ package tools import ( + "strings" "testing" "github.com/stretchr/testify/assert" @@ -76,8 +77,78 @@ func TestResultSummaryTool_NameCategoryPretty(t *testing.T) { } func TestNewTool_DispatchSyntheticTypes(t *testing.T) { - for _, name := range []string{"SessionInit", "HookStart", "HookResponse", "Result"} { + for _, name := range []string{ + "SessionInit", "HookStart", "HookResponse", "Result", + "TokenCount", "TaskStarted", "TaskComplete", "TurnAborted", + "ContextCompacted", "ThreadRolledBack", "ItemCompleted", + "CodexExecCommand", "UserShellCommand", "CodexPatchApply", "MCPToolCall", + "WebSearchEvent", "ViewImage", "GuardianAssessment", "ReviewMode", + "CollabAgentSpawn", "CollabAgentInteraction", "CollabWaiting", + "CollabClose", "QueueOperation", "DeferredToolsDelta", + "AgentListingDelta", "SkillListing", "Budget", "PrLink", + "CompactBoundary", "LocalCommand", "ScheduledTaskFire", + "Informational", "WorktreeState", "Relocated", "Started", + "ClaudeCommand", "GoalStatus", + } { got := NewTool(BaseTool{RawTool: name}) assert.Equal(t, name, got.Name(), "expected NewTool to return the right concrete type for %q", name) } } + +// TestSkillListingTool_PrettyEmitsSingleCount guards against the row rendering +// "count=N count=N" when a listing carries both the names array and the +// redundant skillCount scalar, as Claude Code transcripts do. +func TestSkillListingTool_PrettyEmitsSingleCount(t *testing.T) { + for _, tc := range []struct { + name string + input map[string]any + want string + }{ + { + name: "names and skillCount both present", + input: map[string]any{"names": []any{"a", "b"}, "skillCount": float64(29)}, + want: " count=2", + }, + { + name: "only skillCount present", + input: map[string]any{"skillCount": float64(29)}, + want: " count=29", + }, + { + name: "only names present", + input: map[string]any{"names": []any{"a", "b", "c"}}, + want: " count=3", + }, + } { + t.Run(tc.name, func(t *testing.T) { + pretty := NewTool(BaseTool{RawTool: "SkillListing", Input: tc.input}).Pretty().String() + assert.Equal(t, 1, strings.Count(pretty, "count="), + "expected exactly one count= field in %q", pretty) + assert.Contains(t, pretty, tc.want) + }) + } +} + +func TestUserShellCommandTool_PrettyAndDetail(t *testing.T) { + tool := NewTool(BaseTool{ + RawTool: "UserShellCommand", + Input: map[string]any{ + "command": "gavel proc restart", + "exit_code": 1, + "duration_ms": 2990.9, + "stdout": "Kill sent but port 8088 is still bound", + }, + }) + if _, ok := tool.(*UserShellCommandTool); !ok { + t.Fatalf("NewTool returned %T, want *UserShellCommandTool", tool) + } + pretty := tool.Pretty().String() + assert.Contains(t, pretty, "local command") + assert.Contains(t, pretty, "gavel proc restart") + assert.Contains(t, pretty, "exit=1") + assert.Contains(t, pretty, "3.0s") + detail := tool.Detail() + if assert.NotNil(t, detail) { + assert.Contains(t, detail.String(), "Kill sent but port 8088 is still bound") + } +} diff --git a/pkg/claude/tools/tool.go b/pkg/claude/tools/tool.go index c08e9f4..4dba65c 100644 --- a/pkg/claude/tools/tool.go +++ b/pkg/claude/tools/tool.go @@ -235,12 +235,68 @@ func NewTool(base BaseTool) Tool { return &AskTool{BaseTool: base} case "ExitPlanMode": return &ExitPlanTool{BaseTool: base} + case "Plan": + return &PlanTool{BaseTool: base} case "User": return &UserTool{BaseTool: base} case "Assistant": return &AssistantTool{BaseTool: base} case "Reasoning": return &ReasoningTool{BaseTool: base} + case "Event": + return &EventTool{BaseTool: base} + case "TokenCount": + return &TokenCountTool{BaseTool: base} + case "TaskStarted": + return &TaskStartedTool{BaseTool: base} + case "TaskComplete": + return &TaskCompleteTool{BaseTool: base} + case "TurnAborted", "ContextCompacted", "ThreadRolledBack", "ItemCompleted": + return &LifecycleEventTool{BaseTool: base} + case "MemoryCitation": + return &EventTool{BaseTool: base} + case "CodexExecCommand": + return &CodexExecCommandTool{BaseTool: base} + case "UserShellCommand": + return &UserShellCommandTool{BaseTool: base} + case "CodexPatchApply": + return &CodexPatchApplyTool{BaseTool: base} + case "MCPToolCall": + return &MCPToolCallTool{BaseTool: base} + case "WebSearchEvent": + return &WebSearchEventTool{BaseTool: base} + case "ViewImage": + return &ViewImageTool{BaseTool: base} + case "GuardianAssessment": + return &GuardianAssessmentTool{BaseTool: base} + case "ReviewMode": + return &ReviewModeTool{BaseTool: base} + case "CollabAgentSpawn", "CollabAgentInteraction", "CollabWaiting", "CollabClose": + return &CollabEventTool{BaseTool: base} + case "QueueOperation": + return &QueueOperationTool{BaseTool: base} + case "DeferredToolsDelta": + return &DeferredToolsDeltaTool{BaseTool: base} + case "AgentListingDelta": + return &AgentListingDeltaTool{BaseTool: base} + case "SkillListing": + return &SkillListingTool{BaseTool: base} + case "Budget": + return &BudgetTool{BaseTool: base} + case "PrLink": + return &PrLinkTool{BaseTool: base} + case "CompactBoundary", "LocalCommand", "ScheduledTaskFire", "Informational": + return &ContentEventTool{BaseTool: base} + case "ClaudeCommand": + return &ClaudeCommandTool{BaseTool: base} + case "GoalStatus": + return &GoalStatusTool{BaseTool: base} + case "WorktreeState": + return &WorktreeStateTool{BaseTool: base} + case "Relocated": + return &RelocatedTool{BaseTool: base} + case "Started": + return &StartedTool{BaseTool: base} case "Skill": return &SkillTool{BaseTool: base} case "SessionInit": diff --git a/pkg/claude/tools/user.go b/pkg/claude/tools/user.go index 3b5e72a..2aa4b10 100644 --- a/pkg/claude/tools/user.go +++ b/pkg/claude/tools/user.go @@ -8,7 +8,7 @@ import ( type UserTool struct{ BaseTool } func (t *UserTool) Name() string { return "User" } -func (t *UserTool) Category() string { return "" } +func (t *UserTool) Category() string { return "chat" } func (t *UserTool) FilePath() string { return "" } func (t *UserTool) ExtractPath() string { return "" } func (t *UserTool) Detail() api.Textable { return nil } diff --git a/pkg/claude/tooluse.go b/pkg/claude/tooluse.go index e1f0ad1..45e670d 100644 --- a/pkg/claude/tooluse.go +++ b/pkg/claude/tooluse.go @@ -8,6 +8,7 @@ import ( "strings" "time" + "github.com/flanksource/captain/pkg/ai/assistanttags" "github.com/flanksource/captain/pkg/bash" "github.com/flanksource/captain/pkg/claude/tools" captainCollections "github.com/flanksource/captain/pkg/collections" @@ -96,47 +97,90 @@ const denialCommentSeparator = "the user said:\n" const boilerplatePrefix = "\n\nNote: The user" const askAnswerPrefix = "User has answered your question." -// ExtractToolUses extracts ToolUse records from history entries +// ExtractToolUses extracts history activity rows from transcript entries: +// user/assistant text, reasoning, metadata events, and real tool calls. func ExtractToolUses(entries []HistoryEntry) []ToolUse { var toolUses []ToolUse - // Pass 1: extract tool_use blocks for _, entry := range entries { ts, _ := entry.ParseTimestamp() + var timestamp *time.Time + if !ts.IsZero() { + timestamp = &ts + } - for _, content := range entry.Message.Content { - if content.Type != ContentTypeToolUse { - continue - } - - var inputMap map[string]any - if content.Input != nil { - _ = json.Unmarshal(content.Input, &inputMap) - } + if isVisibleTranscriptEvent(entry.Event) { + toolUses = append(toolUses, ToolUse{ + Tool: tools.EventToolName(entry.Event.Type), + Input: eventInput(entry.Event), + Timestamp: timestamp, + CWD: entry.CWD, + SessionID: entry.SessionID, + ToolUseID: entry.UUID, + RawEntry: entry.RawLine, + }) + } - var timestamp *time.Time - if !ts.IsZero() { - timestamp = &ts - } + for _, content := range entry.Message.Content { + switch content.Type { + case ContentTypeText: + if entry.Message.Role == MessageRoleAssistant { + toolUses = append(toolUses, taggedAssistantToolUses(entry, content.Text, timestamp)...) + continue + } + tool := messageTool(entry.Message.Role) + if tool == "" || content.Text == "" { + continue + } + toolUses = append(toolUses, ToolUse{ + Tool: tool, + Input: map[string]any{"text": content.Text}, + Timestamp: timestamp, + CWD: entry.CWD, + SessionID: entry.SessionID, + IsSidechain: entry.IsSidechain, + AgentID: entry.AgentID, + RawEntry: entry.RawLine, + }) + case ContentTypeThinking, ContentTypeRedactedThinking: + if content.Thinking == "" { + continue + } + toolUses = append(toolUses, ToolUse{ + Tool: "Reasoning", + Input: map[string]any{"text": content.Thinking}, + Timestamp: timestamp, + CWD: entry.CWD, + SessionID: entry.SessionID, + IsSidechain: entry.IsSidechain, + AgentID: entry.AgentID, + RawEntry: entry.RawLine, + }) + case ContentTypeToolUse: + var inputMap map[string]any + if content.Input != nil { + _ = json.Unmarshal(content.Input, &inputMap) + } - var cwd string - if inputMap != nil { - if v, ok := inputMap["cwd"].(string); ok { - cwd = v + var cwd string + if inputMap != nil { + if v, ok := inputMap["cwd"].(string); ok { + cwd = v + } } - } - toolUses = append(toolUses, ToolUse{ - Tool: content.Name, - Input: inputMap, - Timestamp: timestamp, - CWD: cwd, - SessionID: entry.SessionID, - ToolUseID: content.ID, - IsSidechain: entry.IsSidechain, - AgentID: entry.AgentID, - RawEntry: entry.RawLine, - }) + toolUses = append(toolUses, ToolUse{ + Tool: content.Name, + Input: inputMap, + Timestamp: timestamp, + CWD: cwd, + SessionID: entry.SessionID, + ToolUseID: content.ID, + IsSidechain: entry.IsSidechain, + AgentID: entry.AgentID, + RawEntry: entry.RawLine, + }) + } } } @@ -156,6 +200,85 @@ func ExtractToolUses(entries []HistoryEntry) []ToolUse { return expandUserRows(toolUses) } +func taggedAssistantToolUses(entry HistoryEntry, text string, timestamp *time.Time) []ToolUse { + segments := assistanttags.Parse(text) + uses := make([]ToolUse, 0, len(segments)) + for _, segment := range segments { + use := ToolUse{ + Timestamp: timestamp, + CWD: entry.CWD, + SessionID: entry.SessionID, + IsSidechain: entry.IsSidechain, + AgentID: entry.AgentID, + RawEntry: entry.RawLine, + } + switch segment.Kind { + case assistanttags.SegmentPlan: + use.Tool = "Plan" + use.Input = map[string]any{"content": segment.Text, "tag": "proposed_plan"} + case assistanttags.SegmentMemoryCitation: + if segment.Citation == nil { + continue + } + use.Tool = "MemoryCitation" + use.Input = map[string]any{ + "event": "memory_citation", + "source": "claude", + "citation_entries": segment.Citation.CitationEntries, + "rollout_ids": segment.Citation.RolloutIDs, + } + case assistanttags.SegmentText: + use.Tool = "Assistant" + body := segment.Text + if summary, ok := assistanttags.EnvelopeSummary(body); ok { + body = summary + } + use.Input = map[string]any{"text": body} + default: + continue + } + uses = append(uses, use) + } + return uses +} + +func isVisibleTranscriptEvent(event *TranscriptEvent) bool { + if event == nil { + return false + } + switch event.Type { + case "file-history-snapshot", "last-prompt", "permission-mode", "agent-name": + return false + default: + return true + } +} + +func messageTool(role MessageRole) string { + switch role { + case MessageRoleUser: + return "User" + case MessageRoleAssistant: + return "Assistant" + default: + return "" + } +} + +func eventInput(event *TranscriptEvent) map[string]any { + input := map[string]any{"event": event.Type} + if event.Scope != "" { + input["scope"] = event.Scope + } + if event.Subtype != "" { + input["subtype"] = event.Subtype + } + for k, v := range event.Data { + input[k] = v + } + return input +} + // toolResult holds matched result data for a tool call. type toolResult struct { content json.RawMessage @@ -556,6 +679,47 @@ func suggestFilters(filterName string, filters []string, available map[string]st } } +// AbsolutePath anchors a tool-use path to an absolute, cleaned path. +// +// Tool inputs mix absolute paths (Read/Write/Edit carry them) with paths +// relative to the directory the tool ran in (a bash `cat pkg/x.go`), and an +// agent's working directory moves during a session — so a path is only +// unambiguous once anchored to the cwd that produced it. Anchoring prefers that +// cwd and falls back to the project root. +// +// This is the canonical form every consumer should hold: relativising per tool +// use bakes in whichever base that call happened to have, which silently differs +// across one session and leaves the path meaningless to anyone else. Render with +// RelativePath at the point of display instead. +// +// A path that cannot be anchored — no base, or a fragment the shell never +// expanded, e.g. "$DIR/x" — is returned cleaned but otherwise unchanged rather +// than guessed at. +func AbsolutePath(path, cwd, projectRoot string) string { + if path == "" { + return path + } + anchored := func() string { + if filepath.IsAbs(path) { + return filepath.Clean(path) + } + base := cwd + if base == "" { + base = projectRoot + } + if base == "" { + return filepath.Clean(path) + } + return filepath.Join(base, path) + }() + // Clean and Join drop a trailing separator, but that separator is how a + // directory argument (Grep/Glob `pkg/`) is told apart from a file, so keep it. + if strings.HasSuffix(path, "/") && !strings.HasSuffix(anchored, "/") { + anchored += "/" + } + return anchored +} + // RelativePath makes an absolute path relative to projectRoot if possible. // For paths outside the project (more than 1 parent level away), returns absolute path. func RelativePath(path, projectRoot string) string { diff --git a/pkg/claude/tooluse_test.go b/pkg/claude/tooluse_test.go index 19ca79b..f078294 100644 --- a/pkg/claude/tooluse_test.go +++ b/pkg/claude/tooluse_test.go @@ -127,6 +127,58 @@ func TestFilterToolUses(t *testing.T) { } } +func TestExtractToolUses_SplitsTaggedAssistantText(t *testing.T) { + entry := HistoryEntry{ + UUID: "assistant-1", + SessionID: "claude-tagged", + Timestamp: "2026-07-10T10:00:00Z", + Message: Message{ + Role: MessageRoleAssistant, + Content: []ContentBlock{{Type: ContentTypeText, Text: `# Claude fallback +After the plan. + + +MEMORY.md:1-2|note=[shared parser] + + +019f3754-ecfa-7323-a76b-a0205ea30bbe + +`}}, + }, + } + + uses := ExtractToolUses([]HistoryEntry{entry}) + if len(uses) != 3 { + t.Fatalf("uses = %+v, want Plan, Assistant, MemoryCitation", uses) + } + if uses[0].Tool != "Plan" || uses[0].Input["content"] != "# Claude fallback" { + t.Fatalf("plan = %+v", uses[0]) + } + if uses[1].Tool != "Assistant" || uses[1].Input["text"] != "After the plan." { + t.Fatalf("assistant = %+v", uses[1]) + } + if uses[2].Tool != "MemoryCitation" || uses[2].Input["event"] != "memory_citation" { + t.Fatalf("citation = %+v", uses[2]) + } +} + +func TestExtractToolUses_EnvelopeRendersSummary(t *testing.T) { + entry := HistoryEntry{ + Timestamp: "2026-07-08T11:20:00Z", + Message: Message{ + Role: MessageRoleAssistant, + Content: []ContentBlock{{Type: ContentTypeText, Text: `{"endStatus":"completed","plan":{"content":"","path":"/Users/moshe/.codex/plans/x.md","status":"new"},"questions":[],"summary":"Authored the review-banner plan."}`}}, + }, + } + uses := ExtractToolUses([]HistoryEntry{entry}) + if len(uses) != 1 || uses[0].Tool != "Assistant" { + t.Fatalf("uses = %+v, want single Assistant", uses) + } + if uses[0].Input["text"] != "Authored the review-banner plan." { + t.Fatalf("text = %q, want the plain summary", uses[0].Input["text"]) + } +} + func TestExtractToolUses_DetectsDenials(t *testing.T) { denialContent, _ := json.Marshal("The user doesn't want to proceed with this tool use. The tool use was rejected (eg. if it was a file edit, the new_string was NOT written to the file). To tell you how to proceed, the user said:\nkeep commons/logger") diff --git a/pkg/claude/types.go b/pkg/claude/types.go index f407ba2..8c6ba86 100644 --- a/pkg/claude/types.go +++ b/pkg/claude/types.go @@ -10,6 +10,8 @@ const ( HookEventStop HookEventType = "Stop" HookEventSubagentStop HookEventType = "SubagentStop" HookEventUserPromptSubmit HookEventType = "UserPromptSubmit" + HookEventSessionStart HookEventType = "SessionStart" + HookEventSessionEnd HookEventType = "SessionEnd" ) // PermissionMode represents Claude's current permission mode diff --git a/pkg/cli/ai.go b/pkg/cli/ai.go index 123fb4f..bf4d63c 100644 --- a/pkg/cli/ai.go +++ b/pkg/cli/ai.go @@ -3,7 +3,6 @@ package cli import ( "context" "fmt" - "io" "os" "strconv" "strings" @@ -12,6 +11,7 @@ import ( "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/ai/middleware" "github.com/flanksource/captain/pkg/ai/pricing" + "github.com/flanksource/captain/pkg/aiflags" "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/captainconfig" "github.com/flanksource/captain/pkg/claude" @@ -25,21 +25,35 @@ import ( // surfaced as zero-valued defaults rather than failing the command — a missing // or unreadable config should never block `captain ai prompt`. func loadSavedAI() captainconfig.AIDefaults { + return loadSavedConfig().AI +} + +func loadSavedConfig() captainconfig.Config { cfg, _, err := captainconfig.Load() if err != nil { log.Warnf("captainconfig load: %v (continuing with zero defaults)", err) - return captainconfig.AIDefaults{} + return captainconfig.Config{} } - return cfg.AI + return cfg } +// AIProviderOptions binds model selection plus the two knobs that belong to the +// request rather than the model: the API key and the spend budget. +// +// The model flags themselves live in pkg/aiflags — a leaf any clicky CLI can embed +// without inheriting pkg/cli's ~1000 transitive packages. Embedding it here keeps +// captain's flag surface unchanged (clicky promotes embedded flags at any depth) +// while giving downstream repos the same parsing captain uses. type AIProviderOptions struct { - Model string `flag:"model" help:"Model name(s), e.g. claude-sonnet-5 or a comma-separated primary,fallback list like claude-sonnet-5,gpt-4o (defaults to the value saved by 'captain configure')" short:"m"` - Fallback []string `flag:"fallback" help:"Model to try if the primary is unavailable (repeatable; comma-separated allowed)"` - Backend string `flag:"backend" help:"Force backend: anthropic|gemini|openai|deepseek|claude-cli|claude-agent|claude-cmux|codex-cli|codex-cmux|gemini-cli (default: inferred from model or saved by 'captain configure')" short:"b"` - APIKey string `flag:"api-key" help:"API key (env: ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY, GOOGLE_API_KEY, DEEPSEEK_API_KEY)"` - NoCache bool `flag:"no-cache" help:"Disable response caching"` - Budget string `flag:"budget" help:"Max spend in USD, 0=unlimited" default:"0"` + aiflags.ModelFlags + + APIKey string `flag:"api-key" help:"API key (env: ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY, GOOGLE_API_KEY, DEEPSEEK_API_KEY)"` + Budget string `flag:"budget" help:"Max spend in USD, 0=unlimited" default:"0"` +} + +// BudgetUSD parses --budget, failing loud on malformed input. +func (o AIProviderOptions) BudgetUSD() (float64, error) { + return parseFloatFlag("budget", o.Budget) } // parseFloatFlag parses a numeric string flag, returning a descriptive error @@ -56,36 +70,50 @@ func parseFloatFlag(name, val string) (float64, error) { } func (o AIProviderOptions) ToConfig() (ai.Config, error) { - saved := loadSavedAI() - budget, err := parseFloatFlag("budget", o.Budget) + savedCfg := loadSavedConfig() + saved := savedCfg.AI + budget, err := o.BudgetUSD() if err != nil { return ai.Config{}, err } - - model := o.Model - if model == "" { - model = saved.Model - } - backend := o.Backend - if backend == "" { - backend = saved.Backend - } - if backend != "" && !ai.Backend(backend).Valid() { - return ai.Config{}, fmt.Errorf("invalid --backend %q (valid: %s)", backend, ai.BackendList()) - } if budget == 0 { budget = saved.BudgetUSD } - m := api.Model{Name: model, Backend: ai.Backend(backend), Fallbacks: fallbackModelsFromFlags(o.Fallback)} + // One resolve: flags → Model → saved per-provider defaults → catalog. The + // warn-and-continue policy for a broken config stays here (loadSavedConfig), so + // aiflags can hand the error back instead of swallowing it. + m, err := o.ResolveWith(saved) + if err != nil { + return ai.Config{}, err + } return ai.Config{ - Model: m.ExpandCSV(), - Budget: api.Budget{Cost: budget}, - APIKey: o.APIKey, - NoCache: o.NoCache || saved.NoCache, + Model: m, + Budget: api.Budget{Cost: budget}, + APIKey: o.APIKey, + NoCache: o.NoCache || saved.NoCache, + SchemaRepair: schemaRepairConfig(savedCfg.Prompts.SchemaRepair), }, nil } +func schemaRepairConfig(saved captainconfig.SchemaRepairDefaults) api.SchemaRepairConfig { + return api.SchemaRepairConfig{ + Model: api.Model{Name: saved.Model, Backend: api.Backend(saved.Backend)}, + Prompt: strings.TrimSpace(saved.Prompt), + } +} + +func isZeroSchemaRepair(c api.SchemaRepairConfig) bool { + return strings.TrimSpace(c.Prompt) == "" && + c.Model.Name == "" && + c.Model.ID == "" && + c.Model.Backend == "" && + c.Model.Temperature == nil && + c.Model.Effort == "" && + !c.Model.NoCache && + len(c.Model.Fallbacks) == 0 +} + // AIRuntimeOptions binds the per-invocation knobs every AI command shares — // model selection (via embedded AIProviderOptions), generation parameters // (max tokens, temperature, timeout, reasoning), permission/sandbox toggles @@ -99,11 +127,12 @@ func (o AIProviderOptions) ToConfig() (ai.Config, error) { type AIRuntimeOptions struct { AIProviderOptions - MaxTokens int `flag:"max-tokens" help:"Maximum output tokens (0 = saved default or 4096)"` - Temperature string `flag:"temperature" help:"Sampling temperature (0.0-2.0)" default:"0"` - Effort string `flag:"effort" help:"Reasoning effort: low|medium|high|xhigh (codex/genkit; others ignore)"` - MaxTurns int `flag:"max-turns" help:"Max agent turns 0-100, 0 = provider default (claude-agent)"` - Resume string `flag:"resume" help:"Resume an existing session by id (claude-agent, codex)"` + // Effort and Temperature are NOT here: they describe the model and so live on + // the embedded aiflags.ModelFlags, promoted through AIProviderOptions. + // Redeclaring them would bind --effort twice and panic cobra at init. + MaxTokens int `flag:"max-tokens" help:"Maximum output tokens (0 = saved default or 4096)"` + MaxTurns int `flag:"max-turns" help:"Max agent turns 0-100, 0 = provider default (claude-agent)"` + Resume string `flag:"resume" help:"Resume an existing session by id (claude-agent, codex)"` Edit bool `flag:"edit" help:"Safe defaults: acceptEdits + Read/Edit/Write/Glob/Grep allowlist"` AllowedTools []string `flag:"allowed-tools" help:"Override --edit's built-in allowlist (claude only)"` @@ -144,21 +173,25 @@ type AIPromptOptions struct { System string `flag:"system" help:"System prompt" short:"s"` AppendSystem string `flag:"append-system" help:"Append text to the default system prompt"` Var []string `flag:"var" help:"Template variable key=value (repeatable)" short:"V"` + Attach []string `flag:"attach" help:"Attach a local path or URL (repeatable; RFC 4180 comma-separated values allowed)" short:"A"` + MultiModels []string `flag:"multi-models" help:"Run prompt once per runtime selector in parallel, e.g. cli:sonnet-5,cmux:opus (repeatable; comma-separated allowed)" short:"M"` Timeout string `flag:"timeout" help:"Request timeout" default:"120s"` NoStream bool `flag:"no-stream" help:"Disable streaming; print only the final text to stdout"` } type AIPromptResult struct { - Text string `json:"text" pretty:"label=Response"` - Model string `json:"model" pretty:"label=Model"` - Backend string `json:"backend" pretty:"label=Backend"` - Dir string `json:"dir,omitempty" pretty:"label=Dir"` - SessionID string `json:"sessionId,omitempty" pretty:"label=Session"` - Input ai.Request `json:"input" pretty:"-"` - InputTokens int `json:"inputTokens" pretty:"label=Input Tokens"` - Output int `json:"outputTokens" pretty:"label=Output Tokens"` - CostUSD float64 `json:"costUSD,omitempty" pretty:"label=Cost USD"` - Duration string `json:"duration" pretty:"label=Duration"` + Text string `json:"text" pretty:"label=Response"` + StructuredOutput map[string]any `json:"structuredOutput,omitempty" pretty:"-"` + Model string `json:"model" pretty:"label=Model"` + Backend string `json:"backend" pretty:"label=Backend"` + Dir string `json:"dir,omitempty" pretty:"label=Dir"` + SessionID string `json:"sessionId,omitempty" pretty:"label=Session"` + HistoryFile string `json:"historyFile,omitempty" pretty:"label=History"` + Input ai.Request `json:"input" pretty:"-"` + InputTokens int `json:"inputTokens" pretty:"label=Input Tokens"` + Output int `json:"outputTokens" pretty:"label=Output Tokens"` + CostUSD float64 `json:"costUSD,omitempty" pretty:"label=Cost USD"` + Duration string `json:"duration" pretty:"label=Duration"` } // ToRequest translates the runtime knobs into the typed ai.Request, overlaying @@ -199,10 +232,18 @@ func (o AIRuntimeOptions) ToRequest(systemPrompt, appendSystemPrompt, userPrompt maxTokens = 4096 } - effort := o.Effort - if effort == "" { - effort = saved.ReasoningEffort + // Resolve the USD budget onto the request (flag > saved) so backends that read + // req.Budget.Cost — claude-cli, claude-agent — enforce it without relying on a + // later config-side reconciliation that not every path performs (finding A4). + budget, err := parseFloatFlag("budget", o.Budget) + if err != nil { + return ai.Request{}, err + } + if budget == 0 { + budget = saved.BudgetUSD } + + effort := o.Effort if err := api.Effort(effort).Validate(); err != nil { return ai.Request{}, fmt.Errorf("invalid --effort %q: %w", effort, err) } @@ -228,7 +269,7 @@ func (o AIRuntimeOptions) ToRequest(systemPrompt, appendSystemPrompt, userPrompt return ai.Request{ Prompt: api.Prompt{System: systemPrompt, AppendSystem: appendSystemPrompt, User: userPrompt}, Model: api.Model{Temperature: temperaturePtr, Effort: api.Effort(effort), NoCache: o.NoCache || saved.NoCache}, - Budget: api.Budget{MaxTokens: maxTokens, MaxTurns: o.MaxTurns}, + Budget: api.Budget{Cost: budget, MaxTokens: maxTokens, MaxTurns: o.MaxTurns}, Memory: api.Memory{ Skills: o.SkillDirs, SkipHooks: o.NoHooks || saved.NoHooks, @@ -246,45 +287,20 @@ func (o AIRuntimeOptions) ToRequest(systemPrompt, appendSystemPrompt, userPrompt // ToRequest delegates to AIRuntimeOptions.ToRequest, lifting the prompt // fields the prompt-shaped command owns onto the typed request. func (o AIPromptOptions) ToRequest() (ai.Request, error) { - return o.AIRuntimeOptions.ToRequest(o.System, o.AppendSystem, o.Prompt) -} - -// RunAIPrompt is a deprecated alias for `captain prompt run`. It routes through -// the same shared render + execute core (renderPromptSource + executePromptRequest) -// so there is one implementation; the positional `.prompt` file, --prompt/-p, and -// stdin all still work. -func RunAIPrompt(opts AIPromptOptions) (any, error) { - log.Warnf("`captain ai prompt` is deprecated; use `captain prompt run` (it accepts a .prompt file, id, --prompt/-p, or stdin)") - - var stdin string - if claude.IsStdinPiped() { - b, err := io.ReadAll(os.Stdin) - if err != nil { - return nil, fmt.Errorf("read stdin: %w", err) - } - stdin = string(b) - } - - ctx := context.Background() - req, cfg, err := renderPromptSource(ctx, opts.File, opts, "", stdin) + req, err := o.AIRuntimeOptions.ToRequest(o.System, o.AppendSystem, o.Prompt) if err != nil { - return nil, err - } - if req.Prompt.User == "" { - return nil, fmt.Errorf("prompt text required (use --prompt/-p, a file arg, or pipe via stdin)") - } - if cfg.Model.Name == "" { - return nil, fmt.Errorf("no model: pass --model or run 'captain configure' to set a default") - } - if err := req.Validate(); err != nil { - return nil, err + return ai.Request{}, err } - return executePromptRequest(ctx, req, cfg, runtimeTimeout(req.Budget.Timeout), opts.NoStream) + req.Prompt.Attachments, err = attachmentRefsFromFlags(o.Attach) + return req, err } func executePromptRequest(parent context.Context, req ai.Request, cfg ai.Config, timeout time.Duration, noStream bool) (any, error) { ctx, cancel := runContext(parent, req, timeout) defer cancel() + if err := preparePromptAttachments(ctx, &req, cfg); err != nil { + return nil, err + } p, cleanup, err := buildProvider(ctx, &req, cfg) if err != nil { @@ -292,7 +308,7 @@ func executePromptRequest(parent context.Context, req ai.Request, cfg ai.Config, } defer cleanup() - if streamer, ok := p.(ai.StreamingProvider); ok && !noStream { + if streamer, ok := p.(ai.StreamingProvider); ok && !noStream && !req.Prompt.HasSchema() { return runStreaming(ctx, streamer, req) } return runBuffered(ctx, p, req) @@ -414,7 +430,7 @@ func buildProvider(ctx context.Context, req *ai.Request, cfg ai.Config) (ai.Prov cleanup() return nil, func() {}, err } - if p, err = middleware.Wrap(p, middleware.WithLogging(), middleware.WithSchemaValidation()); err != nil { + if p, err = middleware.Wrap(p, middleware.WithLogging(), middleware.WithSchemaValidation(cfg)); err != nil { cleanup() return nil, func() {}, err } @@ -430,16 +446,27 @@ func runBuffered(ctx context.Context, p ai.Provider, req ai.Request) (any, error model := firstNonEmpty(resp.Model, p.GetModel(), req.Name) backend := firstNonEmpty(string(resp.Backend), string(p.GetBackend()), string(req.Backend)) input := resolvedPromptInput(req, model, backend, req.SessionID) + dir := actualRunDir(input) + structuredOutput, err := structuredOutputMap(resp.StructuredData) + if err != nil { + return nil, err + } + text, err := structuredOutputText(resp.Text, structuredOutput) + if err != nil { + return nil, err + } return AIPromptResult{ - Text: resp.Text, - Model: model, - Backend: backend, - Dir: input.Cwd(), - SessionID: input.SessionID, - Input: input, - InputTokens: resp.Usage.InputTokens, - Output: resp.Usage.OutputTokens, - Duration: time.Since(start).Round(time.Millisecond).String(), + Text: text, + StructuredOutput: structuredOutput, + Model: model, + Backend: backend, + Dir: dir, + SessionID: input.SessionID, + HistoryFile: historyFileForRun(api.Backend(backend), input.SessionID, dir), + Input: input, + InputTokens: resp.Usage.InputTokens, + Output: resp.Usage.OutputTokens, + Duration: time.Since(start).Round(time.Millisecond).String(), }, nil } @@ -449,17 +476,20 @@ func runBuffered(ctx context.Context, p ai.Provider, req ai.Request) (any, error func runStreaming(ctx context.Context, sp ai.StreamingProvider, req ai.Request) (any, error) { start := time.Now() var ( - text string - usage ai.Usage - cost float64 - backend = string(sp.GetBackend()) - model = sp.GetModel() - session = req.SessionID + text string + usage ai.Usage + cost float64 + backend = string(sp.GetBackend()) + model = sp.GetModel() + session = req.SessionID + structuredOutput map[string]any + structuredErr error ) renderer := newLineRenderer(os.Stderr, 8) loop, err := ai.RunUntil(ctx, ai.LoopOptions{ Provider: sp, MaxIterations: 1, + MaxCostUSD: req.Budget.Cost, // enforce the USD budget BuildRequest: func(iter int, prev *ai.LoopIteration) (ai.Request, bool) { if iter > 0 { return ai.Request{}, false @@ -478,6 +508,10 @@ func runStreaming(ctx context.Context, sp ai.StreamingProvider, req ai.Request) text += ev.Text } if ev.Kind == ai.EventResult { + if len(ev.StructuredData) > 0 { + text = string(ev.StructuredData) + structuredOutput, structuredErr = structuredOutputMap(ev.StructuredData) + } if ev.Usage != nil { usage = *ev.Usage } @@ -488,24 +522,33 @@ func runStreaming(ctx context.Context, sp ai.StreamingProvider, req ai.Request) if err != nil { return nil, err } + if structuredErr != nil { + return nil, structuredErr + } if loop.StopReason == "error" { return nil, fmt.Errorf("streaming loop stopped: %s", loop.StopReason) } + if loop.StopReason == "max-cost" { + return nil, fmt.Errorf("%w: spent $%.4f of $%.4f budget", ai.ErrBudgetExceeded, loop.TotalCost, req.Budget.Cost) + } if session == "" && len(loop.Iterations) > 0 { session = loop.Iterations[0].SessionID } input := resolvedPromptInput(req, model, backend, session) + dir := actualRunDir(input) return AIPromptResult{ - Text: text, - Model: model, - Backend: backend, - Dir: input.Cwd(), - SessionID: input.SessionID, - Input: input, - InputTokens: usage.InputTokens, - Output: usage.OutputTokens, - CostUSD: cost, - Duration: time.Since(start).Round(time.Millisecond).String(), + Text: text, + StructuredOutput: structuredOutput, + Model: model, + Backend: backend, + Dir: dir, + SessionID: input.SessionID, + HistoryFile: historyFileForRun(api.Backend(backend), input.SessionID, dir), + Input: input, + InputTokens: usage.InputTokens, + Output: usage.OutputTokens, + CostUSD: cost, + Duration: time.Since(start).Round(time.Millisecond).String(), }, nil } @@ -523,6 +566,17 @@ func resolvedPromptInput(req ai.Request, model, backend, sessionID string) ai.Re return out } +func actualRunDir(req ai.Request) string { + if cwd := req.Cwd(); cwd != "" { + return cwd + } + wd, err := os.Getwd() + if err != nil { + return "" + } + return wd +} + // renderEvent writes a human-readable representation of an ai.Event to w. // When the event carries a claude.HistoryEntry in Raw, route through the // shared lineRenderer so live `captain ai prompt` output matches diff --git a/pkg/cli/ai_catalog.go b/pkg/cli/ai_catalog.go index 0fdcf93..467d058 100644 --- a/pkg/cli/ai_catalog.go +++ b/pkg/cli/ai_catalog.go @@ -8,11 +8,8 @@ import "github.com/flanksource/captain/pkg/ai" // (subscription/OAuth via the installed binary), so enumerating their models // must never require the parent provider's API key. // -// The returned ID is the slug the captain provider expects at run time: -// AgentModel when the catalog sets one (e.g. codex's "gpt-5-codex", which the -// app-server sends verbatim) otherwise the catalog ID (e.g. "claude-agent-sonnet", -// which the claude-agent provider de-prefixes itself). claude-cli shares the -// claude-agent provider, so it shares its catalog entries. +// The returned ID is the exact provider model ID Captain sends at runtime. +// CLI/cmux modes share the corresponding agent catalog entries. func agentCatalogModels(b ai.Backend) []ai.ModelDef { return ai.AgentCatalogModels(b) } diff --git a/pkg/cli/ai_catalog_test.go b/pkg/cli/ai_catalog_test.go index 89fa766..1394dfd 100644 --- a/pkg/cli/ai_catalog_test.go +++ b/pkg/cli/ai_catalog_test.go @@ -17,9 +17,9 @@ func installTestCatalog(t *testing.T) { if err := ai.SetModelCatalog([]ai.Model{ {ID: "anthropic/claude-sonnet-4-5", Backend: ai.BackendAnthropic, Label: "Claude Sonnet 4.5"}, - {ID: "claude-agent-opus", Backend: ai.BackendClaudeAgent, Label: "Claude Agent · Opus"}, - {ID: "claude-agent-sonnet", Backend: ai.BackendClaudeAgent, Label: "Claude Agent · Sonnet"}, - {ID: "codex-gpt-5-codex", Backend: ai.BackendCodexCLI, AgentModel: "gpt-5-codex", Label: "Codex · GPT-5"}, + {ID: "claude-opus-4-8", Backend: ai.BackendClaudeAgent, Label: "Claude Agent · Opus 4.8"}, + {ID: "claude-sonnet-5", Backend: ai.BackendClaudeAgent, Label: "Claude Agent · Sonnet 5"}, + {ID: "gpt-5.5", Backend: ai.BackendCodexAgent, Label: "Codex Agent · GPT-5.5"}, }); err != nil { t.Fatalf("SetModelCatalog: %v", err) } @@ -33,18 +33,19 @@ func TestAgentCatalogModels(t *testing.T) { want []ai.ModelDef }{ { - // AgentModel slug wins over the catalog ID: codex receives the - // model string verbatim, so it must be "gpt-5-codex" not - // "codex-gpt-5-codex". backend: ai.BackendCodexCLI, - want: []ai.ModelDef{{ID: "gpt-5-codex", Name: "Codex · GPT-5", Backend: ai.BackendCodexCLI}}, + want: []ai.ModelDef{{ID: "gpt-5.5", Name: "Codex Agent · GPT-5.5", Backend: ai.BackendCodexCLI, CapabilitiesKnown: true}}, + }, + { + backend: ai.BackendCodexAgent, + want: []ai.ModelDef{{ID: "gpt-5.5", Name: "Codex Agent · GPT-5.5", Backend: ai.BackendCodexAgent, CapabilitiesKnown: true}}, }, { // Sorted by ID; genkit (API) anthropic entry excluded. backend: ai.BackendClaudeAgent, want: []ai.ModelDef{ - {ID: "claude-agent-opus", Name: "Claude Agent · Opus", Backend: ai.BackendClaudeAgent}, - {ID: "claude-agent-sonnet", Name: "Claude Agent · Sonnet", Backend: ai.BackendClaudeAgent}, + {ID: "claude-opus-4-8", Name: "Claude Agent · Opus 4.8", Backend: ai.BackendClaudeAgent, CapabilitiesKnown: true}, + {ID: "claude-sonnet-5", Name: "Claude Agent · Sonnet 5", Backend: ai.BackendClaudeAgent, CapabilitiesKnown: true}, }, }, { @@ -52,10 +53,24 @@ func TestAgentCatalogModels(t *testing.T) { // its own backend. backend: ai.BackendClaudeCLI, want: []ai.ModelDef{ - {ID: "claude-agent-opus", Name: "Claude Agent · Opus", Backend: ai.BackendClaudeCLI}, - {ID: "claude-agent-sonnet", Name: "Claude Agent · Sonnet", Backend: ai.BackendClaudeCLI}, + {ID: "claude-opus-4-8", Name: "Claude Agent · Opus 4.8", Backend: ai.BackendClaudeCLI, CapabilitiesKnown: true}, + {ID: "claude-sonnet-5", Name: "Claude Agent · Sonnet 5", Backend: ai.BackendClaudeCLI, CapabilitiesKnown: true}, }, }, + { + // claude-cmux uses the same local Claude catalog and is re-tagged with + // its own backend. + backend: ai.BackendClaudeCmux, + want: []ai.ModelDef{ + {ID: "claude-opus-4-8", Name: "Claude Agent · Opus 4.8", Backend: ai.BackendClaudeCmux, CapabilitiesKnown: true}, + {ID: "claude-sonnet-5", Name: "Claude Agent · Sonnet 5", Backend: ai.BackendClaudeCmux, CapabilitiesKnown: true}, + }, + }, + { + // codex-cmux shares the codex-agent runtime model slug. + backend: ai.BackendCodexCmux, + want: []ai.ModelDef{{ID: "gpt-5.5", Name: "Codex Agent · GPT-5.5", Backend: ai.BackendCodexCmux, CapabilitiesKnown: true}}, + }, { // No catalog entries for gemini-cli: empty, never an error. backend: ai.BackendGeminiCLI, diff --git a/pkg/cli/ai_filters.go b/pkg/cli/ai_filters.go index 71303c8..3cce83e 100644 --- a/pkg/cli/ai_filters.go +++ b/pkg/cli/ai_filters.go @@ -48,6 +48,8 @@ func effortFilterOptions() map[string]api.Textable { capapi.EffortMedium: "Medium", capapi.EffortHigh: "High", capapi.EffortXHigh: "Extra High", + capapi.EffortMax: "Maximum", + capapi.EffortUltra: "Ultra", } out := make(map[string]api.Textable, len(capapi.AllEfforts())) for _, e := range capapi.AllEfforts() { diff --git a/pkg/cli/ai_models.go b/pkg/cli/ai_models.go index 57e479d..12f3f97 100644 --- a/pkg/cli/ai_models.go +++ b/pkg/cli/ai_models.go @@ -13,7 +13,7 @@ import ( type AIModelsOptions struct { Filter string `flag:"filter" help:"Filter models by name substring" short:"f"` - Backend string `flag:"backend" help:"Filter by backend: anthropic|gemini|openai|deepseek|claude-cli|claude-agent|claude-cmux|codex-cli|codex-cmux|gemini-cli" short:"b"` + Backend string `flag:"backend" help:"Filter by backend: anthropic|gemini|openai|deepseek|claude-cli|claude-agent|claude-cmux|codex-cli|codex-agent|codex-cmux|gemini-cli" short:"b"` Limit int `flag:"limit" help:"Maximum models to show" default:"50" short:"l"` All bool `flag:"all" help:"Include all OpenRouter models" short:"a"` } @@ -96,7 +96,7 @@ func runLiveModels(opts AIModelsOptions) (any, error) { // Hide legacy/non-chat IDs unless the user asked for them by // name via --filter. Filtering by user intent overrides the // blacklist so "ai models -f gpt-3.5" still works. - if opts.Filter == "" && ai.IsLegacyModelID(m.ID) { + if opts.Filter == "" && ai.IsLegacyModelIDForBackend(m.ID, r.backend) { continue } diff --git a/pkg/cli/ai_models_test.go b/pkg/cli/ai_models_test.go index ac59c23..966ef37 100644 --- a/pkg/cli/ai_models_test.go +++ b/pkg/cli/ai_models_test.go @@ -132,13 +132,29 @@ func TestIsLegacyModelID(t *testing.T) { "gpt-5-nano": true, "gpt-5-codex": true, "gpt-5-pro": true, + "gpt-5.5-pro": true, + "gpt-5.5-nano": true, + "gpt-5.5-mini": true, + "gpt-5.4-mini": true, + "gpt-5.5-chat-latest": true, + "gpt-5.5-2026-04-23": true, + "gpt-realtime-2.1": true, + "gpt-realtime-2.1-mini": true, + "gpt-image-2": true, + "gpt-audio-1.5": true, + "gpt-5.3-codex": true, + "gpt-5.3-code": true, "o1": true, "o1-pro": true, "o3": true, "o3-pro": true, "o3-mini": true, + "o4-mini": true, + "sora-2": true, "codex-mini-latest": true, "dall-e-3": true, + "image-alpha-001": true, + "audio-alpha-001": true, "dall-e-2": true, "whisper-1": true, "tts-1": true, @@ -174,7 +190,8 @@ func TestIsLegacyModelID(t *testing.T) { "gpt-5.1": false, "gpt-5.2": false, "gpt-5.3": false, - "o4-mini": false, + "gpt-5.4": false, + "gpt-5.5": false, "claude-sonnet-4-5": false, "claude-sonnet-4-6": false, "claude-opus-4-5": false, @@ -203,8 +220,16 @@ func TestRunAIModels_HidesLegacyByDefault(t *testing.T) { "data": []map[string]any{ {"id": "gpt-5"}, {"id": "gpt-5.1"}, + {"id": "gpt-5.6-sol"}, + {"id": "gpt-5.6-terra"}, + {"id": "gpt-5.6-luna"}, {"id": "gpt-5-mini"}, {"id": "gpt-5-codex"}, + {"id": "gpt-5.5-pro"}, + {"id": "gpt-realtime-2.1"}, + {"id": "gpt-image-2"}, + {"id": "gpt-audio-1.5"}, + {"id": "sora-2"}, {"id": "gpt-3.5-turbo"}, {"id": "gpt-4"}, {"id": "gpt-4o-mini"}, @@ -227,7 +252,7 @@ func TestRunAIModels_HidesLegacyByDefault(t *testing.T) { } res := got.(AIModelsResult) - want := []string{"gpt-5", "gpt-5.1", "o4-mini"} + want := []string{"gpt-5", "gpt-5.1", "gpt-5.6-luna", "gpt-5.6-sol", "gpt-5.6-terra"} if len(res.Rows) != len(want) { t.Fatalf("rows = %d, want %d (%v)", len(res.Rows), len(want), res.Rows) } diff --git a/pkg/cli/ai_prompt_file.go b/pkg/cli/ai_prompt_file.go index c67e4d9..45b7283 100644 --- a/pkg/cli/ai_prompt_file.go +++ b/pkg/cli/ai_prompt_file.go @@ -24,8 +24,10 @@ func resolvePromptTemplate(opts AIPromptOptions, stdin string) (tmpl *prompt.Tem return prompt.Load(opts.Prompt), false, nil case strings.TrimSpace(stdin) != "": return prompt.Load(stdin), true, nil + case len(opts.Attach) > 0: + return prompt.Load(""), false, nil default: - return nil, false, fmt.Errorf("prompt required: pass a .prompt file, --prompt/-p text, or pipe via stdin") + return nil, false, fmt.Errorf("prompt or attachment required: pass a .prompt file, --prompt/-p text, --attach/-A, or pipe via stdin") } } @@ -66,6 +68,9 @@ func overlayCLI(base ai.Request, baseCfg ai.Config, o AIPromptOptions) (ai.Reque if bm.Name == "" { bm.Name = baseCfg.Model.Name } + if bm.ID == "" { + bm.ID = baseCfg.Model.ID + } if bm.Backend == "" { bm.Backend = baseCfg.Model.Backend } @@ -75,17 +80,27 @@ func overlayCLI(base ai.Request, baseCfg ai.Config, o AIPromptOptions) (ai.Reque if bm.Effort == "" { bm.Effort = baseCfg.Model.Effort } + identity := selectModelIdentity( + api.Model{Name: bm.Name, ID: bm.ID, Backend: bm.Backend}, + api.Model{Name: o.Model, Backend: api.Backend(o.Backend)}, + ) m := bm - m.Name = firstNonEmpty(o.Model, bm.Name, saved.Model) - m.Backend = api.Backend(firstNonEmpty(o.Backend, string(bm.Backend), saved.Backend)) + m.Name, m.ID, m.Backend = identity.Name, identity.ID, identity.Backend if temperature != 0 { t := temperature m.Temperature = &t } - m.Effort = api.Effort(firstNonEmpty(o.Effort, string(bm.Effort), saved.ReasoningEffort)) + m.Effort = api.Effort(firstNonEmpty(o.Effort, string(bm.Effort))) m.NoCache = o.NoCache || bm.NoCache || saved.NoCache m.Fallbacks = firstFallbacks(o.Fallback, bm.Fallbacks) - req.Model = m.ExpandCSV() + m, err = applyProviderDefaults(m, saved) + if err != nil { + return base, baseCfg, err + } + req.Model, err = ai.ResolveModelSelectors(m) + if err != nil { + return base, baseCfg, err + } req.Budget.MaxTokens = firstPositive(o.MaxTokens, base.Budget.MaxTokens, baseCfg.Budget.MaxTokens, saved.MaxTokens, 4096) req.Budget.Cost = firstPositiveFloat(budget, base.Budget.Cost, baseCfg.Budget.Cost, saved.BudgetUSD) @@ -98,6 +113,13 @@ func overlayCLI(base ai.Request, baseCfg ai.Config, o AIPromptOptions) (ai.Reque if o.AppendSystem != "" { req.Prompt.AppendSystem = o.AppendSystem } + if len(o.Attach) > 0 { + attachments, err := attachmentRefsFromFlags(o.Attach) + if err != nil { + return base, baseCfg, err + } + req.Prompt.Attachments = append(req.Prompt.Attachments, attachments...) + } if o.PermissionMode != "" { req.Permissions.Mode = api.PermissionMode(o.PermissionMode) @@ -196,6 +218,28 @@ func firstNonEmpty(vals ...string) string { return "" } +// selectModelIdentity applies precedence from lowest to highest while keeping a +// model name and backend coupled. A higher-priority name clears a lower-priority +// backend unless that same layer explicitly supplies one. +func selectModelIdentity(layers ...api.Model) api.Model { + var selected api.Model + for _, layer := range layers { + if layer.Name != "" { + selected.Name = layer.Name + selected.ID = layer.ID + selected.Backend = layer.Backend + continue + } + if layer.ID != "" { + selected.ID = layer.ID + } + if layer.Backend != "" { + selected.Backend = layer.Backend + } + } + return selected +} + // firstPositive returns the first value > 0, or 0 when none qualify. func firstPositive(vals ...int) int { for _, v := range vals { diff --git a/pkg/cli/ai_prompt_file_test.go b/pkg/cli/ai_prompt_file_test.go index d88b76c..0c28493 100644 --- a/pkg/cli/ai_prompt_file_test.go +++ b/pkg/cli/ai_prompt_file_test.go @@ -241,13 +241,13 @@ func TestOverlayCLI_FallbackFlagOverridesFrontmatter(t *testing.T) { base := baseFileReq() base.Model.Fallbacks = []api.Model{{Name: "frontmatter-fallback"}} opts := AIPromptOptions{} - opts.Fallback = []string{"cli-fallback-a", "cli-fallback-b,cli-fallback-c"} // repeatable + comma-split + opts.Fallback = []string{"gemini-3.5-flash", "gpt-5.5,claude-sonnet-5"} // repeatable + comma-split req, _, err := overlayCLI(base, ai.Config{}, opts) if err != nil { t.Fatal(err) } - want := []string{"cli-fallback-a", "cli-fallback-b", "cli-fallback-c"} + want := []string{"gemini-3.5-flash", "gpt-5.5", "claude-sonnet-5"} if got := fallbackNames(req.Model.Fallbacks); !reflect.DeepEqual(got, want) { t.Errorf("fallbacks = %v, want CLI flags %v (override frontmatter)", got, want) } diff --git a/pkg/cli/ai_test.go b/pkg/cli/ai_test.go index 9164d23..79f6e52 100644 --- a/pkg/cli/ai_test.go +++ b/pkg/cli/ai_test.go @@ -3,6 +3,7 @@ package cli import ( "context" "encoding/json" + "github.com/flanksource/captain/pkg/aiflags" "os" "path/filepath" "reflect" @@ -30,6 +31,21 @@ func (promptResultProvider) Execute(context.Context, ai.Request) (*ai.Response, }, nil } +type structuredPromptResultProvider struct{} + +func (structuredPromptResultProvider) GetModel() string { return "claude-sonnet-4-6" } + +func (structuredPromptResultProvider) GetBackend() ai.Backend { return ai.BackendAnthropic } + +func (structuredPromptResultProvider) Execute(context.Context, ai.Request) (*ai.Response, error) { + return &ai.Response{ + Text: `{"answer":"42"}`, + StructuredData: json.RawMessage(`{"answer":"42"}`), + Model: "claude-sonnet-4-6", + Backend: ai.BackendAnthropic, + }, nil +} + type promptResultStreamingProvider struct{} func (promptResultStreamingProvider) GetModel() string { return "gpt-5-codex" } @@ -49,6 +65,28 @@ func (promptResultStreamingProvider) ExecuteStream(_ context.Context, _ ai.Reque return events, nil } +type structuredResultStreamingProvider struct { + text string +} + +func (structuredResultStreamingProvider) GetModel() string { return "gpt-5-codex" } + +func (structuredResultStreamingProvider) GetBackend() ai.Backend { return ai.BackendCodexCLI } + +func (structuredResultStreamingProvider) Execute(context.Context, ai.Request) (*ai.Response, error) { + return nil, nil +} + +func (p structuredResultStreamingProvider) ExecuteStream(_ context.Context, _ ai.Request) (<-chan ai.Event, error) { + events := make(chan ai.Event, 2) + if p.text != "" { + events <- ai.Event{Kind: ai.EventText, Text: p.text} + } + events <- ai.Event{Kind: ai.EventResult, StructuredData: json.RawMessage(`{"answer":"42"}`)} + close(events) + return events, nil +} + // isolateSavedAI redirects captainconfig.Path() to an empty file inside // t.TempDir() so loadSavedAI() returns zero defaults rather than leaking // the developer's real ~/.captain.yaml into table-test expectations. @@ -199,7 +237,7 @@ func TestAIRuntimeOptions_ToRequest_ValidationErrors(t *testing.T) { {"temperature above max", func(o *AIPromptOptions) { o.Temperature = "3" }, "0.0-2.0"}, {"temperature below min", func(o *AIPromptOptions) { o.Temperature = "-1" }, "0.0-2.0"}, {"bad permission-mode", func(o *AIPromptOptions) { o.PermissionMode = "yolo" }, "permission-mode"}, - {"bad effort", func(o *AIPromptOptions) { o.Effort = "max" }, "effort"}, + {"bad effort", func(o *AIPromptOptions) { o.Effort = "extreme" }, "effort"}, {"max-turns above max", func(o *AIPromptOptions) { o.MaxTurns = 200 }, "max-turns"}, {"max-turns below min", func(o *AIPromptOptions) { o.MaxTurns = -1 }, "max-turns"}, } @@ -216,14 +254,37 @@ func TestAIRuntimeOptions_ToRequest_ValidationErrors(t *testing.T) { func TestAIProviderOptions_ToConfig_ValidationErrors(t *testing.T) { isolateSavedAI(t) - if _, err := (AIProviderOptions{Model: "claude-x", Backend: "nope"}).ToConfig(); err == nil || !strings.Contains(err.Error(), "backend") { + if _, err := (AIProviderOptions{ModelFlags: aiflags.ModelFlags{Model: "claude-x", Backend: "nope"}}).ToConfig(); err == nil || !strings.Contains(err.Error(), "backend") { t.Fatalf("expected invalid backend error, got %v", err) } - if _, err := (AIProviderOptions{Model: "claude-x", Budget: "free"}).ToConfig(); err == nil || !strings.Contains(err.Error(), "budget") { + if _, err := (AIProviderOptions{ModelFlags: aiflags.ModelFlags{Model: "claude-x"}, Budget: "free"}).ToConfig(); err == nil || !strings.Contains(err.Error(), "budget") { t.Fatalf("expected invalid budget error, got %v", err) } } +func TestAIProviderOptions_ToConfig_LoadsSchemaRepairDefaults(t *testing.T) { + seedSavedAI(t, ` +ai: + model: claude-sonnet-5 + backend: anthropic +prompts: + schemaRepair: + model: gpt-5 + backend: openai + prompt: /repo/prompts/json-repair.prompt +`) + cfg, err := (AIProviderOptions{}).ToConfig() + if err != nil { + t.Fatalf("ToConfig: %v", err) + } + if cfg.SchemaRepair.Model.Name != "gpt-5" || cfg.SchemaRepair.Model.Backend != api.BackendOpenAI { + t.Fatalf("schema repair model = %#v", cfg.SchemaRepair.Model) + } + if cfg.SchemaRepair.Prompt != "/repo/prompts/json-repair.prompt" { + t.Fatalf("schema repair prompt = %q", cfg.SchemaRepair.Prompt) + } +} + // TestAIRuntimeOptions_ToRequest_MaxTokensPrecedence pins the flag > saved > // built-in order, replacing the old magic-4096 sentinel. func TestAIRuntimeOptions_ToRequest_MaxTokensPrecedence(t *testing.T) { @@ -266,7 +327,7 @@ func TestAIRuntimeOptions_ToRequest_MaxTokensPrecedence(t *testing.T) { func TestAIRuntimeOptions_ToRequest_OverlaysSaved(t *testing.T) { seedSavedAI(t, "ai:\n noMCP: true\n noHooks: true\n noSkills: true\n noUser: true\n noProject: true\n noMemory: true\n maxTokens: 16000\n reasoningEffort: low\n") - opts := AIRuntimeOptions{Effort: "high"} // flag overrides saved low + opts := AIRuntimeOptions{AIProviderOptions: AIProviderOptions{ModelFlags: aiflags.ModelFlags{Effort: "high"}}} // flag overrides saved low req, err := opts.ToRequest("sys", "", "user") if err != nil { t.Fatalf("ToRequest: %v", err) @@ -376,6 +437,23 @@ func TestRunBuffered_JSONIncludesFullInputSpec(t *testing.T) { } } +func TestRunBuffered_PreservesStructuredOutput(t *testing.T) { + got, err := runBuffered(context.Background(), structuredPromptResultProvider{}, ai.Request{}) + if err != nil { + t.Fatal(err) + } + result, ok := got.(AIPromptResult) + if !ok { + t.Fatalf("runBuffered returned %T, want AIPromptResult", got) + } + if result.Text != `{"answer":"42"}` { + t.Fatalf("Text = %q, want JSON transcript text", result.Text) + } + if result.StructuredOutput["answer"] != "42" { + t.Fatalf("StructuredOutput = %#v, want decoded answer", result.StructuredOutput) + } +} + func TestRunStreaming_JSONIncludesFullInputSpec(t *testing.T) { req := ai.Request{ Model: api.Model{Name: "gpt-5-codex", Backend: api.BackendCodexCLI, Effort: api.EffortHigh}, @@ -403,13 +481,38 @@ func TestRunStreaming_JSONIncludesFullInputSpec(t *testing.T) { } } +func TestRunStreaming_StructuredResultIsReturnedOnce(t *testing.T) { + for _, tc := range []struct { + name string + text string + }{ + {name: "result only"}, + {name: "replaces prior text", text: "discarded narrative"}, + } { + t.Run(tc.name, func(t *testing.T) { + got, err := runStreaming(context.Background(), structuredResultStreamingProvider{text: tc.text}, ai.Request{}) + if err != nil { + t.Fatal(err) + } + result, ok := got.(AIPromptResult) + if !ok { + t.Fatalf("runStreaming returned %T, want AIPromptResult", got) + } + if result.Text != `{"answer":"42"}` { + t.Fatalf("Text = %q, want authoritative structured JSON once", result.Text) + } + if result.StructuredOutput["answer"] != "42" { + t.Fatalf("StructuredOutput = %#v, want decoded answer", result.StructuredOutput) + } + }) + } +} + func defaultPromptOptions(t *testing.T) AIPromptOptions { t.Helper() return AIPromptOptions{ - AIRuntimeOptions: AIRuntimeOptions{ - Temperature: "0", - }, - Timeout: "120s", - Prompt: "hello", + AIRuntimeOptions: AIRuntimeOptions{AIProviderOptions: AIProviderOptions{ModelFlags: aiflags.ModelFlags{Temperature: "0"}}}, + Timeout: "120s", + Prompt: "hello", } } diff --git a/pkg/cli/analysis.go b/pkg/cli/analysis.go index d756cb4..04b5e76 100644 --- a/pkg/cli/analysis.go +++ b/pkg/cli/analysis.go @@ -2,15 +2,22 @@ package cli import ( "net/url" + "os" "path/filepath" "sort" "strings" + "sync" "github.com/flanksource/captain/pkg/bash" "github.com/flanksource/captain/pkg/claude" "github.com/flanksource/captain/pkg/claude/tools" ) +// ToolAnalysis is the file/network footprint of one tool use. ReadPaths and +// WritePaths are ABSOLUTE: they are data, shared across sessions and processes +// (gavel stages from them), so they are anchored to the cwd of the tool use that +// produced them rather than relativised against whatever base that call had. +// Relativise with DisplayPath at the point of display. type ToolAnalysis struct { ReadPaths []string `json:"readPaths,omitempty"` WritePaths []string `json:"writePaths,omitempty"` @@ -21,8 +28,11 @@ type ToolAnalysis struct { // AnalyzeToolUseLegacy is the old API for callers still using claude.ToolUse. func AnalyzeToolUseLegacy(tu claude.ToolUse, projectRoot string) ToolAnalysis { base := tools.BaseTool{ - RawTool: tu.Tool, - Input: tu.Input, + RawTool: tu.Tool, + Input: tu.Input, + // CWD anchors the tool's relative paths; without it a bash `cat pkg/x.go` + // could only be guessed at from the project root. + CWD: tu.CWD, ProjectRoot: projectRoot, } return AnalyzeToolUse(tools.NewTool(base)) @@ -32,29 +42,38 @@ func AnalyzeToolUse(t tools.Tool) ToolAnalysis { var a ToolAnalysis base := t.Base() - rel := func(path string) string { - return claude.RelativePath(path, base.ProjectRoot) + abs := func(path string) string { + return claude.AbsolutePath(path, base.CWD, base.ProjectRoot) } switch base.RawTool { case "Read": if path := t.FilePath(); path != "" { - a.ReadPaths = append(a.ReadPaths, rel(path)) + a.ReadPaths = append(a.ReadPaths, abs(path)) } case "Grep": if path, ok := base.Input["path"].(string); ok && path != "" { - a.ReadPaths = append(a.ReadPaths, rel(path)) + a.ReadPaths = append(a.ReadPaths, abs(path)) } case "Glob": if path, ok := base.Input["path"].(string); ok && path != "" { - a.ReadPaths = append(a.ReadPaths, rel(path)) + a.ReadPaths = append(a.ReadPaths, abs(path)) } case "Write", "Edit": if path := t.FilePath(); path != "" { - a.WritePaths = append(a.WritePaths, rel(path)) + a.WritePaths = append(a.WritePaths, abs(path)) } case "Bash": - a.analyzeBash(base.Input, rel) + a.analyzeBash(base.Input, abs) + command, _ := base.Input["command"].(string) + for _, path := range extractApplyPatchPaths(command) { + a.WritePaths = appendUnique(a.WritePaths, abs(path)) + } + case "exec": + input, _ := base.Input["input"].(string) + for _, path := range extractApplyPatchPaths(input) { + a.WritePaths = appendUnique(a.WritePaths, abs(path)) + } case "WebFetch": if urlStr, ok := base.Input["url"].(string); ok { if u, err := url.Parse(urlStr); err == nil && u.Host != "" { @@ -82,7 +101,7 @@ func AnalyzeToolUse(t tools.Tool) ToolAnalysis { return a } -func (a *ToolAnalysis) analyzeBash(input map[string]any, rel func(string) string) { +func (a *ToolAnalysis) analyzeBash(input map[string]any, abs func(string) string) { cmd, _ := input["command"].(string) if cmd == "" { return @@ -95,13 +114,13 @@ func (a *ToolAnalysis) analyzeBash(input map[string]any, rel func(string) string _ = err for _, op := range result.Operations { - a.WritePaths = appendUnique(a.WritePaths, rel(op.Path)) + a.WritePaths = appendUnique(a.WritePaths, abs(op.Path)) } for _, path := range result.ReferencedPaths { - a.ReadPaths = appendUnique(a.ReadPaths, rel(path)) + a.ReadPaths = appendUnique(a.ReadPaths, abs(path)) } for _, path := range extractWritePathsFromBash(cmd) { - a.WritePaths = appendUnique(a.WritePaths, rel(path)) + a.WritePaths = appendUnique(a.WritePaths, abs(path)) } binaries := make(map[string]bool) @@ -156,11 +175,55 @@ func appendUnique(slice []string, val string) []string { return append(slice, val) } +// displayBase is the working directory display paths render against, resolved +// once per process. +var displayBase = sync.OnceValue(func() string { + cwd, err := os.Getwd() + if err != nil { + return "" + } + return cwd +}) + +// DisplayPath renders a canonical (absolute) path for humans: relative to the +// working directory when it sits inside it, otherwise the absolute path — a long +// ../../.. chain reads worse than the full path, and a file from another project +// is genuinely elsewhere. +// +// Display only. Never feed the result back into anything that resolves paths: +// that round trip is what made session paths ambiguous in the first place. +func DisplayPath(path string) string { + if path == "" || !filepath.IsAbs(path) { + return path + } + base := displayBase() + if base == "" { + return path + } + rel, err := filepath.Rel(base, path) + if err != nil || strings.HasPrefix(rel, "..") { + return path + } + return rel +} + +// DisplayPaths maps DisplayPath over a slice, for rendering a path list. +func DisplayPaths(paths []string) []string { + if len(paths) == 0 { + return paths + } + out := make([]string, len(paths)) + for i, p := range paths { + out[i] = DisplayPath(p) + } + return out +} + func FormatPathsWithIcons(readPaths, writePaths []string) string { var parts []string seen := make(map[string]bool) for _, p := range readPaths { - dir := pathToDir(p) + dir := pathToDir(DisplayPath(p)) key := "r:" + dir if !seen[key] { seen[key] = true @@ -168,7 +231,7 @@ func FormatPathsWithIcons(readPaths, writePaths []string) string { } } for _, p := range writePaths { - dir := pathToDir(p) + dir := pathToDir(DisplayPath(p)) key := "w:" + dir if !seen[key] { seen[key] = true diff --git a/pkg/cli/analysis_ginkgo_test.go b/pkg/cli/analysis_ginkgo_test.go new file mode 100644 index 0000000..7a76380 --- /dev/null +++ b/pkg/cli/analysis_ginkgo_test.go @@ -0,0 +1,19 @@ +package cli + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/claude" +) + +var _ = Describe("tool write-path analysis", func() { + It("extracts apply_patch paths embedded in Codex custom tool input", func() { + analysis := AnalyzeToolUseLegacy(claude.ToolUse{ + Tool: "exec", + Input: map[string]any{"input": `const patch = "*** Begin Patch\n*** Update File: /repo/old.go\n*** Move to: /repo/new.go\n-old := \"*** Update File: /repo/ignored.go\"\n+new := true\n*** Add File: /repo/nested/added.go\n*** End Patch"; tools.apply_patch(patch);`}, + }, "/repo") + + Expect(analysis.WritePaths).To(ConsistOf("/repo/old.go", "/repo/new.go", "/repo/nested/added.go")) + }) +}) diff --git a/pkg/cli/analysis_test.go b/pkg/cli/analysis_test.go index fafdff8..368a6a8 100644 --- a/pkg/cli/analysis_test.go +++ b/pkg/cli/analysis_test.go @@ -13,19 +13,43 @@ func TestAnalyzeToolUse_Read(t *testing.T) { Input: map[string]any{"file_path": "/home/user/project/src/main.go"}, }, "/home/user/project") - assert.Equal(t, []string{"src/main.go"}, a.ReadPaths) + assert.Equal(t, []string{"/home/user/project/src/main.go"}, a.ReadPaths) assert.Empty(t, a.WritePaths) assert.Empty(t, a.Binaries) assert.Empty(t, a.Domains) } +// TestAnalyzeToolUse_AnchorsRelativePathsToCWD is the case that made session +// paths unusable: a bash command's relative path means nothing without the +// directory that command ran in, and an agent's cwd moves during a session. +func TestAnalyzeToolUse_AnchorsRelativePathsToCWD(t *testing.T) { + a := AnalyzeToolUseLegacy(claude.ToolUse{ + Tool: "Bash", + Input: map[string]any{"command": "touch out.txt"}, + CWD: "/home/user/project/pkg/cli", + }, "/home/user/project") + + assert.Equal(t, []string{"/home/user/project/pkg/cli/out.txt"}, a.WritePaths) +} + +// A tool use with no cwd recorded falls back to the project root rather than +// leaving the path dangling. +func TestAnalyzeToolUse_FallsBackToProjectRootWithoutCWD(t *testing.T) { + a := AnalyzeToolUseLegacy(claude.ToolUse{ + Tool: "Bash", + Input: map[string]any{"command": "touch out.txt"}, + }, "/home/user/project") + + assert.Equal(t, []string{"/home/user/project/out.txt"}, a.WritePaths) +} + func TestAnalyzeToolUse_Grep(t *testing.T) { a := AnalyzeToolUseLegacy(claude.ToolUse{ Tool: "Grep", Input: map[string]any{"pattern": "TODO", "path": "/home/user/project/pkg/"}, }, "/home/user/project") - assert.Equal(t, []string{"pkg/"}, a.ReadPaths) + assert.Equal(t, []string{"/home/user/project/pkg/"}, a.ReadPaths) assert.Empty(t, a.WritePaths) } @@ -35,7 +59,7 @@ func TestAnalyzeToolUse_Glob(t *testing.T) { Input: map[string]any{"pattern": "**/*.go", "path": "/home/user/project/pkg"}, }, "/home/user/project") - assert.Equal(t, []string{"pkg"}, a.ReadPaths) + assert.Equal(t, []string{"/home/user/project/pkg"}, a.ReadPaths) assert.Empty(t, a.WritePaths) } @@ -46,7 +70,7 @@ func TestAnalyzeToolUse_Write(t *testing.T) { }, "/home/user/project") assert.Empty(t, a.ReadPaths) - assert.Equal(t, []string{"pkg/cli/new.go"}, a.WritePaths) + assert.Equal(t, []string{"/home/user/project/pkg/cli/new.go"}, a.WritePaths) } func TestAnalyzeToolUse_Edit(t *testing.T) { @@ -56,7 +80,7 @@ func TestAnalyzeToolUse_Edit(t *testing.T) { }, "/home/user/project") assert.Empty(t, a.ReadPaths) - assert.Equal(t, []string{"cmd/main.go"}, a.WritePaths) + assert.Equal(t, []string{"/home/user/project/cmd/main.go"}, a.WritePaths) } func TestAnalyzeToolUse_Bash(t *testing.T) { @@ -77,8 +101,8 @@ func TestAnalyzeToolUse_Bash(t *testing.T) { { name: "touch creates file", cmd: "touch /home/user/project/output.txt", - readPaths: []string{"output.txt"}, - writePaths: []string{"output.txt"}, + readPaths: []string{"/home/user/project/output.txt"}, + writePaths: []string{"/home/user/project/output.txt"}, binaries: []string{"touch"}, domains: []string{}, }, @@ -116,7 +140,7 @@ func TestAnalyzeToolUse_Bash(t *testing.T) { { name: "redirect creates write path", cmd: "echo data > /home/user/project/out.txt", - writePaths: []string{"out.txt"}, + writePaths: []string{"/home/user/project/out.txt"}, binaries: []string{}, domains: []string{}, }, diff --git a/pkg/cli/apply_patch.go b/pkg/cli/apply_patch.go new file mode 100644 index 0000000..8148b93 --- /dev/null +++ b/pkg/cli/apply_patch.go @@ -0,0 +1,42 @@ +package cli + +import ( + "encoding/json" + "regexp" + "strings" +) + +var ( + applyPatchFileRE = regexp.MustCompile(`(?m)^\*\*\* (?:Add|Update|Delete) File: ([^\r\n]+)\r?$`) + applyPatchMoveRE = regexp.MustCompile(`(?m)^\*\*\* Move to: ([^\r\n]+)\r?$`) + javascriptStringRE = regexp.MustCompile(`"(?:\\.|[^"\\])*"`) +) + +func extractApplyPatchPaths(input string) []string { + payloads := applyPatchPayloads(input) + var paths []string + for _, payload := range payloads { + for _, pattern := range []*regexp.Regexp{applyPatchFileRE, applyPatchMoveRE} { + for _, match := range pattern.FindAllStringSubmatch(payload, -1) { + if path := strings.TrimSpace(match[1]); path != "" { + paths = append(paths, path) + } + } + } + } + return paths +} + +func applyPatchPayloads(input string) []string { + if !strings.Contains(input, "tools.apply_patch") { + return []string{input} + } + var payloads []string + for _, literal := range javascriptStringRE.FindAllString(input, -1) { + var decoded string + if json.Unmarshal([]byte(literal), &decoded) == nil && strings.Contains(decoded, "*** Begin Patch") { + payloads = append(payloads, decoded) + } + } + return payloads +} diff --git a/pkg/cli/attachments.go b/pkg/cli/attachments.go new file mode 100644 index 0000000..89ee891 --- /dev/null +++ b/pkg/cli/attachments.go @@ -0,0 +1,173 @@ +package cli + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/csv" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/attachments" +) + +func attachmentRefsFromFlags(values []string) ([]api.AttachmentRef, error) { + refs := make([]api.AttachmentRef, 0, len(values)) + for _, value := range values { + reader := csv.NewReader(strings.NewReader(value)) + reader.FieldsPerRecord = -1 + fields, err := reader.Read() + if err != nil { + return nil, fmt.Errorf("parse --attach value %q: %w", value, err) + } + for _, field := range fields { + field = strings.TrimSpace(field) + if len(field) >= 2 && field[0] == '\'' && field[len(field)-1] == '\'' { + field = strings.TrimSpace(field[1 : len(field)-1]) + } + if field == "" { + return nil, fmt.Errorf("empty attachment in --attach value %q", value) + } + ref := api.AttachmentRef{Path: field} + switch { + case strings.HasPrefix(field, api.AttachmentIDPrefix): + ref = api.AttachmentRef{ID: field} + case strings.HasPrefix(field, "http://"), strings.HasPrefix(field, "https://"): + ref = api.AttachmentRef{URL: field} + } + if err := ref.Validate(); err != nil { + return nil, err + } + refs = append(refs, ref) + } + } + return refs, nil +} + +type chatAttachmentResolver struct { + store *attachments.Store +} + +func (r chatAttachmentResolver) Resolve(ctx context.Context, inputs []aichat.AttachmentInput) ([]api.AttachmentRef, error) { + limits := r.store.Limits() + if len(inputs) > limits.MaxFiles { + return nil, fmt.Errorf("attachment request has %d files and exceeds %d file limit", len(inputs), limits.MaxFiles) + } + refs := make([]api.AttachmentRef, 0, len(inputs)) + var total int64 + for i, input := range inputs { + id := input.ID + if id == "" { + if _, suffix, ok := strings.Cut(input.URL, "/api/attachments/"); ok { + id = suffix + } + } + if id != "" { + resolved, err := r.store.Resolve(ctx, []api.AttachmentRef{{ + ID: id, Filename: input.Filename, MediaType: input.MediaType, + }}, "") + if err != nil { + return nil, fmt.Errorf("chat attachment %d: %w", i+1, err) + } + refs = append(refs, resolved[0]) + total += resolved[0].Size + if total > limits.MaxRequestBytes { + return nil, fmt.Errorf("attachments total %d bytes exceeds %d byte request limit", total, limits.MaxRequestBytes) + } + continue + } + mediaType, encoded, ok := parseLegacyDataURL(input.URL) + if !ok { + return nil, fmt.Errorf("chat attachment %d must be uploaded before use", i+1) + } + data, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return nil, fmt.Errorf("chat attachment %d has invalid base64 data: %w", i+1, err) + } + ref, err := r.store.Put(bytes.NewReader(data), input.Filename, mediaType) + if err != nil { + return nil, fmt.Errorf("chat attachment %d: %w", i+1, err) + } + refs = append(refs, ref) + total += ref.Size + if total > limits.MaxRequestBytes { + return nil, fmt.Errorf("attachments total %d bytes exceeds %d byte request limit", total, limits.MaxRequestBytes) + } + } + return refs, nil +} + +func parseLegacyDataURL(value string) (mediaType, encoded string, ok bool) { + header, encoded, ok := strings.Cut(value, ",") + if !ok || !strings.HasPrefix(header, "data:") || !strings.HasSuffix(header, ";base64") { + return "", "", false + } + mediaType = strings.TrimSuffix(strings.TrimPrefix(header, "data:"), ";base64") + return mediaType, encoded, mediaType != "" +} + +func preparePromptAttachments(ctx context.Context, req *ai.Request, cfg ai.Config) error { + if err := resolvePromptAttachments(ctx, req); err != nil { + return err + } + if len(req.Prompt.Attachments) == 0 { + return nil + } + model := req.Model + if model.Name == "" { + model = cfg.Model + } + models := append([]api.Model{model}, model.Fallbacks...) + return ai.ValidateAttachmentCompatibility(models, req.Prompt.Attachments) +} + +func resolvePromptAttachments(ctx context.Context, req *ai.Request) error { + if len(req.Prompt.Attachments) == 0 { + return nil + } + allPrepared := true + for _, attachment := range req.Prompt.Attachments { + allPrepared = allPrepared && attachment.IsPrepared() + } + if !allPrepared { + baseDir := req.Cwd() + if baseDir == "" { + var err error + baseDir, err = os.Getwd() + if err != nil { + return fmt.Errorf("resolve attachment working directory: %w", err) + } + } + store, err := newAttachmentStore(baseDir) + if err != nil { + return err + } + resolved, err := store.Resolve(ctx, req.Prompt.Attachments, baseDir) + if err != nil { + return err + } + req.Prompt.Attachments = resolved + } + return nil +} + +func newAttachmentStore(baseDir string) (*attachments.Store, error) { + defaults := loadSavedConfig().Attachments.WithDefaults() + directory := defaults.Directory + if !filepath.IsAbs(directory) { + directory = filepath.Join(baseDir, directory) + } + return attachments.NewStore(attachments.StoreOptions{ + Directory: directory, + Limits: attachments.Limits{ + MaxFileBytes: defaults.MaxFileBytes, + MaxRequestBytes: defaults.MaxRequestBytes, + MaxFiles: defaults.MaxFiles, + }, + }) +} diff --git a/pkg/cli/attachments_gc.go b/pkg/cli/attachments_gc.go new file mode 100644 index 0000000..12b185d --- /dev/null +++ b/pkg/cli/attachments_gc.go @@ -0,0 +1,126 @@ +package cli + +import ( + "context" + "fmt" + "io/fs" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + "time" + + "github.com/flanksource/captain/pkg/api" +) + +type AttachmentsGCOptions struct { + DryRun bool `flag:"dry-run" help:"Report eligible attachments without deleting them"` +} + +func RunAttachmentsGC(opts AttachmentsGCOptions) (any, error) { + cwd, err := os.Getwd() + if err != nil { + return nil, fmt.Errorf("resolve working directory: %w", err) + } + store, err := newAttachmentStore(cwd) + if err != nil { + return nil, err + } + retention, err := parseAttachmentRetention(loadSavedConfig().Attachments.WithDefaults().Retention) + if err != nil { + return nil, err + } + referenced, err := collectAttachmentReferences(filepath.Join(cwd, ".captain"), store.Directory()) + if err != nil { + return nil, err + } + databaseReferences, err := collectDatabaseAttachmentReferences(context.Background()) + if err != nil { + return nil, err + } + for id := range databaseReferences { + referenced[id] = struct{}{} + } + return store.GC(referenced, retention, opts.DryRun) +} + +func parseAttachmentRetention(value string) (time.Duration, error) { + value = strings.TrimSpace(value) + if days, ok := strings.CutSuffix(value, "d"); ok { + count, err := strconv.Atoi(days) + if err != nil || count < 1 { + return 0, fmt.Errorf("invalid attachment retention %q: use a positive duration such as 30d", value) + } + return time.Duration(count) * 24 * time.Hour, nil + } + duration, err := time.ParseDuration(value) + if err != nil || duration <= 0 { + return 0, fmt.Errorf("invalid attachment retention %q: use a positive duration such as 30d", value) + } + return duration, nil +} + +var attachmentIDPattern = regexp.MustCompile(regexp.QuoteMeta(api.AttachmentIDPrefix) + `[0-9a-fA-F]{64}`) + +func collectAttachmentReferences(root, storeDirectory string) (map[string]struct{}, error) { + references := map[string]struct{}{} + err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + if os.IsNotExist(walkErr) { + return nil + } + return walkErr + } + if entry.IsDir() { + if path == storeDirectory { + return filepath.SkipDir + } + return nil + } + data, err := os.ReadFile(path) + if err != nil { + return fmt.Errorf("read attachment reference source %s: %w", path, err) + } + for id := range attachmentReferencesFromContents([]string{string(data)}) { + references[id] = struct{}{} + } + return nil + }) + if err != nil && !os.IsNotExist(err) { + return nil, fmt.Errorf("scan attachment references: %w", err) + } + return references, nil +} + +func collectDatabaseAttachmentReferences(ctx context.Context) (map[string]struct{}, error) { + db, err := captainDB(ctx) + if err != nil { + return nil, fmt.Errorf("open attachment reference database: %w", err) + } + var rows []struct { + Content string + } + if err := db.Gorm().WithContext(ctx).Raw(` + SELECT rendered_spec::text AS content FROM captain_prompt_runs + UNION ALL + SELECT payload::text AS content FROM captain_events + `).Scan(&rows).Error; err != nil { + return nil, fmt.Errorf("scan database attachment references: %w", err) + } + contents := make([]string, len(rows)) + for i := range rows { + contents[i] = rows[i].Content + } + return attachmentReferencesFromContents(contents), nil +} + +func attachmentReferencesFromContents(contents []string) map[string]struct{} { + references := map[string]struct{}{} + for _, content := range contents { + for _, match := range attachmentIDPattern.FindAllString(content, -1) { + references[strings.ToLower(match)] = struct{}{} + } + } + return references +} diff --git a/pkg/cli/attachments_ginkgo_test.go b/pkg/cli/attachments_ginkgo_test.go new file mode 100644 index 0000000..d35c402 --- /dev/null +++ b/pkg/cli/attachments_ginkgo_test.go @@ -0,0 +1,192 @@ +package cli + +import ( + "bytes" + "context" + "encoding/json" + "github.com/flanksource/captain/pkg/aiflags" + "mime/multipart" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/attachments" +) + +var _ = Describe("attachment flags", func() { + It("uses RFC 4180 parsing for repeated and comma-separated values", func() { + refs, err := attachmentRefsFromFlags([]string{`"reports/q1,q2.pdf",https://example.com/chart.png`, "notes.pdf"}) + Expect(err).NotTo(HaveOccurred()) + Expect(refs).To(HaveLen(3)) + Expect(refs[0].Path).To(Equal("reports/q1,q2.pdf")) + Expect(refs[1].URL).To(Equal("https://example.com/chart.png")) + Expect(refs[2].Path).To(Equal("notes.pdf")) + }) + + It("rejects an empty CSV field", func() { + _, err := attachmentRefsFromFlags([]string{"one.pdf,"}) + Expect(err).To(MatchError(ContainSubstring("empty attachment"))) + }) + + It("removes balanced shell quotes around an absolute path", func() { + path := "/Users/moshe/Desktop/Screenshot 2026-07-15 at 7.37.18.png" + refs, err := attachmentRefsFromFlags([]string{"'" + path + "'"}) + Expect(err).NotTo(HaveOccurred()) + Expect(refs).To(Equal([]api.AttachmentRef{{Path: path}})) + }) + + It("maps the canonical prompt action attachment flag without losing quoted commas", func() { + opts, err := actionFlagsToOptions(map[string]string{ + "attach": `"reports/q1,q2.pdf",notes.pdf`, + }) + Expect(err).NotTo(HaveOccurred()) + refs, err := attachmentRefsFromFlags(opts.Attach) + Expect(err).NotTo(HaveOccurred()) + Expect(refs).To(Equal([]api.AttachmentRef{ + {Path: "reports/q1,q2.pdf"}, + {Path: "notes.pdf"}, + })) + }) + + It("renders an attachment-only canonical prompt", func() { + rendered, err := renderPromptCLI(context.Background(), "", AIPromptOptions{ + AIRuntimeOptions: AIRuntimeOptions{AIProviderOptions: AIProviderOptions{ModelFlags: aiflags.ModelFlags{Model: "gemini-2.5-pro"}}}, + Attach: []string{"diagram.png"}, + }, "", "") + Expect(err).NotTo(HaveOccurred()) + Expect(rendered.ValidationError).To(BeEmpty()) + Expect(rendered.Input.Prompt.Attachments).To(Equal([]api.AttachmentRef{{Path: "diagram.png"}})) + }) + + It("carries prompt workbench attachments into the rendered backend request", func() { + attachment := api.AttachmentRef{ + ID: api.AttachmentIDPrefix + strings.Repeat("a", 64), + Filename: "diagram.png", + MediaType: "image/png", + Size: 512, + } + rendered, err := renderPrompt(context.Background(), "", PromptRenderRequest{Spec: &api.Spec{ + Model: api.Model{Name: "gemini-2.5-pro", Backend: api.BackendGemini}, + Prompt: api.Prompt{Attachments: []api.AttachmentRef{attachment}}, + }}) + + Expect(err).NotTo(HaveOccurred()) + Expect(rendered.ValidationError).To(BeEmpty()) + Expect(rendered.Input.Prompt.Attachments).To(Equal([]api.AttachmentRef{attachment})) + }) + + It("reports an unsupported batch attachment as one failed model", func() { + withGinkgoCaptainDB() + + dir := GinkgoT().TempDir() + path := filepath.Join(dir, "diagram.png") + content := append([]byte("\x89PNG\r\n\x1a\n"), make([]byte, 512)...) + Expect(os.WriteFile(path, content, 0o600)).To(Succeed()) + + originalExecute := executePromptRequestFunc + DeferCleanup(func() { executePromptRequestFunc = originalExecute }) + executePromptRequestFunc = func(ctx context.Context, req ai.Request, cfg ai.Config, _ time.Duration, _ bool) (any, error) { + if err := preparePromptAttachments(ctx, &req, cfg); err != nil { + return nil, err + } + return AIPromptResult{Text: "ok", Model: req.Model.Name, Backend: string(req.Model.Backend)}, nil + } + + req := ai.Request{Prompt: api.Prompt{ + User: "What is this image of?", Attachments: []api.AttachmentRef{{Path: path}}, + }} + req.SetCwd(dir) + result, err := executeSyncBatch(context.Background(), PromptRenderResult{ + Name: "attachment batch", Input: req, + }, AIPromptOptions{MultiModels: []string{"*:sol"}}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.Status).To(Equal("partial")) + Expect(result.Succeeded).To(Equal(3)) + Expect(result.Failed).To(Equal(1)) + Expect(result.Runs).To(ContainElement(And( + HaveField("Backend", string(api.BackendCodexCmux)), + HaveField("Status", "failed"), + HaveField("Error", ContainSubstring("does not accept image/png attachments")), + ))) + }) +}) + +var _ = Describe("attachment HTTP API", func() { + It("uploads once and serves the durable blob by ID", func() { + store, err := attachments.NewStore(attachments.StoreOptions{Directory: filepath.Join(GinkgoT().TempDir(), "attachments")}) + Expect(err).NotTo(HaveOccurred()) + var body bytes.Buffer + writer := multipart.NewWriter(&body) + part, err := writer.CreateFormFile("file", "diagram.png") + Expect(err).NotTo(HaveOccurred()) + content := append([]byte("\x89PNG\r\n\x1a\n"), make([]byte, 512)...) + _, err = part.Write(content) + Expect(err).NotTo(HaveOccurred()) + Expect(writer.Close()).To(Succeed()) + + upload := httptest.NewRequest(http.MethodPost, "/api/attachments", &body) + upload.Header.Set("Content-Type", writer.FormDataContentType()) + uploadResponse := httptest.NewRecorder() + handleAttachmentUpload(store)(uploadResponse, upload) + Expect(uploadResponse.Code).To(Equal(http.StatusCreated)) + var ref api.AttachmentRef + Expect(json.Unmarshal(uploadResponse.Body.Bytes(), &ref)).To(Succeed()) + + download := httptest.NewRequest(http.MethodGet, "/api/attachments/"+ref.ID, nil) + download.SetPathValue("id", ref.ID) + downloadResponse := httptest.NewRecorder() + handleAttachmentGet(store)(downloadResponse, download) + Expect(downloadResponse.Code).To(Equal(http.StatusOK)) + Expect(downloadResponse.Body.Bytes()).To(Equal(content)) + Expect(downloadResponse.Header().Get("Content-Type")).To(Equal("image/png")) + }) +}) + +var _ = Describe("chat attachment resolver", func() { + It("migrates a legacy data URL into the durable store", func() { + store, err := attachments.NewStore(attachments.StoreOptions{Directory: filepath.Join(GinkgoT().TempDir(), "attachments")}) + Expect(err).NotTo(HaveOccurred()) + refs, err := (chatAttachmentResolver{store: store}).Resolve(context.Background(), []aichat.AttachmentInput{{ + URL: "data:image/png;base64,iVBORw0KGgo=", Filename: "legacy.png", MediaType: "image/png", + }}) + Expect(err).NotTo(HaveOccurred()) + Expect(refs).To(HaveLen(1)) + Expect(refs[0].ID).To(HavePrefix(api.AttachmentIDPrefix)) + Expect(refs[0].IsPrepared()).To(BeTrue()) + }) + + It("enforces aggregate request limits while migrating legacy parts", func() { + store, err := attachments.NewStore(attachments.StoreOptions{ + Directory: filepath.Join(GinkgoT().TempDir(), "attachments"), + Limits: attachments.Limits{ + MaxFileBytes: 1024, + MaxRequestBytes: 1024, + MaxFiles: 1, + }, + }) + Expect(err).NotTo(HaveOccurred()) + _, err = (chatAttachmentResolver{store: store}).Resolve(context.Background(), []aichat.AttachmentInput{ + {URL: "data:image/png;base64,iVBORw0KGgo=", Filename: "one.png", MediaType: "image/png"}, + {URL: "data:image/png;base64,iVBORw0KGgo=", Filename: "two.png", MediaType: "image/png"}, + }) + Expect(err).To(MatchError(ContainSubstring("exceeds 1 file limit"))) + }) +}) + +var _ = Describe("attachment garbage collection references", func() { + It("retains durable IDs found in database JSON", func() { + id := api.AttachmentIDPrefix + "ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789" + references := attachmentReferencesFromContents([]string{`{"prompt":{"attachments":[{"id":"` + id + `"}]}}`}) + Expect(references).To(HaveKey(strings.ToLower(id))) + }) +}) diff --git a/pkg/cli/attachments_http.go b/pkg/cli/attachments_http.go new file mode 100644 index 0000000..93a8629 --- /dev/null +++ b/pkg/cli/attachments_http.go @@ -0,0 +1,88 @@ +package cli + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/attachments" +) + +func handleAttachmentUpload(store *attachments.Store) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + r.Body = http.MaxBytesReader(w, r.Body, store.Limits().MaxRequestBytes) + reader, err := r.MultipartReader() + if err != nil { + writeAttachmentError(w, http.StatusBadRequest, fmt.Errorf("read multipart upload: %w", err)) + return + } + var uploaded *api.AttachmentRef + for { + part, err := reader.NextPart() + if err == io.EOF { + break + } + if err != nil { + writeAttachmentError(w, http.StatusBadRequest, fmt.Errorf("read multipart upload: %w", err)) + return + } + if part.FormName() != "file" || part.FileName() == "" { + _ = part.Close() + continue + } + if uploaded != nil { + _ = part.Close() + writeAttachmentError(w, http.StatusBadRequest, fmt.Errorf("upload exactly one file per request")) + return + } + ref, err := store.Put(part, part.FileName(), part.Header.Get("Content-Type")) + _ = part.Close() + if err != nil { + writeAttachmentError(w, http.StatusBadRequest, err) + return + } + uploaded = &ref + } + if uploaded == nil { + writeAttachmentError(w, http.StatusBadRequest, fmt.Errorf("multipart field file is required")) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + _ = json.NewEncoder(w).Encode(uploaded) + } +} + +func handleAttachmentGet(store *attachments.Store) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + file, err := store.Open(r.PathValue("id")) + if err != nil { + writeAttachmentError(w, http.StatusNotFound, err) + return + } + defer file.Close() + sample := make([]byte, 512) + read, err := file.Read(sample) + if err != nil && err != io.EOF { + writeAttachmentError(w, http.StatusInternalServerError, err) + return + } + if _, err := file.Seek(0, io.SeekStart); err != nil { + writeAttachmentError(w, http.StatusInternalServerError, err) + return + } + w.Header().Set("Content-Type", http.DetectContentType(sample[:read])) + w.Header().Set("X-Content-Type-Options", "nosniff") + if _, err := io.Copy(w, file); err != nil { + return + } + } +} + +func writeAttachmentError(w http.ResponseWriter, status int, err error) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) +} diff --git a/pkg/cli/changes.go b/pkg/cli/changes.go index 601f513..5fbc24a 100644 --- a/pkg/cli/changes.go +++ b/pkg/cli/changes.go @@ -6,6 +6,8 @@ import ( "sort" "time" + "github.com/flanksource/clicky/api" + "github.com/flanksource/captain/pkg/claude" ) @@ -20,13 +22,25 @@ type ChangesOptions struct { Ignored bool `flag:"ignored" help:"Include gitignored / out-of-repo files; --ignored=true to show them"` } +// SessionPath is an absolute path to a file a session touched. It serialises +// absolute — consumers resolve it from another process and another working +// directory (gavel stages from these) — and renders relative to the working +// directory, which is all a human reading the table wants to see. +type SessionPath string + +func (p SessionPath) Pretty() api.Text { + return api.Text{Content: DisplayPath(string(p))} +} + +func (p SessionPath) String() string { return string(p) } + // ChangedFile describes a single file modified during a session. type ChangedFile struct { - Path string `json:"path" pretty:"label=File,table"` - Edits int `json:"edits" pretty:"label=Edits,table"` - Tools string `json:"tools" pretty:"label=Tools,table"` - Agent string `json:"agent,omitempty" pretty:"label=Agent,table"` - Last string `json:"last,omitempty" pretty:"label=Last Modified,table"` + Path SessionPath `json:"path" pretty:"label=File,table"` + Edits int `json:"edits" pretty:"label=Edits,table"` + Tools string `json:"tools" pretty:"label=Tools,table"` + Agent string `json:"agent,omitempty" pretty:"label=Agent,table"` + Last string `json:"last,omitempty" pretty:"label=Last Modified,table"` } // ChangesResult lists the files modified by a single session. @@ -156,7 +170,7 @@ func buildChangesResult(sessionID string, uses []claude.ToolUse) ChangesResult { Files: make([]ChangedFile, 0, len(files)), } for path, a := range files { - row := ChangedFile{Path: path, Edits: a.edits, Tools: joinSorted(a.tools), Agent: joinSorted(a.agents)} + row := ChangedFile{Path: SessionPath(path), Edits: a.edits, Tools: joinSorted(a.tools), Agent: joinSorted(a.agents)} if a.last != nil { row.Last = a.last.Format("2006-01-02 15:04") } @@ -184,7 +198,7 @@ func agentLabel(tu claude.ToolUse) string { func (r *ChangesResult) filter(pf *pathFilter) { kept := r.Files[:0] for _, f := range r.Files { - if pf.keep(f.Path) { + if pf.keep(string(f.Path)) { kept = append(kept, f) } } diff --git a/pkg/cli/changes_test.go b/pkg/cli/changes_test.go index b3819ba..c5f2862 100644 --- a/pkg/cli/changes_test.go +++ b/pkg/cli/changes_test.go @@ -1,6 +1,8 @@ package cli import ( + "os" + "path/filepath" "testing" "time" @@ -45,12 +47,30 @@ func TestBuildChangesResult_AggregatesWritePaths(t *testing.T) { // Only main.go (Write+Edit) and util.go (Write) count; readme.md was only Read. assert.Equal(t, 2, result.FileCount) + // Paths are absolute: consumers (gavel staging) resolve them from another + // process whose working directory is unrelated to the session's. // main.go has the most edits, so it sorts first. - assert.Equal(t, "main.go", result.Files[0].Path) + assert.Equal(t, SessionPath(root+"/main.go"), result.Files[0].Path) assert.Equal(t, 2, result.Files[0].Edits) assert.Equal(t, "Edit, Write", result.Files[0].Tools) assert.Equal(t, now.Format("2006-01-02 15:04"), result.Files[0].Last) - assert.Equal(t, "util.go", result.Files[1].Path) + assert.Equal(t, SessionPath(root+"/util.go"), result.Files[1].Path) assert.Equal(t, 1, result.Files[1].Edits) } + +// TestSessionPathRendersRelativeToWorkingDir pins the display half of the +// contract: the value stays absolute, only its rendering is relative. +func TestSessionPathRendersRelativeToWorkingDir(t *testing.T) { + cwd, err := os.Getwd() + assert.NoError(t, err) + + inside := SessionPath(filepath.Join(cwd, "pkg", "cli", "changes.go")) + assert.Equal(t, "pkg/cli/changes.go", inside.Pretty().Content) + assert.Equal(t, filepath.Join(cwd, "pkg", "cli", "changes.go"), inside.String(), "the value itself stays absolute") + + // A file outside the working directory reads better absolute than as a + // ../../.. chain. + outside := SessionPath("/tmp/elsewhere/x.go") + assert.Equal(t, "/tmp/elsewhere/x.go", outside.Pretty().Content) +} diff --git a/pkg/cli/chat_thread_store.go b/pkg/cli/chat_thread_store.go index fb9a24c..13aaf5c 100644 --- a/pkg/cli/chat_thread_store.go +++ b/pkg/cli/chat_thread_store.go @@ -13,7 +13,7 @@ import ( "sync" "time" - "github.com/flanksource/clicky/aichat" + "github.com/flanksource/captain/pkg/aichat" ) type fileThreadStore struct { @@ -152,7 +152,10 @@ func (s *fileThreadStore) AddUsage(_ context.Context, id string, usage aichat.Tu } thread.TotalInputTokens += usage.InputTokens thread.TotalOutputTokens += usage.OutputTokens - thread.TotalCostUsd += usage.CostUSD + thread.TotalReasoningTokens += usage.ReasoningTokens + thread.TotalCacheReadTokens += usage.CacheReadTokens + thread.TotalCacheWriteTokens += usage.CacheWriteTokens + thread.TotalCostUSD += usage.CostUSD thread.LastContextTokens = usage.InputTokens thread.UpdatedAt = time.Now() if err := s.saveLocked(state); err != nil { diff --git a/pkg/cli/chat_thread_store_test.go b/pkg/cli/chat_thread_store_test.go index a404ae9..df36a59 100644 --- a/pkg/cli/chat_thread_store_test.go +++ b/pkg/cli/chat_thread_store_test.go @@ -5,7 +5,7 @@ import ( "path/filepath" "testing" - "github.com/flanksource/clicky/aichat" + "github.com/flanksource/captain/pkg/aichat" ) func TestFileThreadStorePersistsThreads(t *testing.T) { @@ -57,8 +57,8 @@ func TestFileThreadStorePersistsThreads(t *testing.T) { if len(got.Messages) != 1 || got.Messages[0].Parts[0].Text != "continue" { t.Errorf("Messages = %+v", got.Messages) } - if got.TotalInputTokens != 10 || got.TotalOutputTokens != 5 || got.TotalCostUsd != 0.25 { - t.Errorf("usage totals = input %d output %d cost %f", got.TotalInputTokens, got.TotalOutputTokens, got.TotalCostUsd) + if got.TotalInputTokens != 10 || got.TotalOutputTokens != 5 || got.TotalCostUSD != 0.25 { + t.Errorf("usage totals = input %d output %d cost %f", got.TotalInputTokens, got.TotalOutputTokens, got.TotalCostUSD) } list, err := reloaded.List(ctx) diff --git a/pkg/cli/command_alias.go b/pkg/cli/command_alias.go new file mode 100644 index 0000000..921d020 --- /dev/null +++ b/pkg/cli/command_alias.go @@ -0,0 +1,43 @@ +package cli + +import ( + "fmt" + + "github.com/spf13/cobra" +) + +type CommandAliasOptions struct { + Name string + Short string + Hidden bool + Root *cobra.Command + Target []string +} + +func NewCommandAlias(opts CommandAliasOptions) *cobra.Command { + if opts.Name == "" { + panic("command alias name is required") + } + if opts.Root == nil { + panic(fmt.Sprintf("command alias %s root is required", opts.Name)) + } + if len(opts.Target) == 0 { + panic(fmt.Sprintf("command alias %s target is required", opts.Name)) + } + return &cobra.Command{ + Use: opts.Name, + Short: opts.Short, + Hidden: opts.Hidden, + DisableFlagParsing: true, + RunE: func(_ *cobra.Command, args []string) error { + forwarded := make([]string, 0, len(opts.Target)+len(args)) + forwarded = append(forwarded, opts.Target...) + forwarded = append(forwarded, args...) + opts.Root.SetArgs(forwarded) + silenceErrors := opts.Root.SilenceErrors + opts.Root.SilenceErrors = true + defer func() { opts.Root.SilenceErrors = silenceErrors }() + return opts.Root.Execute() + }, + } +} diff --git a/pkg/cli/command_alias_ginkgo_test.go b/pkg/cli/command_alias_ginkgo_test.go new file mode 100644 index 0000000..1351b23 --- /dev/null +++ b/pkg/cli/command_alias_ginkgo_test.go @@ -0,0 +1,61 @@ +package cli + +import ( + "bytes" + "errors" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/spf13/cobra" +) + +var _ = Describe("command aliases", func() { + It("redispatches raw arguments through the canonical command", func() { + root := &cobra.Command{Use: "captain", SilenceErrors: true, SilenceUsage: true} + prompt := &cobra.Command{Use: "prompt"} + var received []string + prompt.AddCommand(&cobra.Command{ + Use: "run", + DisableFlagParsing: true, + RunE: func(_ *cobra.Command, args []string) error { + received = append([]string(nil), args...) + return nil + }, + }) + root.AddCommand(prompt) + ai := &cobra.Command{Use: "ai"} + ai.AddCommand(NewCommandAlias(CommandAliasOptions{ + Name: "prompt", Root: root, Target: []string{"prompt", "run"}, + })) + root.AddCommand(ai) + + args := []string{"--attach", "diagram.png", "-p", "describe it"} + root.SetArgs(append([]string{"ai", "prompt"}, args...)) + Expect(root.Execute()).To(Succeed()) + Expect(received).To(Equal(args)) + }) + + It("prints a canonical command error once", func() { + root := &cobra.Command{Use: "captain", SilenceUsage: true} + var output bytes.Buffer + root.SetErr(&output) + prompt := &cobra.Command{Use: "prompt"} + prompt.AddCommand(&cobra.Command{ + Use: "run", + RunE: func(_ *cobra.Command, _ []string) error { + return errors.New("canonical failure") + }, + }) + root.AddCommand(prompt) + ai := &cobra.Command{Use: "ai"} + ai.AddCommand(NewCommandAlias(CommandAliasOptions{ + Name: "prompt", Root: root, Target: []string{"prompt", "run"}, + })) + root.AddCommand(ai) + root.SetArgs([]string{"ai", "prompt"}) + + Expect(root.Execute()).To(MatchError("canonical failure")) + Expect(strings.Count(output.String(), "Error: canonical failure")).To(Equal(1)) + }) +}) diff --git a/pkg/cli/configure.go b/pkg/cli/configure.go index 087a1b0..3228a51 100644 --- a/pkg/cli/configure.go +++ b/pkg/cli/configure.go @@ -9,10 +9,21 @@ import ( "github.com/charmbracelet/huh" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/captainconfig" + clickyrpc "github.com/flanksource/clicky/rpc" + "github.com/flanksource/clicky/text" ) -type ConfigureOptions struct{} +type ConfigureOptions struct { + Provider string `flag:"provider" args:"true" help:"API provider to configure: anthropic|openai|gemini|deepseek"` + Token text.SensitiveString `flag:"token" hidden:"true" help:"Provider API token (prefer the secure interactive prompt)"` + Test bool `flag:"test" help:"Test the current or supplied token without saving it"` + Agent string `flag:"agent" help:"Default runtime agent for this provider"` + Model string `flag:"model" help:"Default model for this provider"` + Effort string `flag:"effort" help:"Default reasoning effort, or default to use the model default"` + Active bool `flag:"active" help:"Use this provider for completely flagless runs"` +} type ConfigureResult struct { Path string `json:"path" pretty:"label=Saved To"` @@ -39,15 +50,33 @@ const ( // the form predictable and the test below straightforward. var allToggles = []string{toggleCaching, toggleMCP, toggleHooks, toggleSkills, toggleUser, toggleProject, toggleMemory} -func RunConfigure(opts ConfigureOptions) (any, error) { +func RunConfigure(ctx context.Context, opts ConfigureOptions) (any, error) { + if _, ok := clickyrpc.RequestFromContext(ctx); ok { + return nil, fmt.Errorf("configure is unavailable over generated RPC; use the guarded provider configuration APIs") + } + if strings.TrimSpace(opts.Provider) != "" { + return runProviderConfigure(ctx, opts) + } + if !opts.Token.IsEmpty() || opts.Test || opts.Agent != "" || opts.Model != "" || opts.Effort != "" || opts.Active { + return nil, fmt.Errorf("provider is required when using token or provider-default flags") + } + return runConfigureWizard() +} + +func runConfigureWizard() (any, error) { current, _, err := captainconfig.Load() if err != nil { return nil, err } - backend := defaultString(current.AI.Backend, string(ai.BackendAnthropic)) - model := defaultString(current.AI.Model, defaultModelFor(ai.Backend(backend))) - effort := defaultString(current.AI.ReasoningEffort, "high") + activeProvider := api.Backend(current.AI.ActiveProvider()) + currentDefaults, err := effectiveProviderDefaults(current.AI, activeProvider) + if err != nil { + return nil, err + } + backend := currentDefaults.Agent + model := currentDefaults.Model + effort := defaultString(currentDefaults.Effort, "high") budget := floatToInput(current.AI.BudgetUSD) maxTokens := intToInput(current.AI.MaxTokens) timeout := defaultString(current.AI.Timeout, "120s") @@ -71,12 +100,10 @@ func RunConfigure(opts ConfigureOptions) (any, error) { Value(&model), huh.NewSelect[string](). Title("Reasoning effort"). - Description("Honoured by codex-cli and the API backends (thinking budget); CLI wrappers may ignore."). - Options( - huh.NewOption("low", "low"), - huh.NewOption("medium", "medium"), - huh.NewOption("high", "high"), - ). + Description("Honoured by codex-agent and the API backends (thinking budget); CLI wrappers may ignore."). + OptionsFunc(func() []huh.Option[string] { + return effortHuhOptionsFor(ai.Backend(backend), model) + }, []any{&backend, &model}). Value(&effort), huh.NewInput(). Title("Budget (USD)"). @@ -116,6 +143,13 @@ func RunConfigure(opts ConfigureOptions) (any, error) { Timeout: timeout, Enabled: enabled, }) + for provider, defaults := range current.AI.Providers { + if _, exists := cfg.AI.Providers[provider]; !exists { + cfg.AI.Providers[provider] = defaults + } + } + cfg.Prompts = current.Prompts + cfg.Attachments = current.Attachments if err := captainconfig.Save(cfg); err != nil { return nil, err } @@ -123,9 +157,9 @@ func RunConfigure(opts ConfigureOptions) (any, error) { path, _ := captainconfig.Path() return ConfigureResult{ Path: path, - Backend: cfg.AI.Backend, - Model: cfg.AI.Model, - ReasoningEffort: cfg.AI.ReasoningEffort, + Backend: backend, + Model: model, + ReasoningEffort: effort, BudgetUSD: budget, MaxTokens: maxTokens, Timeout: cfg.AI.Timeout, @@ -155,21 +189,23 @@ func buildConfigFromForm(in formInputs) captainconfig.Config { enabled[e] = true } + provider := api.Backend(in.Backend).Provider() return captainconfig.Config{ AI: captainconfig.AIDefaults{ - Backend: in.Backend, - Model: in.Model, - ReasoningEffort: in.ReasoningEffort, - BudgetUSD: budget, - MaxTokens: maxTokens, - Timeout: strings.TrimSpace(in.Timeout), - NoCache: !enabled[toggleCaching], - NoMCP: !enabled[toggleMCP], - NoHooks: !enabled[toggleHooks], - NoSkills: !enabled[toggleSkills], - NoUser: !enabled[toggleUser], - NoProject: !enabled[toggleProject], - NoMemory: !enabled[toggleMemory], + DefaultProvider: string(provider), + Providers: map[string]captainconfig.ProviderDefaults{ + string(provider): {Agent: in.Backend, Model: in.Model, ReasoningEffort: in.ReasoningEffort}, + }, + BudgetUSD: budget, + MaxTokens: maxTokens, + Timeout: strings.TrimSpace(in.Timeout), + NoCache: !enabled[toggleCaching], + NoMCP: !enabled[toggleMCP], + NoHooks: !enabled[toggleHooks], + NoSkills: !enabled[toggleSkills], + NoUser: !enabled[toggleUser], + NoProject: !enabled[toggleProject], + NoMemory: !enabled[toggleMemory], }, } } @@ -211,6 +247,7 @@ func backendOptions() []huh.Option[string] { huh.NewOption("Claude Agent (SDK)", string(ai.BackendClaudeAgent)), huh.NewOption("Claude cmux", string(ai.BackendClaudeCmux)), huh.NewOption("Codex CLI", string(ai.BackendCodexCLI)), + huh.NewOption("Codex Agent (app-server)", string(ai.BackendCodexAgent)), huh.NewOption("Codex cmux", string(ai.BackendCodexCmux)), huh.NewOption("Gemini CLI", string(ai.BackendGeminiCLI)), } @@ -265,28 +302,50 @@ func modelHuhOptions(models []ai.ModelDef) []huh.Option[string] { } // defaultModelFor returns a hard-coded picker default per backend that seeds the -// form. CLI/agent backends use the catalog slug their picker actually lists -// (agentCatalogModels) so the seeded default is a selectable option. API +// form. CLI/agent backends use exact provider model IDs from the catalog so the +// seeded default is a selectable option. API // backends have no "default" flag on /v1/models, so we use the most-current id // we expect each provider to keep stable; the user can pick anything else. func defaultModelFor(b ai.Backend) string { switch b { case ai.BackendAnthropic: return "claude-sonnet-5" - case ai.BackendClaudeCLI, ai.BackendClaudeAgent: - return "claude-agent-sonnet" + case ai.BackendClaudeCLI, ai.BackendClaudeAgent, ai.BackendClaudeCmux: + return "claude-sonnet-5" case ai.BackendOpenAI: - return "gpt-5.5" + return "gpt-5.6" case ai.BackendDeepSeek: return "deepseek-reasoner" - case ai.BackendCodexCLI: - return "gpt-5-codex" + case ai.BackendCodexCLI, ai.BackendCodexAgent, ai.BackendCodexCmux: + return "gpt-5.6-sol" case ai.BackendGemini, ai.BackendGeminiCLI: return "gemini-3.5-flash" } return "" } +func effortHuhOptions() []huh.Option[string] { + return effortOptions(api.AllEfforts()) +} + +func effortHuhOptionsFor(backend ai.Backend, model string) []huh.Option[string] { + if supported, _, ok := ai.ModelEfforts(backend, model); ok { + if len(supported) == 0 { + return []huh.Option[string]{huh.NewOption("Backend default", "")} + } + return effortOptions(supported) + } + return effortHuhOptions() +} + +func effortOptions(efforts []api.Effort) []huh.Option[string] { + out := make([]huh.Option[string], 0, len(efforts)) + for _, effort := range efforts { + out = append(out, huh.NewOption(string(effort), string(effort))) + } + return out +} + func defaultString(v, fallback string) string { if strings.TrimSpace(v) == "" { return fallback diff --git a/pkg/cli/configure_provider.go b/pkg/cli/configure_provider.go new file mode 100644 index 0000000..6328d52 --- /dev/null +++ b/pkg/cli/configure_provider.go @@ -0,0 +1,228 @@ +package cli + +import ( + "context" + "fmt" + "os" + "strings" + + "github.com/charmbracelet/huh" + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" + "github.com/flanksource/captain/pkg/credentials" +) + +type ConfigureTokenResult struct { + Path string `json:"path" pretty:"label=Vault"` + Provider string `json:"provider" pretty:"label=Provider"` + MaskedToken string `json:"maskedToken" pretty:"label=Token"` + Source string `json:"source" pretty:"label=Source"` + ModelCount int `json:"modelCount" pretty:"label=Models"` + Saved bool `json:"saved" pretty:"label=Saved"` +} + +type ConfigureProviderDefaultsResult struct { + Path string `json:"path" pretty:"label=Saved To"` + Provider string `json:"provider" pretty:"label=Provider"` + Agent string `json:"agent" pretty:"label=Agent"` + Model string `json:"model" pretty:"label=Model"` + Effort string `json:"effort,omitempty" pretty:"label=Effort"` + Active bool `json:"active" pretty:"label=Active Provider"` +} + +var ( + configureTokenModels = ai.ListModelsWithAPIKey + configureDefaultsModels = modelsForConfigurableAgent +) + +func runProviderConfigure(ctx context.Context, opts ConfigureOptions) (any, error) { + provider := api.Backend(strings.TrimSpace(opts.Provider)) + if !configurableAPIBackend(provider) { + return nil, fmt.Errorf("provider must be one of: anthropic, openai, gemini, deepseek (got %q)", opts.Provider) + } + hasDefaults := opts.Agent != "" || opts.Model != "" || opts.Effort != "" || opts.Active + if hasDefaults { + if !opts.Token.IsEmpty() || opts.Test { + return nil, fmt.Errorf("token flags cannot be combined with agent, model, effort, or active defaults") + } + return runProviderDefaultsConfigure(ctx, provider, opts) + } + return runProviderTokenConfigure(ctx, provider, opts) +} + +func runProviderDefaultsConfigure(ctx context.Context, provider api.Backend, opts ConfigureOptions) (ConfigureProviderDefaultsResult, error) { + saved, _, err := captainconfig.Load() + if err != nil { + return ConfigureProviderDefaultsResult{}, err + } + current, err := effectiveProviderDefaults(saved.AI, provider) + if err != nil { + return ConfigureProviderDefaultsResult{}, err + } + if opts.Active && opts.Agent == "" && opts.Model == "" && opts.Effort == "" { + if err := captainconfig.Update(func(cfg *captainconfig.Config) error { + cfg.AI.DefaultProvider = string(provider) + return nil + }); err != nil { + return ConfigureProviderDefaultsResult{}, err + } + path, _ := captainconfig.Path() + return ConfigureProviderDefaultsResult{ + Path: path, Provider: string(provider), Agent: current.Agent, + Model: current.Model, Effort: current.Effort, Active: true, + }, nil + } + next := current + selectionChanged := false + if opts.Agent != "" { + next.Agent = strings.TrimSpace(opts.Agent) + if next.Agent != current.Agent && opts.Model == "" { + next.Model = defaultModelFor(api.Backend(next.Agent)) + } + selectionChanged = next.Agent != current.Agent + } + if opts.Model != "" { + next.Model = strings.TrimSpace(opts.Model) + selectionChanged = selectionChanged || next.Model != current.Model + } + if opts.Effort != "" { + if strings.EqualFold(strings.TrimSpace(opts.Effort), "default") { + next.Effort = "" + } else { + next.Effort = strings.TrimSpace(opts.Effort) + } + } else if selectionChanged { + next.Effort = defaultEffortFor(api.Backend(next.Agent), next.Model) + } + if err := validateProviderDefaults(ctx, provider, next); err != nil { + return ConfigureProviderDefaultsResult{}, err + } + if err := captainconfig.Update(func(cfg *captainconfig.Config) error { + if cfg.AI.Providers == nil { + cfg.AI.Providers = map[string]captainconfig.ProviderDefaults{} + } + cfg.AI.Providers[string(provider)] = captainconfig.ProviderDefaults{ + Agent: next.Agent, Model: next.Model, ReasoningEffort: next.Effort, + } + if opts.Active { + cfg.AI.DefaultProvider = string(provider) + } + return nil + }); err != nil { + return ConfigureProviderDefaultsResult{}, err + } + path, _ := captainconfig.Path() + return ConfigureProviderDefaultsResult{ + Path: path, Provider: string(provider), Agent: next.Agent, Model: next.Model, + Effort: next.Effort, Active: opts.Active || saved.AI.ActiveProvider() == string(provider), + }, nil +} + +func defaultEffortFor(agent api.Backend, model string) string { + _, effort, ok := ai.ModelEfforts(agent, model) + if !ok { + return "" + } + return string(effort) +} + +func validateProviderDefaults(ctx context.Context, provider api.Backend, defaults ProviderDefaultView) error { + agent := api.Backend(strings.TrimSpace(defaults.Agent)) + if agent.Provider() != provider { + return fmt.Errorf("agent %q does not belong to provider %q", agent, provider) + } + models, err := configureDefaultsModels(ctx, agent) + if err != nil { + return fmt.Errorf("list models for %s: %w", agent, err) + } + model := strings.TrimSpace(defaults.Model) + found := false + for _, candidate := range models { + if candidate.ID == model { + found = true + break + } + } + if !found { + return fmt.Errorf("model %q is not available for agent %q", model, agent) + } + if err := ai.ValidateModelEffort(agent, model, api.Effort(defaults.Effort)); err != nil { + return err + } + return nil +} + +func modelsForConfigurableAgent(ctx context.Context, agent api.Backend) ([]ai.ModelDef, error) { + if agent.Kind() == "api" { + return ai.ListModels(ctx, agent) + } + return agentCatalogModels(agent), nil +} + +func runProviderTokenConfigure(ctx context.Context, backend api.Backend, opts ConfigureOptions) (ConfigureTokenResult, error) { + vault, err := credentials.DefaultVault() + if err != nil { + return ConfigureTokenResult{}, err + } + token := strings.TrimSpace(opts.Token.Value()) + source := "candidate" + if token == "" && opts.Test { + resolved, err := ai.ResolveAPIKey(backend) + if err != nil { + return ConfigureTokenResult{}, err + } + token, source = resolved.Token, resolved.Source + if token == "" { + return ConfigureTokenResult{}, fmt.Errorf("no credential configured for %s", backend) + } + } else if token == "" { + token, err = promptConfigureToken(backend) + if err != nil { + return ConfigureTokenResult{}, err + } + } + models, err := configureTokenModels(ctx, backend, token) + if err != nil { + return ConfigureTokenResult{}, fmt.Errorf("validate %s credential: %w", backend, err) + } + result := ConfigureTokenResult{ + Path: vault.Path(), Provider: string(backend), MaskedToken: ai.MaskKey(token), + Source: source, ModelCount: len(models), Saved: !opts.Test, + } + if !opts.Test { + if err := vault.Set(string(backend), token); err != nil { + return ConfigureTokenResult{}, err + } + result.Source = credentials.SourceVault + } + return result, nil +} + +func configurableAPIBackend(backend api.Backend) bool { + return backend.Provider() == backend +} + +func promptConfigureToken(backend api.Backend) (string, error) { + info, err := os.Stdin.Stat() + if err != nil { + return "", fmt.Errorf("inspect stdin: %w", err) + } + if info.Mode()&os.ModeCharDevice == 0 { + return "", fmt.Errorf("token is required in non-interactive mode; pass --token") + } + var token string + form := huh.NewForm(huh.NewGroup( + huh.NewInput().Title(fmt.Sprintf("%s API token", backend)).EchoMode(huh.EchoModePassword). + Value(&token).Validate(func(value string) error { + if strings.TrimSpace(value) == "" { + return fmt.Errorf("token cannot be empty") + } + return nil + }), + )) + if err := form.Run(); err != nil { + return "", err + } + return strings.TrimSpace(token), nil +} diff --git a/pkg/cli/configure_test.go b/pkg/cli/configure_test.go index 2545e53..b72b8e8 100644 --- a/pkg/cli/configure_test.go +++ b/pkg/cli/configure_test.go @@ -1,15 +1,184 @@ package cli import ( + "context" + "errors" + "net/http/httptest" + "path/filepath" "reflect" "sort" "strings" "testing" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/captainconfig" + "github.com/flanksource/captain/pkg/credentials" + clickyrpc "github.com/flanksource/clicky/rpc" + "github.com/flanksource/clicky/text" ) +func TestRunConfigureRejectsGeneratedRPCInvocation(t *testing.T) { + ctx := clickyrpc.ContextWithRequest(context.Background(), httptest.NewRequest("POST", "/api/v1/configure", nil)) + if _, err := RunConfigure(ctx, ConfigureOptions{Provider: "openai", Test: true}); err == nil { + t.Fatal("generated configure RPC must direct callers to the guarded provider-token endpoints") + } +} + +func TestRunProviderConfigureTestsBeforeSaving(t *testing.T) { + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) + previous := configureTokenModels + configureTokenModels = func(_ context.Context, backend ai.Backend, token string) ([]ai.ModelDef, error) { + if backend != ai.BackendOpenAI || token != "candidate-secret" { + t.Fatalf("validation got backend=%s token=%q", backend, token) + } + return []ai.ModelDef{{ID: "gpt-example"}}, nil + } + t.Cleanup(func() { configureTokenModels = previous }) + + result, err := runProviderTokenConfigure(context.Background(), ai.BackendOpenAI, ConfigureOptions{ + Token: text.SensitiveString("candidate-secret"), + }) + if err != nil { + t.Fatalf("runProviderConfigure: %v", err) + } + if result.Provider != "openai" || result.ModelCount != 1 || !result.Saved || result.MaskedToken != "cand…cret" { + t.Fatalf("result = %+v", result) + } + vault, _ := credentials.DefaultVault() + values, err := vault.Load() + if err != nil || values["openai"] != "candidate-secret" { + t.Fatalf("vault = %v, err=%v", values, err) + } +} + +func TestRunProviderConfigureFailedValidationPreservesToken(t *testing.T) { + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) + vault, _ := credentials.DefaultVault() + if err := vault.Set("anthropic", "existing-secret"); err != nil { + t.Fatalf("seed vault: %v", err) + } + previous := configureTokenModels + configureTokenModels = func(context.Context, ai.Backend, string) ([]ai.ModelDef, error) { + return nil, errors.New("credential rejected") + } + t.Cleanup(func() { configureTokenModels = previous }) + + _, err := runProviderTokenConfigure(context.Background(), ai.BackendAnthropic, ConfigureOptions{ + Token: text.SensitiveString("invalid-secret"), + }) + if err == nil { + t.Fatal("expected validation error") + } + values, loadErr := vault.Load() + if loadErr != nil || values["anthropic"] != "existing-secret" { + t.Fatalf("vault = %v, err=%v", values, loadErr) + } +} + +func TestRunProviderConfigureTestCurrentDoesNotWrite(t *testing.T) { + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) + t.Setenv("DEEPSEEK_API_KEY", "environment-secret") + previous := configureTokenModels + configureTokenModels = func(_ context.Context, backend ai.Backend, token string) ([]ai.ModelDef, error) { + if backend != ai.BackendDeepSeek || token != "environment-secret" { + t.Fatalf("validation got backend=%s token=%q", backend, token) + } + return []ai.ModelDef{{ID: "deepseek-example"}}, nil + } + t.Cleanup(func() { configureTokenModels = previous }) + + result, err := runProviderTokenConfigure(context.Background(), ai.BackendDeepSeek, ConfigureOptions{Test: true}) + if err != nil { + t.Fatalf("runProviderConfigure: %v", err) + } + if result.Saved || result.Source != credentials.SourceEnvironment { + t.Fatalf("result = %+v", result) + } + vault, _ := credentials.DefaultVault() + values, err := vault.Load() + if err != nil || len(values) != 0 { + t.Fatalf("test-only run wrote vault: %v, err=%v", values, err) + } +} + +func TestRunProviderDefaultsConfigureSavesPartialDefaultsAndActiveProvider(t *testing.T) { + captainconfig.SetPathForTesting(filepath.Join(t.TempDir(), ".captain.yaml")) + t.Cleanup(func() { captainconfig.SetPathForTesting("") }) + if err := captainconfig.Save(captainconfig.Config{AI: captainconfig.AIDefaults{ + DefaultProvider: "anthropic", + Providers: map[string]captainconfig.ProviderDefaults{ + "anthropic": {Agent: "claude-agent", Model: "claude-existing", ReasoningEffort: "low"}, + "openai": {Agent: "openai", Model: "gpt-existing", ReasoningEffort: "medium"}, + }, + }}); err != nil { + t.Fatalf("seed config: %v", err) + } + previous := configureDefaultsModels + configureDefaultsModels = func(_ context.Context, agent api.Backend) ([]ai.ModelDef, error) { + if agent != api.BackendCodexAgent { + t.Fatalf("agent = %s", agent) + } + return []ai.ModelDef{{ID: "gpt-5.6-sol"}}, nil + } + t.Cleanup(func() { configureDefaultsModels = previous }) + + result, err := runProviderDefaultsConfigure(context.Background(), api.OpenAIProvider, ConfigureOptions{ + Agent: "codex-agent", Active: true, + }) + if err != nil { + t.Fatalf("runProviderDefaultsConfigure: %v", err) + } + if result.Agent != "codex-agent" || result.Model != "gpt-5.6-sol" || !result.Active { + t.Fatalf("result = %+v", result) + } + got, _, err := captainconfig.Load() + if err != nil { + t.Fatalf("load config: %v", err) + } + if got.AI.DefaultProvider != "openai" || got.AI.Providers["openai"].Agent != "codex-agent" || got.AI.Providers["anthropic"].Model != "claude-existing" { + t.Fatalf("saved config = %+v", got.AI) + } +} + +func TestRunProviderDefaultsConfigureDefaultEffortClearsSavedEffort(t *testing.T) { + captainconfig.SetPathForTesting(filepath.Join(t.TempDir(), ".captain.yaml")) + t.Cleanup(func() { captainconfig.SetPathForTesting("") }) + if err := captainconfig.Save(captainconfig.Config{AI: captainconfig.AIDefaults{ + Providers: map[string]captainconfig.ProviderDefaults{ + "openai": {Agent: "codex-agent", Model: "gpt-5.6-sol", ReasoningEffort: "high"}, + }, + }}); err != nil { + t.Fatalf("seed config: %v", err) + } + previous := configureDefaultsModels + configureDefaultsModels = func(context.Context, api.Backend) ([]ai.ModelDef, error) { + return []ai.ModelDef{{ID: "gpt-5.6-sol"}}, nil + } + t.Cleanup(func() { configureDefaultsModels = previous }) + + result, err := runProviderDefaultsConfigure(context.Background(), api.OpenAIProvider, ConfigureOptions{Effort: "default"}) + if err != nil { + t.Fatalf("runProviderDefaultsConfigure: %v", err) + } + got, _, loadErr := captainconfig.Load() + if loadErr != nil || result.Effort != "" || got.AI.Providers["openai"].ReasoningEffort != "" { + t.Fatalf("result=%+v config=%+v err=%v", result, got.AI, loadErr) + } +} + +func TestRunProviderConfigureRejectsCredentialAndDefaultFlagsTogether(t *testing.T) { + _, err := runProviderConfigure(context.Background(), ConfigureOptions{ + Provider: "openai", Model: "gpt-5.6", Token: text.SensitiveString("candidate-secret"), + }) + if err == nil || !strings.Contains(err.Error(), "cannot be combined") { + t.Fatalf("error = %v", err) + } +} + func TestBuildConfigFromForm_TogglesInvert(t *testing.T) { in := formInputs{ Backend: "anthropic", @@ -23,19 +192,20 @@ func TestBuildConfigFromForm_TogglesInvert(t *testing.T) { got := buildConfigFromForm(in) want := captainconfig.Config{ AI: captainconfig.AIDefaults{ - Backend: "anthropic", - Model: "claude-sonnet-4-6", - ReasoningEffort: "medium", - BudgetUSD: 1.5, - MaxTokens: 8192, - Timeout: "180s", - NoCache: false, - NoMCP: false, - NoHooks: false, - NoSkills: true, - NoUser: true, - NoProject: true, - NoMemory: true, + DefaultProvider: "anthropic", + Providers: map[string]captainconfig.ProviderDefaults{ + "anthropic": {Agent: "anthropic", Model: "claude-sonnet-4-6", ReasoningEffort: "medium"}, + }, + BudgetUSD: 1.5, + MaxTokens: 8192, + Timeout: "180s", + NoCache: false, + NoMCP: false, + NoHooks: false, + NoSkills: true, + NoUser: true, + NoProject: true, + NoMemory: true, }, } if !reflect.DeepEqual(got, want) { @@ -104,6 +274,8 @@ func TestModelOptionsFor_NoKeyShowsErrorRow(t *testing.T) { // sentinel option carrying the error so the user can fix their environment // without leaving the wizard. CLI/agent backends are covered separately — // they list from the catalog and never require a key. + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) t.Setenv("OPENAI_API_KEY", "") t.Setenv("ANTHROPIC_API_KEY", "") t.Setenv("GEMINI_API_KEY", "") @@ -134,20 +306,20 @@ func TestModelOptionsFor_CLIBackendsUseCatalogWithoutKey(t *testing.T) { if len(opts) != 1 { t.Fatalf("codex-cli picker = %+v, want a single catalog option", opts) } - // huh.Option.Key is the display label; the selectable value is the runtime - // slug the codex provider expects verbatim. - if opts[0].Value != "gpt-5-codex" { - t.Errorf("codex-cli option value = %q, want catalog slug gpt-5-codex", opts[0].Value) + if opts[0].Value != "gpt-5.5" { + t.Errorf("codex-cli option value = %q, want exact model gpt-5.5", opts[0].Value) } } func TestDefaultModelFor_HardcodedPerBackend(t *testing.T) { cases := map[ai.Backend]string{ ai.BackendAnthropic: "claude-sonnet-5", - ai.BackendClaudeCLI: "claude-agent-sonnet", - ai.BackendClaudeAgent: "claude-agent-sonnet", - ai.BackendOpenAI: "gpt-5.5", - ai.BackendCodexCLI: "gpt-5-codex", + ai.BackendClaudeCLI: "claude-sonnet-5", + ai.BackendClaudeAgent: "claude-sonnet-5", + ai.BackendOpenAI: "gpt-5.6", + ai.BackendCodexCLI: "gpt-5.6-sol", + ai.BackendCodexAgent: "gpt-5.6-sol", + ai.BackendCodexCmux: "gpt-5.6-sol", ai.BackendGemini: "gemini-3.5-flash", ai.BackendGeminiCLI: "gemini-3.5-flash", } @@ -158,6 +330,27 @@ func TestDefaultModelFor_HardcodedPerBackend(t *testing.T) { } } +func TestEffortHuhOptionsForModel(t *testing.T) { + luna := effortHuhOptionsFor(ai.BackendCodexAgent, "luna") + for _, option := range luna { + if option.Value == "ultra" { + t.Fatalf("Luna options should not include ultra: %+v", luna) + } + } + sol := effortHuhOptionsFor(ai.BackendCodexAgent, "sol") + foundMax := false + for _, option := range sol { + foundMax = foundMax || option.Value == "max" + } + if !foundMax { + t.Fatalf("Sol options should include max: %+v", sol) + } + noEffort := effortHuhOptionsFor(ai.BackendDeepSeek, "deepseek-v4-pro") + if len(noEffort) != 1 || noEffort[0].Value != "" { + t.Fatalf("DeepSeek options = %+v, want only the backend default", noEffort) + } +} + func TestValidators(t *testing.T) { if err := validateFloat(""); err != nil { t.Errorf("validateFloat(\"\") = %v, want nil (blank allowed)", err) diff --git a/pkg/cli/database.go b/pkg/cli/database.go new file mode 100644 index 0000000..13a2a6c --- /dev/null +++ b/pkg/cli/database.go @@ -0,0 +1,207 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + "sync" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/monitor" + commonsdb "github.com/flanksource/commons-db/db" + "github.com/spf13/pflag" + "gorm.io/gorm" +) + +const databaseURLFlag = "db-url" + +var databaseURL string + +// BindDatabaseURLFlag exposes the process database as a root persistent flag. +// The explicit CLI value wins over environment variables and config files. +func BindDatabaseURLFlag(flags *pflag.FlagSet) { + flags.StringVar(&databaseURL, databaseURLFlag, "", "PostgreSQL database URL (overrides environment and db.json)") +} + +// captainDBState memoizes the process-wide native database handle. The +// database is mandatory: session/plan/prompt surfaces read it exclusively, so +// failing to open it is a loud error rather than a degraded mode. +var captainDBState struct { + mu sync.Mutex + opened bool + migrated bool + db *database.DB + dsn string + source string + err error +} + +type captainDatabaseMode uint8 + +const ( + captainDatabaseNoMigrations captainDatabaseMode = iota + captainDatabaseWithMigrations +) + +// captainDB opens the native Captain database without running migrations. +func captainDB(ctx context.Context) (*database.DB, error) { + return captainDBForMode(ctx, captainDatabaseNoMigrations) +} + +func captainServeDB(ctx context.Context) (*database.DB, error) { + return captainDBForMode(ctx, captainDatabaseWithMigrations) +} + +func captainDBForMode(ctx context.Context, mode captainDatabaseMode) (*database.DB, error) { + captainDBState.mu.Lock() + defer captainDBState.mu.Unlock() + if captainDBState.opened { + if mode == captainDatabaseWithMigrations && !captainDBState.migrated { + return nil, errors.New("captain serve cannot migrate after the process database was opened without migrations") + } + return captainDBState.db, captainDBState.err + } + captainDBState.db, captainDBState.dsn, captainDBState.source, captainDBState.err = openCaptainDB(ctx, mode) + captainDBState.opened = true + captainDBState.migrated = captainDBState.err == nil && mode == captainDatabaseWithMigrations + return captainDBState.db, captainDBState.err +} + +// ConfigureNativeDatabase injects a host-owned GORM pool before Captain's CLI +// database is first used. Hosts such as Gavel use this to keep Captain session, +// prompt, and plan APIs on the same process-owned database. Reconfiguring the +// same pool is idempotent; replacing an initialized pool is rejected because +// callers may already hold handles backed by it. +func ConfigureNativeDatabase(gormDB *gorm.DB) error { + db, err := database.Use(gormDB) + if err != nil { + return err + } + + captainDBState.mu.Lock() + defer captainDBState.mu.Unlock() + if captainDBState.opened { + if captainDBState.err == nil && captainDBState.db != nil && captainDBState.db.Gorm() == gormDB { + return nil + } + return fmt.Errorf("native Captain database is already configured with a different pool") + } + captainDBState.db = db + captainDBState.dsn = "" + captainDBState.source = "host-provided database" + captainDBState.err = nil + captainDBState.opened = true + captainDBState.migrated = true + return nil +} + +// setCaptainDBForTest injects (or, with nil, resets) the process-wide handle +// so tests run against their own embedded database instead of a configured +// DSN. Production code never calls this. +func setCaptainDBForTest(db *database.DB) { + captainDBState.mu.Lock() + defer captainDBState.mu.Unlock() + captainDBState.db = db + captainDBState.dsn = "" + captainDBState.source = "" + captainDBState.err = nil + captainDBState.opened = db != nil + captainDBState.migrated = db != nil +} + +func openCaptainDB(ctx context.Context, mode captainDatabaseMode) (*database.DB, string, string, error) { + dsn, source, err := captainDSN() + if err != nil { + return nil, "", "", err + } + log.Debugf("captain database using %s", source) + options := []database.Option{database.WithDSN(dsn)} + if mode == captainDatabaseWithMigrations { + options = append(options, database.WithMigrations()) + } + db, err := database.Open(ctx, options...) + if err != nil { + return nil, "", "", fmt.Errorf("open captain database (%s): %w", source, err) + } + return db, dsn, source, nil +} + +func captainDatabaseIdentity() (dsn, source string) { + captainDBState.mu.Lock() + defer captainDBState.mu.Unlock() + return captainDBState.dsn, captainDBState.source +} + +// serveMonitorState holds the serve process's live monitor so prompt-run code +// can register freshly launched transcripts for immediate tailing. +var serveMonitorState struct { + mu sync.RWMutex + mon *monitor.Monitor +} + +func setServeMonitor(mon *monitor.Monitor) { + serveMonitorState.mu.Lock() + serveMonitorState.mon = mon + serveMonitorState.mu.Unlock() +} + +func serveMonitor() *monitor.Monitor { + serveMonitorState.mu.RLock() + defer serveMonitorState.mu.RUnlock() + return serveMonitorState.mon +} + +func captainHostID() string { + return database.LocalHostID() +} + +// monitorDiscoverProcesses is indirected so cli tests can fake live-process +// discovery; nil selects the monitor's real ps-based discovery. +var monitorDiscoverProcesses func() ([]monitor.Process, error) + +// freshenSessionDB runs a one-shot monitor pass (ps poll + incremental +// transcript scan) before a CLI read when no live monitor holds the writer +// lock. With serve running it is a fast no-op. +func freshenSessionDB(ctx context.Context) (*database.DB, error) { + db, err := captainDB(ctx) + if err != nil { + return nil, err + } + config := monitor.Config{DB: db, HostID: captainHostID(), DiscoverProcesses: monitorDiscoverProcesses} + if err := monitor.RunOnce(ctx, config); err != nil { + return nil, fmt.Errorf("refresh session database: %w", err) + } + return db, nil +} + +// captainDSN resolves the database connection: explicit env DSNs first, then a +// gavel-shared database, finally captain's own shared embedded postgres. +func captainDSN() (dsn, source string, err error) { + if dsn := strings.TrimSpace(databaseURL); dsn != "" { + return dsn, "--" + databaseURLFlag, nil + } + for _, env := range []string{gavelDBEnvDSN, gavelCacheEnvDSN, captainSessionEnvDSN} { + if dsn := strings.TrimSpace(os.Getenv(env)); dsn != "" { + return dsn, env, nil + } + } + dsn, source, err = gavelConfiguredSessionDSN() + if err != nil { + return "", "", fmt.Errorf("resolve gavel database: %w", err) + } + if dsn != "" { + return dsn, source, nil + } + dir, err := sessionDBDir() + if err != nil { + return "", "", fmt.Errorf("resolve captain database directory: %w", err) + } + // Shared embedded-postgres daemon: leave running for other captain processes. + dsn, _, err = commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{DataDir: dir}) + if err != nil { + return "", "", fmt.Errorf("start captain embedded database: %w", err) + } + return dsn, "captain embedded database", nil +} diff --git a/pkg/cli/database_ginkgo_test.go b/pkg/cli/database_ginkgo_test.go new file mode 100644 index 0000000..8400e2c --- /dev/null +++ b/pkg/cli/database_ginkgo_test.go @@ -0,0 +1,31 @@ +package cli + +import ( + "os" + "path/filepath" + + "github.com/flanksource/captain/pkg/database" + commonsdb "github.com/flanksource/commons-db/db" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func withGinkgoCaptainDB() { + if os.Getenv("CAPTAIN_DB_EMBEDDED_TEST") == "" { + Skip("set CAPTAIN_DB_EMBEDDED_TEST=1 to run embedded-postgres cli tests") + } + dsn, stop, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: filepath.Join(GinkgoT().TempDir(), "postgres"), + Database: "captain_cli", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { Expect(stop()).To(Succeed()) }) + + db, err := database.Open(GinkgoT().Context(), database.WithDSN(dsn), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + setCaptainDBForTest(db) + DeferCleanup(func() { + setCaptainDBForTest(nil) + Expect(db.Close()).To(Succeed()) + }) +} diff --git a/pkg/cli/database_mode_ginkgo_test.go b/pkg/cli/database_mode_ginkgo_test.go new file mode 100644 index 0000000..eb09717 --- /dev/null +++ b/pkg/cli/database_mode_ginkgo_test.go @@ -0,0 +1,40 @@ +package cli + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/flanksource/captain/pkg/database" +) + +var _ = Describe("Captain database migration mode", Serial, func() { + AfterEach(func() { + setCaptainDBForTest(nil) + }) + + It("rejects serve startup after a non-migrating handle was installed", func(ctx SpecContext) { + db, err := database.Use(&gorm.DB{}) + Expect(err).NotTo(HaveOccurred()) + captainDBState.mu.Lock() + captainDBState.db = db + captainDBState.opened = true + captainDBState.migrated = false + captainDBState.mu.Unlock() + + _, err = captainServeDB(ctx) + + Expect(err).To(MatchError("captain serve cannot migrate after the process database was opened without migrations")) + }) + + It("reuses a migration-initialized handle", func(ctx SpecContext) { + db, err := database.Use(&gorm.DB{}) + Expect(err).NotTo(HaveOccurred()) + setCaptainDBForTest(db) + + opened, err := captainServeDB(ctx) + + Expect(err).NotTo(HaveOccurred()) + Expect(opened).To(BeIdenticalTo(db)) + }) +}) diff --git a/pkg/cli/database_test.go b/pkg/cli/database_test.go new file mode 100644 index 0000000..81b0a12 --- /dev/null +++ b/pkg/cli/database_test.go @@ -0,0 +1,143 @@ +package cli + +import ( + "os" + "path/filepath" + "testing" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/monitor" + commonsdb "github.com/flanksource/commons-db/db" + "github.com/spf13/pflag" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestConfigureNativeDatabase(t *testing.T) { + setCaptainDBForTest(nil) + t.Cleanup(func() { setCaptainDBForTest(nil) }) + + first := &gorm.DB{} + require.NoError(t, ConfigureNativeDatabase(first)) + require.NoError(t, ConfigureNativeDatabase(first), "reconfiguring the same pool should be idempotent") + + db, err := captainDB(t.Context()) + require.NoError(t, err) + require.Same(t, first, db.Gorm()) + + err = ConfigureNativeDatabase(&gorm.DB{}) + require.EqualError(t, err, "native Captain database is already configured with a different pool") +} + +func TestConfigureNativeDatabaseRejectsNil(t *testing.T) { + setCaptainDBForTest(nil) + t.Cleanup(func() { setCaptainDBForTest(nil) }) + + err := ConfigureNativeDatabase(nil) + require.EqualError(t, err, "captain database GORM pool is nil") +} + +func TestCaptainDSNPrecedence(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + databaseURL = "" + t.Cleanup(func() { databaseURL = "" }) + + t.Run("db-url flag wins", func(t *testing.T) { + flags := pflag.NewFlagSet("captain", pflag.ContinueOnError) + BindDatabaseURLFlag(flags) + require.NoError(t, flags.Parse([]string{"--db-url", "postgres://flag/captain"})) + t.Cleanup(func() { databaseURL = "" }) + + t.Setenv(gavelDBEnvDSN, "postgres://primary/gavel") + t.Setenv(gavelCacheEnvDSN, "postgres://cache/gavel") + t.Setenv(captainSessionEnvDSN, "postgres://captain/db") + dsn, source, err := captainDSN() + require.NoError(t, err) + require.Equal(t, "postgres://flag/captain", dsn) + require.Equal(t, "--db-url", source) + }) + + t.Run("gavel primary env wins", func(t *testing.T) { + t.Setenv(gavelDBEnvDSN, "postgres://primary/gavel") + t.Setenv(gavelCacheEnvDSN, "postgres://cache/gavel") + t.Setenv(captainSessionEnvDSN, "postgres://captain/db") + dsn, source, err := captainDSN() + require.NoError(t, err) + require.Equal(t, "postgres://primary/gavel", dsn) + require.Equal(t, gavelDBEnvDSN, source) + }) + + t.Run("cache env is next", func(t *testing.T) { + t.Setenv(gavelDBEnvDSN, "") + t.Setenv(gavelCacheEnvDSN, "postgres://cache/gavel") + t.Setenv(captainSessionEnvDSN, "postgres://captain/db") + dsn, source, err := captainDSN() + require.NoError(t, err) + require.Equal(t, "postgres://cache/gavel", dsn) + require.Equal(t, gavelCacheEnvDSN, source) + }) + + t.Run("captain env is next", func(t *testing.T) { + t.Setenv(gavelDBEnvDSN, "") + t.Setenv(gavelCacheEnvDSN, "") + t.Setenv(captainSessionEnvDSN, "postgres://captain/db") + dsn, source, err := captainDSN() + require.NoError(t, err) + require.Equal(t, "postgres://captain/db", dsn) + require.Equal(t, captainSessionEnvDSN, source) + }) + + t.Run("gavel db.json mode=dsn is used without env", func(t *testing.T) { + t.Setenv(gavelDBEnvDSN, "") + t.Setenv(gavelCacheEnvDSN, "") + t.Setenv(captainSessionEnvDSN, "") + dir := filepath.Join(home, ".config", "gavel") + require.NoError(t, os.MkdirAll(dir, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "db.json"), + []byte(`{"mode":"dsn","dsn":"postgres://from/config"}`), 0o644)) + dsn, source, err := captainDSN() + require.NoError(t, err) + require.Equal(t, "postgres://from/config", dsn) + require.Contains(t, source, "db.json") + }) + + t.Run("invalid db.json mode fails loudly", func(t *testing.T) { + t.Setenv(gavelDBEnvDSN, "") + t.Setenv(gavelCacheEnvDSN, "") + t.Setenv(captainSessionEnvDSN, "") + dir := filepath.Join(home, ".config", "gavel") + require.NoError(t, os.MkdirAll(dir, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "db.json"), + []byte(`{"mode":"bogus"}`), 0o644)) + _, _, err := captainDSN() + require.Error(t, err) + }) +} + +// withTestCaptainDB starts an isolated embedded postgres, injects it as the +// process-wide captain database, and fakes live-process discovery so tests +// never touch a configured DSN or the host's real processes. +func withTestCaptainDB(t *testing.T, processes ...monitor.Process) *database.DB { + t.Helper() + if os.Getenv("CAPTAIN_DB_EMBEDDED_TEST") == "" { + t.Skip("set CAPTAIN_DB_EMBEDDED_TEST=1 to run embedded-postgres cli tests") + } + dsn, stop, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: filepath.Join(t.TempDir(), "postgres"), + Database: "captain_cli", + }) + require.NoError(t, err) + db, err := database.Open(t.Context(), database.WithDSN(dsn), database.WithMigrations()) + require.NoError(t, err) + + setCaptainDBForTest(db) + monitorDiscoverProcesses = func() ([]monitor.Process, error) { return processes, nil } + t.Cleanup(func() { + setCaptainDBForTest(nil) + monitorDiscoverProcesses = nil + require.NoError(t, db.Close()) + require.NoError(t, stop()) + }) + return db +} diff --git a/pkg/cli/event_renderer.go b/pkg/cli/event_renderer.go new file mode 100644 index 0000000..a4b2120 --- /dev/null +++ b/pkg/cli/event_renderer.go @@ -0,0 +1,17 @@ +package cli + +import ( + "os" + + "github.com/flanksource/captain/pkg/ai" +) + +// NewEventRenderer returns the canonical stateful terminal callback used for +// live Captain events. It shares the same history-backed row renderer as the +// Captain CLI, including session boundaries and structured tool rows. +func NewEventRenderer(output *os.File) func(int, ai.Event) { + renderer := newLineRenderer(output, 8) + return func(_ int, event ai.Event) { + renderEvent(output, renderer, event) + } +} diff --git a/pkg/cli/event_renderer_ginkgo_test.go b/pkg/cli/event_renderer_ginkgo_test.go new file mode 100644 index 0000000..0a14938 --- /dev/null +++ b/pkg/cli/event_renderer_ginkgo_test.go @@ -0,0 +1,50 @@ +package cli + +import ( + "io" + "os" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/claude" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Captain event renderer", func() { + It("keeps text deltas contiguous and renders a command once", func() { + reader, writer, err := os.Pipe() + Expect(err).NotTo(HaveOccurred()) + + render := NewEventRenderer(writer) + for _, delta := range []string{"a", " keyed", " H", "MAC", " so", " the", " token"} { + render(0, ai.Event{Kind: ai.EventText, Text: delta, Model: "gpt-5.6-sol"}) + } + render(0, ai.Event{ + Kind: ai.EventToolUse, + Tool: "Bash", + Input: map[string]any{"command": "pwd"}, + ToolCallID: "cmd-1", + SessionID: "thread-1", + Model: "gpt-5.6-sol", + Raw: claude.ToolUse{ + Tool: "Bash", + Input: map[string]any{"command": "pwd"}, + ToolUseID: "cmd-1", + SessionID: "thread-1", + Source: "codex", + Model: "gpt-5.6-sol", + }, + }) + + Expect(writer.Close()).To(Succeed()) + output, err := io.ReadAll(reader) + Expect(err).NotTo(HaveOccurred()) + Expect(reader.Close()).To(Succeed()) + + text := string(output) + Expect(text).To(ContainSubstring("a keyed HMAC so the token")) + Expect(strings.Count(text, "pwd")).To(Equal(1)) + Expect(text).NotTo(ContainSubstring("[gpt-5.6-sol]")) + }) +}) diff --git a/pkg/cli/gavel_dsn.go b/pkg/cli/gavel_dsn.go new file mode 100644 index 0000000..1bd44d4 --- /dev/null +++ b/pkg/cli/gavel_dsn.go @@ -0,0 +1,188 @@ +package cli + +import ( + "encoding/json" + "errors" + "fmt" + "net" + "os" + "path/filepath" + "strconv" + "strings" + "syscall" + "time" + + commonsdb "github.com/flanksource/commons-db/db" +) + +const ( + captainSessionEnvDSN = "CAPTAIN_SESSION_DB_URL" + gavelDBEnvDSN = "GAVEL_DB_DSN" + gavelCacheEnvDSN = "GAVEL_GITHUB_CACHE_DSN" + + gavelDBModeDSN = "dsn" + gavelDBModeEmbedded = "embedded" +) + +// sessionDBDir is the embedded-postgres data directory (shared across processes). +func sessionDBDir() (string, error) { + cache, err := os.UserCacheDir() + if err != nil { + return "", err + } + return filepath.Join(cache, "captain", "session-db"), nil +} + +type gavelDBConfig struct { + Mode string `json:"mode"` + DSN string `json:"dsn,omitempty"` +} + +// gavelConfiguredSessionDSN resolves a gavel-shared database from +// ~/.config/gavel/db.json: an explicit DSN, or gavel's embedded postgres +// (reusing a running instance before starting one). +func gavelConfiguredSessionDSN() (string, string, error) { + cfg, path, err := loadGavelDBConfig() + if err != nil { + return "", "", err + } + switch cfg.Mode { + case "": + return "", "", nil + case gavelDBModeDSN: + if strings.TrimSpace(cfg.DSN) == "" { + return "", "", fmt.Errorf("%s has mode=%s but empty dsn", path, gavelDBModeDSN) + } + return cfg.DSN, path, nil + case gavelDBModeEmbedded: + running, err := findRunningGavelEmbeddedPostgres() + if err != nil { + return "", "", err + } + if running != nil { + return gavelEmbeddedDSN(running.Port), path, nil + } + dataDir, err := gavelEmbeddedDataDir() + if err != nil { + return "", "", err + } + dsn, _, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: dataDir, + Database: "gavel", + }) + if err != nil { + return "", "", err + } + return dsn, path, nil + default: + return "", "", fmt.Errorf("%s has unsupported mode %q", path, cfg.Mode) + } +} + +func loadGavelDBConfig() (gavelDBConfig, string, error) { + path, err := gavelDBConfigPath() + if err != nil { + return gavelDBConfig{}, "", err + } + b, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return gavelDBConfig{}, path, nil + } + return gavelDBConfig{}, path, fmt.Errorf("read %s: %w", path, err) + } + var cfg gavelDBConfig + if err := json.Unmarshal(b, &cfg); err != nil { + return gavelDBConfig{}, path, fmt.Errorf("parse %s: %w", path, err) + } + return cfg, path, nil +} + +func gavelDBConfigPath() (string, error) { + dir, err := gavelStateDir() + if err != nil { + return "", err + } + return filepath.Join(dir, "db.json"), nil +} + +func gavelEmbeddedDataDir() (string, error) { + dir, err := gavelStateDir() + if err != nil { + return "", err + } + return filepath.Join(dir, "embedded-pg"), nil +} + +func gavelStateDir() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("resolve home dir: %w", err) + } + dir := filepath.Join(home, ".config", "gavel") + if err := os.MkdirAll(dir, 0o755); err != nil { + return "", fmt.Errorf("create gavel state dir %s: %w", dir, err) + } + return dir, nil +} + +type runningGavelEmbeddedPostgres struct { + PID int + Port int +} + +func findRunningGavelEmbeddedPostgres() (*runningGavelEmbeddedPostgres, error) { + dataDir, err := gavelEmbeddedDataDir() + if err != nil { + return nil, err + } + pidPath := filepath.Join(dataDir, "data", "postmaster.pid") + raw, err := os.ReadFile(pidPath) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + return nil, fmt.Errorf("read %s: %w", pidPath, err) + } + lines := strings.Split(string(raw), "\n") + const postmasterLinePort = 3 + if len(lines) <= postmasterLinePort { + return nil, fmt.Errorf("%s has %d lines, need >%d", pidPath, len(lines), postmasterLinePort) + } + pid, err := strconv.Atoi(strings.TrimSpace(lines[0])) + if err != nil || pid <= 0 { + return nil, fmt.Errorf("%s: invalid pid %q: %w", pidPath, lines[0], err) + } + port, err := strconv.Atoi(strings.TrimSpace(lines[postmasterLinePort])) + if err != nil || port <= 0 || port > 65535 { + return nil, fmt.Errorf("%s: invalid port %q: %w", pidPath, lines[postmasterLinePort], err) + } + if !processAlive(pid) || !tcpPortReachable("localhost", port) { + return nil, nil + } + return &runningGavelEmbeddedPostgres{PID: pid, Port: port}, nil +} + +func gavelEmbeddedDSN(port int) string { + return fmt.Sprintf("postgres://postgres:postgres@localhost:%d/gavel?sslmode=disable", port) +} + +func processAlive(pid int) bool { + if pid <= 0 { + return false + } + p, err := os.FindProcess(pid) + if err != nil { + return false + } + return p.Signal(syscall.Signal(0)) == nil +} + +func tcpPortReachable(host string, port int) bool { + conn, err := net.DialTimeout("tcp", net.JoinHostPort(host, strconv.Itoa(port)), 200*time.Millisecond) + if err != nil { + return false + } + _ = conn.Close() + return true +} diff --git a/pkg/cli/history.go b/pkg/cli/history.go index 90d0ae2..ead5c1e 100644 --- a/pkg/cli/history.go +++ b/pkg/cli/history.go @@ -91,7 +91,7 @@ func RunHistory(opts HistoryOptions) (any, error) { IncludeAgents: opts.Agents, } - if len(opts.Categories) == 0 && opts.TextFilter == "" && !opts.Last { + if len(opts.Categories) == 0 && opts.TextFilter == "" && !opts.Last && !usesDefaultHiddenHistoryTools(opts) { filter.Limit = opts.Limit } @@ -143,6 +143,8 @@ func RunHistory(opts HistoryOptions) (any, error) { tl = filterToolsByPath(tl, newPathFilter(opts.Plans, opts.Ignored)) } + tl = collapseRepeatedTitles(tl) + if opts.Last { tl = lastSessionTools(tl) // --last means "the whole most-recent session" — don't let the row @@ -181,6 +183,31 @@ func lastSessionTools(tl []tools.Tool) []tools.Tool { return tl[start:] } +// collapseRepeatedTitles drops SessionTitle rows that repeat the title already +// shown for their session. Claude Code rewrites the ai-title record on nearly +// every turn, so a long session otherwise renders dozens of identical +// "🏷 title …" rows (65 of them, for 3 distinct titles, on a real session). A +// title that genuinely changes still gets its own row — including a change back +// to a title used earlier in the session. +func collapseRepeatedTitles(tl []tools.Tool) []tools.Tool { + shown := map[sessionKey]string{} + out := make([]tools.Tool, 0, len(tl)) + for _, t := range tl { + title, ok := t.(*tools.SessionTitleTool) + if !ok { + out = append(out, t) + continue + } + key := keyForTool(t) + if prev, seen := shown[key]; seen && prev == title.Str("aiTitle") { + continue + } + shown[key] = title.Str("aiTitle") + out = append(out, t) + } + return out +} + // gatherToolUses collects Claude and Codex tool uses for the working directory, // tags each with its source, and applies the filter (including any session ID) // to both. It is the shared front-end for the history and changes commands. @@ -568,6 +595,9 @@ func filterTools(tl []tools.Tool, opts HistoryOptions, classifier *bash.Category cat := classifyTool(t, classifier) categorySet[cat] = struct{}{} + if hideDefaultHistoryTool(t, opts) { + continue + } if len(opts.Categories) > 0 && !matchCategoryFilters(categoryFilterCandidates(t, cat), opts.Categories) { continue } @@ -590,6 +620,9 @@ func filterTools(tl []tools.Tool, opts HistoryOptions, classifier *bash.Category categories = append(categories, c) } for _, filter := range opts.Categories { + if strings.HasPrefix(strings.TrimSpace(filter), "!") { + continue + } if similar := captainCollections.FindSimilar(filter, categories, 3); len(similar) > 0 { fmt.Fprintf(os.Stderr, "category %q matched nothing. Did you mean: %s?\n", filter, strings.Join(similar, ", ")) } @@ -599,6 +632,48 @@ func filterTools(tl []tools.Tool, opts HistoryOptions, classifier *bash.Category return result } +func hideDefaultHistoryTool(t tools.Tool, opts HistoryOptions) bool { + base := t.Base() + switch { + case t.Name() == "TokenCount", base.RawTool == "TokenCount": + return usesDefaultHiddenHistoryTools(opts) + default: + return false + } +} + +func usesDefaultHiddenHistoryTools(opts HistoryOptions) bool { + return !hasExplicitToolFilter(opts.Tools, "TokenCount") && toolFiltersMayInclude(opts.Tools, "TokenCount") +} + +func hasExplicitToolFilter(filters []string, tool string) bool { + for _, filter := range filters { + for _, part := range strings.Split(filter, ",") { + part = strings.TrimSpace(part) + if part == "" || strings.HasPrefix(part, "!") { + continue + } + if collections.MatchItems(tool, part) { + return true + } + } + } + return false +} + +func toolFiltersMayInclude(filters []string, tool string) bool { + var parts []string + for _, filter := range filters { + for _, part := range strings.Split(filter, ",") { + part = strings.TrimSpace(part) + if part != "" { + parts = append(parts, part) + } + } + } + return collections.MatchItems(tool, parts...) +} + func matchCategoryFilters(candidates []string, filters []string) bool { filters = normalizeCategoryFilters(filters) if len(filters) == 0 { @@ -647,29 +722,40 @@ func matchesAnyCategoryCandidate(candidates []string, pattern string) bool { func categoryFilterCandidates(t tools.Tool, category string) []string { base := t.Base() - return uniqueNonEmpty( + values := []string{ category, t.Name(), base.RawTool, - messageAlias(t.Name()), - ) + } + values = append(values, chatCategoryAliases(t.Name())...) + values = append(values, chatCategoryAliases(base.RawTool)...) + return uniqueNonEmpty(values...) } func toolUseCategoryFilterCandidates(tu claude.ToolUse, category string) []string { - return uniqueNonEmpty( + values := []string{ category, tu.Tool, tu.DisplayTool(), - messageAlias(tu.DisplayTool()), - ) + } + values = append(values, chatCategoryAliases(tu.Tool)...) + values = append(values, chatCategoryAliases(tu.DisplayTool())...) + return uniqueNonEmpty(values...) } -func messageAlias(tool string) string { +func chatCategoryAliases(tool string) []string { + if tools.IsEventToolName(tool) { + return []string{"chat", "event"} + } switch strings.ToLower(tool) { - case "assistant", "reasoning": - return "message" + case "system", "assistant", "reasoning": + return []string{"chat", "message"} + case "user": + return []string{"chat", "message"} + case "event": + return []string{"chat", "event"} default: - return "" + return nil } } @@ -730,14 +816,30 @@ func matchesHistoryTextFilter(t tools.Tool, category, filter string) bool { return false } +// shellCommand returns the raw command line for tools that execute a shell. +// Claude's Bash carries it as "command"; Codex's exec carries it as "input" +// (matching AnalyzeToolUse). Without the exec arm every Codex shell call +// classifies as `other` regardless of what it ran. +func shellCommand(base *tools.BaseTool) (string, bool) { + switch base.RawTool { + case "Bash": + cmd, ok := base.Input["command"].(string) + return cmd, ok + case "exec": + cmd, ok := base.Input["input"].(string) + return cmd, ok + } + return "", false +} + func classifyTool(t tools.Tool, classifier *bash.CategoryClassifier) string { if cat := t.Category(); cat != "" { return cat } base := t.Base() cat := classifier.ClassifyToolWithPath(base.RawTool, t.FilePath()) - if cat == bash.CategoryOther && base.RawTool == "Bash" { - if rawCmd, ok := base.Input["command"].(string); ok { + if cat == bash.CategoryOther { + if rawCmd, ok := shellCommand(base); ok { cat = classifier.ClassifyBash(rawCmd) } } @@ -747,7 +849,7 @@ func classifyTool(t tools.Tool, classifier *bash.CategoryClassifier) string { func approvedStatus(t tools.Tool) string { base := t.Base() name := t.Name() - if base.RawTool == "ExitPlanMode" || base.RawTool == "User" || name == "Plan" { + if base.RawTool == "ExitPlanMode" || isChatHistoryTool(base.RawTool) || name == "Plan" { return "" } if base.Denied { @@ -846,6 +948,9 @@ func runHistorySummary(toolUses []claude.ToolUse, opts HistoryOptions, classifie } func classifyToolUse(tu claude.ToolUse, classifier *bash.CategoryClassifier) string { + if isChatHistoryTool(tu.Tool) { + return "chat" + } category := classifier.ClassifyToolWithPath(tu.Tool, tu.FilePath()) if category == bash.CategoryOther && tu.Tool == "Bash" { if rawCmd, ok := tu.Input["command"].(string); ok { @@ -855,6 +960,18 @@ func classifyToolUse(tu claude.ToolUse, classifier *bash.CategoryClassifier) str return string(category) } +func isChatHistoryTool(tool string) bool { + if tools.IsEventToolName(tool) { + return true + } + switch tool { + case "System", "User", "Assistant", "Reasoning", "Event": + return true + default: + return false + } +} + func matchesToolUseTextFilter(tu claude.ToolUse, category, filter string) bool { filter = strings.ToLower(strings.TrimSpace(filter)) if filter == "" { diff --git a/pkg/cli/history_cost_test.go b/pkg/cli/history_cost_test.go index 88f4ade..ee1da13 100644 --- a/pkg/cli/history_cost_test.go +++ b/pkg/cli/history_cost_test.go @@ -32,7 +32,7 @@ func TestRunHistory_CostWithoutClaudeFlag(t *testing.T) { "cwd": project, "message": map[string]any{ "role": "assistant", - "model": "claude-opus-4", + "model": "claude-opus-4-5", "usage": map[string]any{"input_tokens": 1000, "output_tokens": 500}, "content": []any{map[string]any{ "type": "tool_use", diff --git a/pkg/cli/history_test.go b/pkg/cli/history_test.go index 7f1e458..35e38f7 100644 --- a/pkg/cli/history_test.go +++ b/pkg/cli/history_test.go @@ -1 +1,125 @@ package cli + +import ( + "testing" + + "github.com/flanksource/captain/pkg/claude/tools" +) + +func titleTool(sessionID, title string) tools.Tool { + return tools.NewTool(tools.BaseTool{ + RawTool: "SessionTitle", + Source: "claude", + SessionID: sessionID, + Input: map[string]any{"aiTitle": title}, + }) +} + +func bashTool(sessionID, command string) tools.Tool { + return tools.NewTool(tools.BaseTool{ + RawTool: "Bash", + Source: "claude", + SessionID: sessionID, + Input: map[string]any{"command": command}, + }) +} + +func collapsedTitles(t *testing.T, tl []tools.Tool) []string { + t.Helper() + var titles []string + for _, tool := range collapseRepeatedTitles(tl) { + if title, ok := tool.(*tools.SessionTitleTool); ok { + titles = append(titles, title.Str("aiTitle")) + } + } + return titles +} + +// TestCollapseRepeatedTitles_ObservedRunPattern reproduces the shape of session +// 929f3d1b: 82 ai-title records covering only 3 distinct titles, each rewritten +// once per turn and therefore interleaved with real tool calls. Only the first +// row of each run should survive. +func TestCollapseRepeatedTitles_ObservedRunPattern(t *testing.T) { + const sessionID = "929f3d1b" + runs := []struct { + title string + count int + }{ + {"Update GitHub scraper for OAuth and security settings", 23}, + {"github-scraper-org-settings", 57}, + {"pr-coderabbit-review", 2}, + } + + var tl []tools.Tool + for _, run := range runs { + for i := 0; i < run.count; i++ { + tl = append(tl, titleTool(sessionID, run.title), bashTool(sessionID, "go test ./...")) + } + } + + got := collapsedTitles(t, tl) + want := []string{runs[0].title, runs[1].title, runs[2].title} + if len(got) != len(want) { + t.Fatalf("collapsed to %d title rows (%v), want %d", len(got), got, len(want)) + } + for i := range want { + if got[i] != want[i] { + t.Errorf("title[%d] = %q, want %q", i, got[i], want[i]) + } + } + + if kept := len(collapseRepeatedTitles(tl)); kept != len(want)+82 { + t.Errorf("non-title rows must pass through untouched: kept %d rows, want %d", kept, len(want)+82) + } +} + +// TestCollapseRepeatedTitles_ChangesStillRender covers the cases the collapse +// must not swallow: a title reverting to an earlier value, and the same title +// used by a different session. +func TestCollapseRepeatedTitles_ChangesStillRender(t *testing.T) { + tests := []struct { + name string + in []tools.Tool + want []string + }{ + { + name: "revert to an earlier title gets its own row", + in: []tools.Tool{ + titleTool("S1", "first"), + titleTool("S1", "second"), + titleTool("S1", "second"), + titleTool("S1", "first"), + }, + want: []string{"first", "second", "first"}, + }, + { + name: "identical titles in different sessions both render", + in: []tools.Tool{ + titleTool("S1", "shared"), + titleTool("S1", "shared"), + titleTool("S2", "shared"), + titleTool("S2", "shared"), + }, + want: []string{"shared", "shared"}, + }, + { + name: "no titles at all is a passthrough", + in: []tools.Tool{bashTool("S1", "ls")}, + want: nil, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := collapsedTitles(t, tc.in) + if len(got) != len(tc.want) { + t.Fatalf("got %d titles (%v), want %d (%v)", len(got), got, len(tc.want), tc.want) + } + for i := range tc.want { + if got[i] != tc.want[i] { + t.Errorf("title[%d] = %q, want %q", i, got[i], tc.want[i]) + } + } + }) + } +} diff --git a/pkg/cli/hook_monitor.go b/pkg/cli/hook_monitor.go new file mode 100644 index 0000000..68ae182 --- /dev/null +++ b/pkg/cli/hook_monitor.go @@ -0,0 +1,139 @@ +package cli + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" + "github.com/flanksource/captain/pkg/codexconfig" + "github.com/flanksource/captain/pkg/monitor" +) + +type HookMonitorNotifyOptions struct { + Provider string // claude | codex + URL string // serve base URL override; default monitor.ServeBaseURL() +} + +// hookNotifyTimeout bounds hook delivery so an agent turn is never held up: +// deliver within it or drop the event (recon reconciles drops). +const hookNotifyTimeout = time.Second + +// RunHookMonitorNotify is the hook receiver both providers invoke: Claude Code +// pipes the payload on stdin, codex appends it as the final argv argument. It +// always succeeds with empty stdout — a hook failure must never block or slow +// an agent turn, and Claude injects UserPromptSubmit hook stdout as context. +func RunHookMonitorNotify(opts HookMonitorNotifyOptions, args []string) error { + ev, err := readHookNotifyEvent(opts.Provider, args) + if err != nil { + fmt.Fprintf(os.Stderr, "captain hook monitor notify: %v\n", err) + return nil + } + baseURL := opts.URL + if baseURL == "" { + baseURL = api.ServeBaseURL() + } + ctx, cancel := context.WithTimeout(context.Background(), hookNotifyTimeout) + defer cancel() + if err := monitor.PostHookEvent(ctx, baseURL, ev); err != nil { + // Degraded mode by design: captain serve is down or slow. The event is + // dropped; startup/daily recon and the stale reaper converge the DB. + fmt.Fprintf(os.Stderr, "captain hook monitor notify: %v\n", err) + } + return nil +} + +type HookMonitorInstallOptions struct { + Timeout int `flag:"timeout" help:"Hook timeout in seconds" default:"10"` + URL string `flag:"url" help:"Captain serve base URL baked into the hook command (for non-default ports)"` +} + +// monitorHookEvents are the Claude Code lifecycle events captain subscribes to +// for session monitoring: start/end bind and tear down the session, the middle +// three signal progress worth ingesting. +var monitorHookEvents = []claude.HookEventType{ + claude.HookEventSessionStart, + claude.HookEventUserPromptSubmit, + claude.HookEventStop, + claude.HookEventSubagentStop, + claude.HookEventSessionEnd, +} + +// RunHookMonitorInstall installs the session-monitoring hooks user-wide: +// Claude Code lifecycle hooks in ~/.claude/settings.json and codex's notify +// program in ~/.codex/config.toml. +func RunHookMonitorInstall(opts HookMonitorInstallOptions) (any, error) { + captainPath, err := os.Executable() + if err != nil { + captainPath = "captain" + } + urlSuffix := "" + if opts.URL != "" { + urlSuffix = " --url " + opts.URL + } + + target, err := ensureUserClaudeSettings() + if err != nil { + return nil, err + } + claudeCommand := fmt.Sprintf("%s hook monitor notify --provider claude%s", captainPath, urlSuffix) + var results []string + for _, event := range monitorHookEvents { + result, err := installHook(target, string(event), "", claudeCommand, "hook monitor notify --provider claude", opts.Timeout) + if err != nil { + return nil, err + } + results = append(results, result) + } + + codexArgv := []string{captainPath, "hook", "monitor", "notify", "--provider", "codex"} + if opts.URL != "" { + codexArgv = append(codexArgv, "--url", opts.URL) + } + codexResult, err := codexconfig.SetNotify(codexArgv) + if err != nil { + return nil, err + } + results = append(results, codexResult) + + results = append(results, "Events are delivered to a running `captain serve`; when it is down they are dropped and the daily recon backfills them.") + return strings.Join(results, "\n"), nil +} + +// ensureUserClaudeSettings returns the user-level settings.json path, creating +// an empty file when missing — monitoring hooks are user-level infrastructure +// and must be installable on a fresh machine. +func ensureUserClaudeSettings() (string, error) { + target := filepath.Join(claude.GetClaudeHome(), "settings.json") + if _, err := os.Stat(target); os.IsNotExist(err) { + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + return "", fmt.Errorf("ensure %s: %w", filepath.Dir(target), err) + } + if err := os.WriteFile(target, []byte("{}\n"), 0o644); err != nil { + return "", fmt.Errorf("create %s: %w", target, err) + } + } + return target, nil +} + +func readHookNotifyEvent(provider string, args []string) (monitor.HookEvent, error) { + switch provider { + case "claude": + if !claude.IsStdinPiped() { + return monitor.HookEvent{}, fmt.Errorf("claude hook payload must be piped on stdin") + } + data, err := os.ReadFile("/dev/stdin") + if err != nil { + return monitor.HookEvent{}, fmt.Errorf("reading stdin: %w", err) + } + return monitor.ParseClaudeHookPayload(data) + case "codex": + return monitor.ParseCodexNotifyPayload(args) + default: + return monitor.HookEvent{}, fmt.Errorf("unknown hook provider %q", provider) + } +} diff --git a/pkg/cli/hook_monitor_test.go b/pkg/cli/hook_monitor_test.go new file mode 100644 index 0000000..bc777dc --- /dev/null +++ b/pkg/cli/hook_monitor_test.go @@ -0,0 +1,52 @@ +package cli + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/flanksource/captain/pkg/codexconfig" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRunHookMonitorInstall(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + codexPath := filepath.Join(home, ".codex", "config.toml") + codexconfig.SetPathForTesting(codexPath) + t.Cleanup(func() { codexconfig.SetPathForTesting("") }) + + readHookEvents := func() map[string][]any { + t.Helper() + data, err := os.ReadFile(filepath.Join(home, ".claude", "settings.json")) + require.NoError(t, err) + var settings struct { + Hooks map[string][]any `json:"hooks"` + } + require.NoError(t, json.Unmarshal(data, &settings)) + return settings.Hooks + } + + _, err := RunHookMonitorInstall(HookMonitorInstallOptions{Timeout: 10}) + require.NoError(t, err) + + hooks := readHookEvents() + for _, event := range []string{"SessionStart", "UserPromptSubmit", "Stop", "SubagentStop", "SessionEnd"} { + assert.Len(t, hooks[event], 1, "event %s must be installed", event) + } + + codexData, err := os.ReadFile(codexPath) + require.NoError(t, err) + assert.Contains(t, string(codexData), `"hook","monitor","notify","--provider","codex"`) + + t.Run("second install is idempotent", func(t *testing.T) { + _, err := RunHookMonitorInstall(HookMonitorInstallOptions{Timeout: 10}) + require.NoError(t, err) + hooks := readHookEvents() + for _, event := range []string{"SessionStart", "UserPromptSubmit", "Stop", "SubagentStop", "SessionEnd"} { + assert.Len(t, hooks[event], 1, "event %s must not be duplicated", event) + } + }) +} diff --git a/pkg/cli/info.go b/pkg/cli/info.go index 24c68f2..5273887 100644 --- a/pkg/cli/info.go +++ b/pkg/cli/info.go @@ -4,11 +4,8 @@ import ( "fmt" "os" "path/filepath" - "sort" - "strings" "time" - "github.com/flanksource/captain/pkg/ai/history" "github.com/flanksource/captain/pkg/claude" "github.com/flanksource/clicky" "github.com/flanksource/clicky/api" @@ -129,22 +126,26 @@ func shortID(id string) string { } type InfoResult struct { - CWD string `json:"cwd"` - ProjectRoot string `json:"projectRoot"` - ProjectName string `json:"projectName"` - MarkerFile string `json:"markerFile"` - ClaudeDir string `json:"claudeDir,omitempty"` - CodexDir string `json:"codexDir,omitempty"` - Claude SourceStats `json:"claude,omitempty"` - Codex SourceStats `json:"codex,omitempty"` - TotalSessions int `json:"totalSessions"` - TotalToolCalls int `json:"totalToolCalls"` + CWD string `json:"cwd"` + CurrentSession *EnvironmentSessionInfo `json:"currentSession,omitempty"` + ProjectRoot string `json:"projectRoot"` + ProjectName string `json:"projectName"` + MarkerFile string `json:"markerFile"` + ClaudeDir string `json:"claudeDir,omitempty"` + CodexDir string `json:"codexDir,omitempty"` + Claude SourceStats `json:"claude,omitempty"` + Codex SourceStats `json:"codex,omitempty"` + TotalSessions int `json:"totalSessions"` + TotalToolCalls int `json:"totalToolCalls"` showClaude bool showCodex bool } func (r InfoResult) Pretty() api.Text { + if r.CurrentSession != nil { + return r.CurrentSession.Pretty(r.CWD) + } t := api.Text{}. Add(icons.Folder).Space(). Append(displayName(r), "font-bold text-blue-600") @@ -208,13 +209,13 @@ func infoRow(label, value string) api.Text { Append(value, "text-gray-700") } -func RunInfo(opts InfoOptions) (any, error) { +func runInfoDiscovery(opts InfoOptions) (InfoResult, error) { path := opts.Path if path == "" { var err error path, err = os.Getwd() if err != nil { - return nil, err + return InfoResult{}, err } } @@ -260,259 +261,3 @@ func RunInfo(opts InfoOptions) (any, error) { return result, nil } - -func collectClaudeStats(path string, searchAll, includeAgents bool, result *InfoResult) SourceStats { - stats := SourceStats{} - - projectsDir := claude.GetProjectsDir() - normalized := claude.NormalizePath(path) - claudeProjectDir := filepath.Join(projectsDir, normalized) - if _, err := os.Stat(claudeProjectDir); err == nil { - result.ClaudeDir = claudeProjectDir - } - - sessionFiles, err := claude.FindSessionFiles(projectsDir, path, searchAll) - if err != nil || len(sessionFiles) == 0 { - return stats - } - if !searchAll { - result.ClaudeDir = filepath.Dir(sessionFiles[0]) - } - stats.SessionCount = len(sessionFiles) - - for _, sessionFile := range sessionFiles { - entries, err := claude.ReadHistoryFile(sessionFile) - if err != nil { - continue - } - - session := SessionInfo{ID: sessionIDFromFile(sessionFile)} - for _, entry := range entries { - ts, err := entry.ParseTimestamp() - if err == nil { - updateRange(&stats, ts) - if session.StartedAt == nil || ts.Before(*session.StartedAt) { - t := ts - session.StartedAt = &t - } - if session.EndedAt == nil || ts.After(*session.EndedAt) { - t := ts - session.EndedAt = &t - } - } - if session.Model == "" && entry.Message.Model != "" { - session.Model = entry.Message.Model - } - if session.Version == "" && entry.Version != "" { - session.Version = entry.Version - } - if session.GitBranch == "" && entry.GitBranch != "" { - session.GitBranch = entry.GitBranch - } - if session.CWD == "" && entry.CWD != "" { - session.CWD = entry.CWD - } - if session.ID == "" && entry.SessionID != "" { - session.ID = entry.SessionID - } - toolCount := len(entry.Message.GetToolUses()) - session.ToolCalls += toolCount - stats.TotalToolCalls += toolCount - } - stats.Sessions = append(stats.Sessions, session) - } - - if includeAgents { - addAgentToolCalls(&stats, projectsDir, path, searchAll) - } - - sortSessionsRecent(stats.Sessions) - return stats -} - -// addAgentToolCalls folds nested sub-agent (Task/Agent) tool calls into the -// total and into their parent session's count, without inflating SessionCount — -// sub-agents belong to the session that spawned them, not to new sessions. -func addAgentToolCalls(stats *SourceStats, projectsDir, path string, searchAll bool) { - agentFiles, err := claude.FindAgentTranscripts(projectsDir, path, searchAll) - if err != nil { - return - } - byID := make(map[string]*SessionInfo, len(stats.Sessions)) - for i := range stats.Sessions { - byID[stats.Sessions[i].ID] = &stats.Sessions[i] - } - for _, af := range agentFiles { - entries, err := claude.ReadHistoryFile(af) - if err != nil { - continue - } - count, parentID := 0, "" - for _, entry := range entries { - count += len(entry.Message.GetToolUses()) - if parentID == "" && entry.SessionID != "" { - parentID = entry.SessionID - } - if ts, err := entry.ParseTimestamp(); err == nil { - updateRange(stats, ts) - } - } - stats.TotalToolCalls += count - if s := byID[parentID]; s != nil { - s.ToolCalls += count - } - } -} - -func collectCodexStats(projectRoot string, searchAll bool, result *InfoResult) SourceStats { - stats := SourceStats{} - - home, err := os.UserHomeDir() - if err == nil { - codexSessionsDir := filepath.Join(home, ".codex", "sessions") - if _, err := os.Stat(codexSessionsDir); err == nil { - result.CodexDir = codexSessionsDir - } - } - - codexFiles, err := history.FindCodexSessionFiles() - if err != nil || len(codexFiles) == 0 { - return stats - } - - for _, file := range codexFiles { - uses, err := history.ExtractCodexToolUses(file) - if err != nil || len(uses) == 0 { - continue - } - if !searchAll && !codexSessionMatchesProject(uses, projectRoot) { - continue - } - stats.SessionCount++ - stats.TotalToolCalls += len(uses) - - session := SessionInfo{ToolCalls: len(uses)} - for _, u := range uses { - if u.Timestamp != nil { - updateRange(&stats, *u.Timestamp) - if session.StartedAt == nil || u.Timestamp.Before(*session.StartedAt) { - t := *u.Timestamp - session.StartedAt = &t - } - if session.EndedAt == nil || u.Timestamp.After(*session.EndedAt) { - t := *u.Timestamp - session.EndedAt = &t - } - } - if session.ID == "" && u.SessionID != "" { - session.ID = u.SessionID - } - if session.CWD == "" && u.CWD != "" { - session.CWD = u.CWD - } - } - - if meta, err := history.ReadCodexSessionInfo(file); err == nil && meta != nil { - if session.ID == "" { - session.ID = meta.ID - } - session.Provider = meta.ModelProvider - session.Version = meta.CLIVersion - session.GitBranch = meta.GitBranch - session.Model = meta.Model - session.ReasoningEffort = meta.ReasoningEffort - if session.StartedAt == nil && meta.StartedAt != nil { - session.StartedAt = meta.StartedAt - } - } - - stats.Sessions = append(stats.Sessions, session) - } - sortSessionsRecent(stats.Sessions) - return stats -} - -// markCurrentSession flags the most-recent session (sessions are sorted -// recent-first) as the detected current session. -func markCurrentSession(sessions []SessionInfo) { - if len(sessions) > 0 { - sessions[0].Current = true - } -} - -func sortSessionsRecent(s []SessionInfo) { - sort.Slice(s, func(i, j int) bool { - ti := startedAtSort(s[i]) - tj := startedAtSort(s[j]) - return ti.After(tj) - }) -} - -func startedAtSort(s SessionInfo) time.Time { - if s.StartedAt != nil { - return *s.StartedAt - } - return time.Time{} -} - -func sessionIDFromFile(path string) string { - base := filepath.Base(path) - return strings.TrimSuffix(base, filepath.Ext(base)) -} - -// codexSessionMatchesProject returns true when at least one tool use in the -// session reports a CWD that lives within the given project root. Codex stores -// the cwd in the session_meta event, which is propagated onto every ToolUse. -func codexSessionMatchesProject(uses []history.ToolUse, projectRoot string) bool { - if projectRoot == "" { - return true - } - for _, u := range uses { - if codexCWDMatchesProject(u.CWD, projectRoot) { - return true - } - } - return false -} - -func codexCWDMatchesProject(cwd, projectRoot string) bool { - if projectRoot == "" { - return true - } - if cwd == "" { - return false - } - rootAbs := canonicalPath(projectRoot) - cwdAbs := canonicalPath(cwd) - if cwdAbs == rootAbs { - return true - } - rel, err := filepath.Rel(rootAbs, cwdAbs) - return err == nil && rel != ".." && !startsWithParent(rel) -} - -func canonicalPath(path string) string { - abs, err := filepath.Abs(path) - if err != nil { - abs = path - } - if resolved, err := filepath.EvalSymlinks(abs); err == nil { - return resolved - } - return abs -} - -func startsWithParent(rel string) bool { - return len(rel) >= 2 && rel[:2] == ".." -} - -func updateRange(stats *SourceStats, ts time.Time) { - if stats.HistoryStart == nil || ts.Before(*stats.HistoryStart) { - t := ts - stats.HistoryStart = &t - } - if stats.HistoryEnd == nil || ts.After(*stats.HistoryEnd) { - t := ts - stats.HistoryEnd = &t - } -} diff --git a/pkg/cli/info_discovery.go b/pkg/cli/info_discovery.go new file mode 100644 index 0000000..a574343 --- /dev/null +++ b/pkg/cli/info_discovery.go @@ -0,0 +1,246 @@ +package cli + +import ( + "os" + "path/filepath" + "sort" + "strings" + "time" + + "github.com/flanksource/captain/pkg/ai/history" + "github.com/flanksource/captain/pkg/claude" +) + +func collectClaudeStats(path string, searchAll, includeAgents bool, result *InfoResult) SourceStats { + stats := SourceStats{} + projectsDir := claude.GetProjectsDir() + normalized := claude.NormalizePath(path) + claudeProjectDir := filepath.Join(projectsDir, normalized) + if _, err := os.Stat(claudeProjectDir); err == nil { + result.ClaudeDir = claudeProjectDir + } + sessionFiles, err := claude.FindSessionFiles(projectsDir, path, searchAll) + if err != nil || len(sessionFiles) == 0 { + return stats + } + if !searchAll { + result.ClaudeDir = filepath.Dir(sessionFiles[0]) + } + stats.SessionCount = len(sessionFiles) + for _, sessionFile := range sessionFiles { + entries, err := claude.ReadHistoryFile(sessionFile) + if err != nil { + continue + } + session := SessionInfo{ID: sessionIDFromFile(sessionFile)} + for _, entry := range entries { + ts, err := entry.ParseTimestamp() + if err == nil { + updateRange(&stats, ts) + if session.StartedAt == nil || ts.Before(*session.StartedAt) { + t := ts + session.StartedAt = &t + } + if session.EndedAt == nil || ts.After(*session.EndedAt) { + t := ts + session.EndedAt = &t + } + } + if session.Model == "" && entry.Message.Model != "" { + session.Model = entry.Message.Model + } + if session.Version == "" && entry.Version != "" { + session.Version = entry.Version + } + if session.GitBranch == "" && entry.GitBranch != "" { + session.GitBranch = entry.GitBranch + } + if session.CWD == "" && entry.CWD != "" { + session.CWD = entry.CWD + } + if session.ID == "" && entry.SessionID != "" { + session.ID = entry.SessionID + } + toolCount := len(entry.Message.GetToolUses()) + session.ToolCalls += toolCount + stats.TotalToolCalls += toolCount + } + stats.Sessions = append(stats.Sessions, session) + } + if includeAgents { + addAgentToolCalls(&stats, projectsDir, path, searchAll) + } + sortSessionsRecent(stats.Sessions) + return stats +} + +func addAgentToolCalls(stats *SourceStats, projectsDir, path string, searchAll bool) { + agentFiles, err := claude.FindAgentTranscripts(projectsDir, path, searchAll) + if err != nil { + return + } + byID := make(map[string]*SessionInfo, len(stats.Sessions)) + for i := range stats.Sessions { + byID[stats.Sessions[i].ID] = &stats.Sessions[i] + } + for _, af := range agentFiles { + entries, err := claude.ReadHistoryFile(af) + if err != nil { + continue + } + count, parentID := 0, "" + for _, entry := range entries { + count += len(entry.Message.GetToolUses()) + if parentID == "" && entry.SessionID != "" { + parentID = entry.SessionID + } + if ts, err := entry.ParseTimestamp(); err == nil { + updateRange(stats, ts) + } + } + stats.TotalToolCalls += count + if s := byID[parentID]; s != nil { + s.ToolCalls += count + } + } +} + +func collectCodexStats(projectRoot string, searchAll bool, result *InfoResult) SourceStats { + stats := SourceStats{} + home, err := os.UserHomeDir() + if err == nil { + codexSessionsDir := filepath.Join(home, ".codex", "sessions") + if _, err := os.Stat(codexSessionsDir); err == nil { + result.CodexDir = codexSessionsDir + } + } + codexFiles, err := history.FindCodexSessionFiles() + if err != nil || len(codexFiles) == 0 { + return stats + } + for _, file := range codexFiles { + uses, err := history.ExtractCodexToolUses(file) + if err != nil || len(uses) == 0 { + continue + } + if !searchAll && !codexSessionMatchesProject(uses, projectRoot) { + continue + } + stats.SessionCount++ + stats.TotalToolCalls += len(uses) + session := SessionInfo{ToolCalls: len(uses)} + for _, u := range uses { + if u.Timestamp != nil { + updateRange(&stats, *u.Timestamp) + if session.StartedAt == nil || u.Timestamp.Before(*session.StartedAt) { + t := *u.Timestamp + session.StartedAt = &t + } + if session.EndedAt == nil || u.Timestamp.After(*session.EndedAt) { + t := *u.Timestamp + session.EndedAt = &t + } + } + if session.ID == "" && u.SessionID != "" { + session.ID = u.SessionID + } + if session.CWD == "" && u.CWD != "" { + session.CWD = u.CWD + } + } + if meta, err := history.ReadCodexSessionInfo(file); err == nil && meta != nil { + if session.ID == "" { + session.ID = meta.ID + } + session.Provider = meta.ModelProvider + session.Version = meta.CLIVersion + session.GitBranch = meta.GitBranch + session.Model = meta.Model + session.ReasoningEffort = meta.ReasoningEffort + if session.StartedAt == nil && meta.StartedAt != nil { + session.StartedAt = meta.StartedAt + } + } + stats.Sessions = append(stats.Sessions, session) + } + sortSessionsRecent(stats.Sessions) + return stats +} + +func markCurrentSession(sessions []SessionInfo) { + if len(sessions) > 0 { + sessions[0].Current = true + } +} + +func sortSessionsRecent(s []SessionInfo) { + sort.Slice(s, func(i, j int) bool { + return startedAtSort(s[i]).After(startedAtSort(s[j])) + }) +} + +func startedAtSort(s SessionInfo) time.Time { + if s.StartedAt != nil { + return *s.StartedAt + } + return time.Time{} +} + +func sessionIDFromFile(path string) string { + base := filepath.Base(path) + return strings.TrimSuffix(base, filepath.Ext(base)) +} + +func codexSessionMatchesProject(uses []history.ToolUse, projectRoot string) bool { + if projectRoot == "" { + return true + } + for _, u := range uses { + if codexCWDMatchesProject(u.CWD, projectRoot) { + return true + } + } + return false +} + +func codexCWDMatchesProject(cwd, projectRoot string) bool { + if projectRoot == "" { + return true + } + if cwd == "" { + return false + } + rootAbs := canonicalPath(projectRoot) + cwdAbs := canonicalPath(cwd) + if cwdAbs == rootAbs { + return true + } + rel, err := filepath.Rel(rootAbs, cwdAbs) + return err == nil && rel != ".." && !startsWithParent(rel) +} + +func canonicalPath(path string) string { + abs, err := filepath.Abs(path) + if err != nil { + abs = path + } + if resolved, err := filepath.EvalSymlinks(abs); err == nil { + return resolved + } + return abs +} + +func startsWithParent(rel string) bool { + return len(rel) >= 2 && rel[:2] == ".." +} + +func updateRange(stats *SourceStats, ts time.Time) { + if stats.HistoryStart == nil || ts.Before(*stats.HistoryStart) { + t := ts + stats.HistoryStart = &t + } + if stats.HistoryEnd == nil || ts.After(*stats.HistoryEnd) { + t := ts + stats.HistoryEnd = &t + } +} diff --git a/pkg/cli/info_session.go b/pkg/cli/info_session.go new file mode 100644 index 0000000..7a80890 --- /dev/null +++ b/pkg/cli/info_session.go @@ -0,0 +1,174 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/clicky/api" + "github.com/flanksource/clicky/api/icons" + clickyrpc "github.com/flanksource/clicky/rpc" + "github.com/google/uuid" +) + +// EnvironmentSessionInfo identifies the agent session exported to this process +// and the Captain database sessions bound to that identity. +type EnvironmentSessionInfo struct { + Source string `json:"source"` + SessionID string `json:"sessionId,omitempty"` + Marker string `json:"marker"` + CaptainSessions []SessionRecord `json:"captainSessions"` +} + +type infoSessionStore interface { + sessionOverviewStore + ListSessionOverviewsByProviderSessionID(context.Context, string) ([]database.SessionOverview, error) +} + +type infoRuntime struct { + getenv func(string) string + getwd func() (string, error) + openDB func(context.Context) (infoSessionStore, error) +} + +// RunInfo reports the active environment session before falling back to +// project-scoped Claude and Codex transcript discovery. +func RunInfo(ctx context.Context, opts InfoOptions) (InfoResult, error) { + return runInfo(ctx, opts, infoRuntime{ + getenv: os.Getenv, + getwd: os.Getwd, + openDB: func(ctx context.Context) (infoSessionStore, error) { return captainDB(ctx) }, + }) +} + +func runInfo(ctx context.Context, opts InfoOptions, runtime infoRuntime) (InfoResult, error) { + if _, ok := clickyrpc.RequestFromContext(ctx); ok && opts.Path != "" { + workspace, err := runtime.getwd() + if err != nil { + return InfoResult{}, fmt.Errorf("resolve info workspace: %w", err) + } + opts.Path, err = resolveCatalogDir(workspace, opts.Path) + if err != nil { + return InfoResult{}, fmt.Errorf("resolve info path: %w", err) + } + } + if !infoUsesEnvironment(opts) { + return runInfoDiscovery(opts) + } + current := detectEnvironmentSession(runtime.getenv) + if current == nil { + return runInfoDiscovery(opts) + } + cwd, err := runtime.getwd() + if err != nil { + return InfoResult{}, err + } + current.CaptainSessions = []SessionRecord{} + if current.SessionID == "" { + return InfoResult{CWD: cwd, CurrentSession: current}, nil + } + store, err := runtime.openDB(ctx) + if err != nil { + return InfoResult{}, fmt.Errorf("resolve %s session %q: %w", current.Source, current.SessionID, err) + } + resolved, err := resolveEnvironmentSession(ctx, store, *current) + if err != nil { + return InfoResult{}, err + } + return InfoResult{CWD: cwd, CurrentSession: &resolved}, nil +} + +func infoUsesEnvironment(opts InfoOptions) bool { + return opts.Path == "" && !opts.All && !opts.Claude && !opts.Codex +} + +// CurrentEnvironmentSession returns the active agent session identified by +// Captain's provider environment markers. +func CurrentEnvironmentSession() *EnvironmentSessionInfo { + return detectEnvironmentSession(os.Getenv) +} + +func detectEnvironmentSession(getenv func(string) string) *EnvironmentSessionInfo { + type marker struct { + source string + name string + id bool + value string + } + markers := []marker{ + {source: "codex", name: "CODEX_THREAD_ID", id: true}, + {source: "codex", name: "CODEX_SESSION_ID", id: true}, + {source: "codex", name: "CODEX_SANDBOX"}, + {source: "claude", name: "CLAUDE_CODE_SESSION_ID", id: true}, + {source: "claude", name: "CLAUDE_SESSION_ID", id: true}, + {source: "claude", name: "CLAUDECODE", value: "1"}, + {source: "gemini", name: "GEMINI_SESSION_ID", id: true}, + {source: "gemini", name: "GEMINI_CLI", value: "1"}, + {source: "captain", name: "CAPTAIN_SESSION_ID", id: true}, + } + for _, candidate := range markers { + value := strings.TrimSpace(getenv(candidate.name)) + if value == "" || candidate.value != "" && value != candidate.value { + continue + } + result := &EnvironmentSessionInfo{Source: candidate.source, Marker: candidate.name} + if candidate.id { + result.SessionID = value + } + return result + } + return nil +} + +func resolveEnvironmentSession(ctx context.Context, store infoSessionStore, current EnvironmentSessionInfo) (EnvironmentSessionInfo, error) { + var ( + overviews []database.SessionOverview + err error + ) + if current.Source == "captain" { + if _, parseErr := uuid.Parse(current.SessionID); parseErr != nil { + return EnvironmentSessionInfo{}, fmt.Errorf("%s must contain a valid Captain UUID: %w", current.Marker, parseErr) + } + overviews, err = resolveOverviewsByIdentity(ctx, store, current.SessionID) + if errors.Is(err, database.ErrSessionNotFound) { + overviews, err = nil, nil + } + } else { + overviews, err = store.ListSessionOverviewsByProviderSessionID(ctx, current.SessionID) + } + if err != nil { + return EnvironmentSessionInfo{}, fmt.Errorf("resolve %s session from %s: %w", current.Source, current.Marker, err) + } + current.CaptainSessions = make([]SessionRecord, len(overviews)) + for i := range overviews { + current.CaptainSessions[i] = recordFromOverview(overviews[i]) + } + return current, nil +} + +func (s EnvironmentSessionInfo) Pretty(cwd string) api.Text { + name := s.Source + if name != "" { + name = strings.ToUpper(name[:1]) + name[1:] + } + t := api.Text{}. + Add(icons.AI).Space(). + Append(name+" session", "font-bold text-blue-600"). + NewLine().Add(infoRow("CWD", cwd)). + NewLine().Add(infoRow("Detected by", s.Marker)) + if s.SessionID != "" { + t = t.NewLine().Add(infoRow("Session", s.SessionID)) + } + t = t.NewLine().NewLine().Append("Captain sessions", "font-bold text-blue-600") + if len(s.CaptainSessions) == 0 { + return t.NewLine().Append(" (no matching sessions found)", "text-gray-500 italic") + } + rows := make([]sessionLiveRow, len(s.CaptainSessions)) + for i := range s.CaptainSessions { + rows[i] = sessionLiveRow{SessionRecord: s.CaptainSessions[i]} + } + return t.NewLine().Add(api.NewTableFrom(rows)) +} diff --git a/pkg/cli/info_session_ginkgo_test.go b/pkg/cli/info_session_ginkgo_test.go new file mode 100644 index 0000000..2395c5a --- /dev/null +++ b/pkg/cli/info_session_ginkgo_test.go @@ -0,0 +1,236 @@ +package cli + +import ( + "context" + "errors" + "net/http/httptest" + "os" + "path/filepath" + + "github.com/flanksource/captain/pkg/database" + clickyrpc "github.com/flanksource/clicky/rpc" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type infoSessionStoreStub struct { + providerID string + providerOverviews []database.SessionOverview + providerErr error + identities []string + identityOverviews []database.SessionOverview + identityErr error + threadRoots []uuid.UUID + threadOverviews []database.SessionOverview +} + +func (s *infoSessionStoreStub) ListSessionOverviewsByProviderSessionID(_ context.Context, id string) ([]database.SessionOverview, error) { + s.providerID = id + return s.providerOverviews, s.providerErr +} + +func (s *infoSessionStoreStub) ListSessionOverviewsByIdentity(_ context.Context, identity string) ([]database.SessionOverview, error) { + s.identities = append(s.identities, identity) + return s.identityOverviews, s.identityErr +} + +func (s *infoSessionStoreStub) ListThreadSessionOverviews(_ context.Context, rootID uuid.UUID) ([]database.SessionOverview, error) { + s.threadRoots = append(s.threadRoots, rootID) + return s.threadOverviews, nil +} + +var _ = Describe("captain info environment session", func() { + DescribeTable("detects the active agent before discovery", + func(env map[string]string, expected *EnvironmentSessionInfo) { + actual := detectEnvironmentSession(func(key string) string { return env[key] }) + Expect(actual).To(Equal(expected)) + }, + Entry("codex thread", map[string]string{ + "CODEX_THREAD_ID": "codex-thread", "CLAUDE_CODE_SESSION_ID": "claude-session", + }, &EnvironmentSessionInfo{Source: "codex", SessionID: "codex-thread", Marker: "CODEX_THREAD_ID"}), + Entry("codex compatibility session", map[string]string{ + "CODEX_SESSION_ID": "codex-session", + }, &EnvironmentSessionInfo{Source: "codex", SessionID: "codex-session", Marker: "CODEX_SESSION_ID"}), + Entry("codex provider marker", map[string]string{ + "CODEX_SANDBOX": "seatbelt", + }, &EnvironmentSessionInfo{Source: "codex", Marker: "CODEX_SANDBOX"}), + Entry("claude code session", map[string]string{ + "CLAUDE_CODE_SESSION_ID": "claude-session", + }, &EnvironmentSessionInfo{Source: "claude", SessionID: "claude-session", Marker: "CLAUDE_CODE_SESSION_ID"}), + Entry("legacy claude session", map[string]string{ + "CLAUDE_SESSION_ID": "legacy-claude", + }, &EnvironmentSessionInfo{Source: "claude", SessionID: "legacy-claude", Marker: "CLAUDE_SESSION_ID"}), + Entry("claude provider marker", map[string]string{ + "CLAUDECODE": "1", + }, &EnvironmentSessionInfo{Source: "claude", Marker: "CLAUDECODE"}), + Entry("gemini session", map[string]string{ + "GEMINI_SESSION_ID": "gemini-session", + }, &EnvironmentSessionInfo{Source: "gemini", SessionID: "gemini-session", Marker: "GEMINI_SESSION_ID"}), + Entry("gemini provider marker", map[string]string{ + "GEMINI_CLI": "1", + }, &EnvironmentSessionInfo{Source: "gemini", Marker: "GEMINI_CLI"}), + Entry("captain session", map[string]string{ + "CAPTAIN_SESSION_ID": "055781c7-360a-4eb2-80be-452b3937fcfe", + }, &EnvironmentSessionInfo{Source: "captain", SessionID: "055781c7-360a-4eb2-80be-452b3937fcfe", Marker: "CAPTAIN_SESSION_ID"}), + Entry("false boolean markers", map[string]string{ + "CLAUDECODE": "0", "GEMINI_CLI": "false", + }, nil), + Entry("blank markers", map[string]string{ + "CODEX_THREAD_ID": " ", "CLAUDE_CODE_SESSION_ID": "", "CAPTAIN_SESSION_ID": " ", + }, nil), + ) + + It("exports the active environment session detector", func() { + for _, marker := range []string{ + "CODEX_THREAD_ID", "CODEX_SESSION_ID", "CODEX_SANDBOX", + "CLAUDE_CODE_SESSION_ID", "CLAUDE_SESSION_ID", "CLAUDECODE", + "GEMINI_SESSION_ID", "GEMINI_CLI", "CAPTAIN_SESSION_ID", + } { + GinkgoT().Setenv(marker, "") + } + GinkgoT().Setenv("CODEX_THREAD_ID", "codex-thread") + + Expect(CurrentEnvironmentSession()).To(Equal(&EnvironmentSessionInfo{ + Source: "codex", SessionID: "codex-thread", Marker: "CODEX_THREAD_ID", + })) + }) + + DescribeTable("preserves explicit discovery flags", + func(opts InfoOptions) { Expect(infoUsesEnvironment(opts)).To(BeFalse()) }, + Entry("all", InfoOptions{All: true}), + Entry("claude", InfoOptions{Claude: true}), + Entry("codex", InfoOptions{Codex: true}), + Entry("path", InfoOptions{Path: "/repo"}), + ) + + It("uses the environment fast path for a bare invocation", func() { + Expect(infoUsesEnvironment(InfoOptions{})).To(BeTrue()) + }) + + It("confines generated RPC discovery to the server workspace", func(ctx SpecContext) { + workspace := GinkgoT().TempDir() + outside := GinkgoT().TempDir() + request := httptest.NewRequest("POST", "/api/v1/info", nil) + requestContext := clickyrpc.ContextWithRequest(ctx, request) + + _, err := runInfo(requestContext, InfoOptions{Path: outside}, infoRuntime{ + getenv: func(string) string { return "" }, + getwd: func() (string, error) { return workspace, nil }, + }) + + Expect(err).To(MatchError(ContainSubstring("escapes workspace root"))) + }) + + It("rejects generated RPC discovery through a workspace symlink", func(ctx SpecContext) { + workspace := GinkgoT().TempDir() + outside := GinkgoT().TempDir() + link := filepath.Join(workspace, "outside") + Expect(os.Symlink(outside, link)).To(Succeed()) + request := httptest.NewRequest("POST", "/api/v1/info", nil) + requestContext := clickyrpc.ContextWithRequest(ctx, request) + + _, err := runInfo(requestContext, InfoOptions{Path: link}, infoRuntime{ + getenv: func(string) string { return "" }, + getwd: func() (string, error) { return workspace, nil }, + }) + + Expect(err).To(MatchError(ContainSubstring("escapes workspace root"))) + }) + + It("returns a provider-only session without opening the database", func(ctx SpecContext) { + databaseOpened := false + result, err := runInfo(ctx, InfoOptions{}, infoRuntime{ + getenv: func(key string) string { + if key == "GEMINI_CLI" { + return "1" + } + return "" + }, + getwd: func() (string, error) { return "/repo", nil }, + openDB: func(context.Context) (infoSessionStore, error) { + databaseOpened = true + return nil, errors.New("unexpected database open") + }, + }) + + Expect(err).NotTo(HaveOccurred()) + Expect(databaseOpened).To(BeFalse()) + Expect(result.CurrentSession).To(Equal(&EnvironmentSessionInfo{ + Source: "gemini", Marker: "GEMINI_CLI", CaptainSessions: []SessionRecord{}, + })) + }) + + It("returns an empty Captain session list when an exact provider ID is not stored", func(ctx SpecContext) { + store := &infoSessionStoreStub{} + result, err := runInfo(ctx, InfoOptions{}, infoRuntime{ + getenv: func(key string) string { + if key == "CODEX_THREAD_ID" { + return "codex-thread" + } + return "" + }, + getwd: func() (string, error) { return "/repo", nil }, + openDB: func(context.Context) (infoSessionStore, error) { return store, nil }, + }) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.CurrentSession.CaptainSessions).To(BeEmpty()) + Expect(result.CurrentSession.CaptainSessions).NotTo(BeNil()) + }) + + It("returns every exact provider-session database match", func(ctx SpecContext) { + firstID := uuid.MustParse("055781c7-360a-4eb2-80be-452b3937fcfe") + secondID := uuid.MustParse("7ca78c55-e280-50ff-a19a-9f355a6fc55e") + providerID := "019f7c25-9adf-7901-add9-8c46693472fb" + store := &infoSessionStoreStub{providerOverviews: []database.SessionOverview{ + {ID: firstID, ProviderSessionID: &providerID, Source: "codex"}, + {ID: secondID, ProviderSessionID: &providerID, Source: "captain"}, + }} + session := EnvironmentSessionInfo{Source: "codex", SessionID: providerID, Marker: "CODEX_THREAD_ID"} + + resolved, err := resolveEnvironmentSession(ctx, store, session) + + Expect(err).NotTo(HaveOccurred()) + Expect(store.providerID).To(Equal(providerID)) + Expect(resolved.CaptainSessions).To(HaveLen(2)) + Expect(resolved.CaptainSessions[0].Key).To(Equal(firstID.String())) + Expect(resolved.CaptainSessions[1].Key).To(Equal(secondID.String())) + }) + + It("expands a Captain root session through the existing thread resolver", func(ctx SpecContext) { + rootID := uuid.MustParse("055781c7-360a-4eb2-80be-452b3937fcfe") + childID := uuid.MustParse("7ca78c55-e280-50ff-a19a-9f355a6fc55e") + store := &infoSessionStoreStub{ + identityOverviews: []database.SessionOverview{{ID: rootID, Source: "captain"}}, + threadOverviews: []database.SessionOverview{ + {ID: rootID, Source: "captain"}, {ID: childID, RootSessionID: &rootID, Source: "codex"}, + }, + } + session := EnvironmentSessionInfo{Source: "captain", SessionID: rootID.String(), Marker: "CAPTAIN_SESSION_ID"} + + resolved, err := resolveEnvironmentSession(ctx, store, session) + + Expect(err).NotTo(HaveOccurred()) + Expect(store.identities).To(Equal([]string{rootID.String()})) + Expect(store.threadRoots).To(Equal([]uuid.UUID{rootID})) + Expect(resolved.CaptainSessions).To(HaveLen(2)) + }) + + It("fails loudly for an invalid Captain session ID", func(ctx SpecContext) { + _, err := resolveEnvironmentSession(ctx, &infoSessionStoreStub{}, EnvironmentSessionInfo{ + Source: "captain", SessionID: "not-a-uuid", Marker: "CAPTAIN_SESSION_ID", + }) + + Expect(err).To(MatchError(ContainSubstring("CAPTAIN_SESSION_ID"))) + }) + + It("propagates exact provider lookup failures", func(ctx SpecContext) { + lookupErr := errors.New("database unavailable") + _, err := resolveEnvironmentSession(ctx, &infoSessionStoreStub{providerErr: lookupErr}, EnvironmentSessionInfo{ + Source: "codex", SessionID: "codex-thread", Marker: "CODEX_THREAD_ID", + }) + + Expect(err).To(MatchError(ContainSubstring("database unavailable"))) + }) +}) diff --git a/pkg/cli/model_selection_ginkgo_test.go b/pkg/cli/model_selection_ginkgo_test.go new file mode 100644 index 0000000..48c5180 --- /dev/null +++ b/pkg/cli/model_selection_ginkgo_test.go @@ -0,0 +1,119 @@ +package cli + +import ( + "context" + "os" + "path/filepath" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/aiflags" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("CLI model selection", func() { + BeforeEach(func() { + path := filepath.Join(GinkgoT().TempDir(), ".captain.yaml") + Expect(os.WriteFile(path, []byte("ai:\n model: opus\n backend: claude-agent\n"), 0o600)).To(Succeed()) + captainconfig.SetPathForTesting(path) + DeferCleanup(func() { captainconfig.SetPathForTesting("") }) + }) + + It("does not attach the saved Claude backend to an explicit model", func() { + opts := AIPromptOptions{} + opts.Model = "gemini-3.5-flash" + + req, cfg, err := overlayCLI(ai.Request{}, ai.Config{}, opts) + + Expect(err).NotTo(HaveOccurred()) + Expect(req.Model.Name).To(Equal("gemini-3.5-flash")) + Expect(req.Model.Backend).To(Equal(api.BackendGemini)) + Expect(cfg.Model).To(Equal(req.Model)) + }) + + It("uses the same precedence for provider options", func() { + cfg, err := (AIProviderOptions{ModelFlags: aiflags.ModelFlags{Model: "gemini-3.5-flash"}}).ToConfig() + + Expect(err).NotTo(HaveOccurred()) + Expect(cfg.Model.Name).To(Equal("gemini-3.5-flash")) + Expect(cfg.Model.Backend).To(Equal(api.BackendGemini)) + }) + + It("keeps the saved model and backend paired when there is no override", func() { + cfg, err := (AIProviderOptions{}).ToConfig() + + Expect(err).NotTo(HaveOccurred()) + Expect(cfg.Model.Name).To(Equal("claude-opus-5")) + Expect(cfg.Model.Backend).To(Equal(api.BackendClaudeAgent)) + }) + + It("does not retain a prompt backend when a structured spec replaces its model", func() { + req := ai.Request{Model: api.Model{Name: "opus", Backend: api.BackendClaudeAgent}} + cfg := ai.Config{Model: req.Model} + + overlayRuntimeSpec(&req, &cfg, api.Spec{Model: api.Model{Name: "gemini-3.5-flash"}}) + Expect(applyPromptDefaults(&req, &cfg)).To(Succeed()) + resolved, err := ai.ResolveModelSelectors(cfg.Model) + + Expect(err).NotTo(HaveOccurred()) + Expect(resolved.Name).To(Equal("gemini-3.5-flash")) + Expect(resolved.Backend).To(Equal(api.BackendGemini)) + }) + + It("omits single-runtime identity from multi-model parent labels", func() { + labels := promptTaskLabels(PromptRenderResult{ + Name: "test", Model: "opus", Backend: string(api.BackendClaudeAgent), + }, "multi") + + Expect(labels).To(HaveKeyWithValue("mode", "multi")) + Expect(labels).NotTo(HaveKey("model")) + Expect(labels).NotTo(HaveKey("backend")) + }) + + It("passes a concrete Gemini runtime to a bare multi-model execution", func() { + originalExecute := executePromptRequestFunc + DeferCleanup(func() { executePromptRequestFunc = originalExecute }) + var executed ai.Config + executePromptRequestFunc = func(_ context.Context, _ ai.Request, cfg ai.Config, _ time.Duration, _ bool) (any, error) { + executed = cfg + return AIPromptResult{Text: "ok", Model: cfg.Model.Name, Backend: string(cfg.Model.Backend)}, nil + } + + result, err := executeSyncBatch( + context.Background(), + testRenderedPrompt(api.Model{Name: "opus", Backend: api.BackendClaudeAgent}), + AIPromptOptions{AIRuntimeOptions: AIRuntimeOptions{AIProviderOptions: AIProviderOptions{}}, MultiModels: []string{"gemini-3.5-flash"}}, + ) + + Expect(err).NotTo(HaveOccurred()) + Expect(executed.Model.Name).To(Equal("gemini-3.5-flash")) + Expect(executed.Model.Backend).To(Equal(api.BackendGemini)) + Expect(result.Runs).To(HaveLen(1)) + Expect(result.Runs[0].Selector).To(Equal("gemini:gemini-3.5-flash")) + }) + + It("replaces the configured runtime identity for a multi-model variant", func() { + configured := api.Model{Name: "gpt-5.6-luna", Backend: api.BackendCodexCLI}.Capabilities() + selected := api.Model{Name: "gemini-3.5-flash", Backend: api.BackendGemini, Effort: api.EffortHigh}.Capabilities() + + variant := renderVariant(testRenderedPrompt(configured), selected, nil).Config.Model + + Expect(variant.Validate()).To(Succeed()) + Expect(variant).To(Equal(selected)) + }) + + It("preserves runtime fallbacks when no CLI fallback override is set", func() { + selected := api.Model{ + Name: "gemini-3.5-flash", + Backend: api.BackendGemini, + Fallbacks: api.ModelList{{Name: "gemini-3-flash", Backend: api.BackendGemini}}, + } + + variant := renderVariant(testRenderedPrompt(api.Model{}), selected, nil).Config.Model + + Expect(variant.Fallbacks).To(Equal(selected.Fallbacks)) + }) +}) diff --git a/pkg/cli/pathfilter.go b/pkg/cli/pathfilter.go index 3ea94c5..6274ddb 100644 --- a/pkg/cli/pathfilter.go +++ b/pkg/cli/pathfilter.go @@ -33,18 +33,14 @@ func (f *pathFilter) keep(path string) bool { if path == "" { return true } - // changes/transcript paths are relativized to the project root; absolutize - // (against the working dir) so plan-path and gitignore checks are reliable. - abs := path - if !filepath.IsAbs(abs) { - if a, err := filepath.Abs(abs); err == nil { - abs = a - } - } - if !f.includePlans && strings.Contains(filepath.ToSlash(abs), "/.claude/plans/") { + // Paths arrive absolute (ToolAnalysis anchors them to the tool use's cwd), so + // the plan-path and gitignore checks can use them directly. They previously + // had to be re-absolutized against this process's working directory, which + // silently mis-resolved any path relativised against a different base. + if !f.includePlans && strings.Contains(filepath.ToSlash(path), "/.claude/plans/") { return false } - if !f.includeIgnored && f.isIgnored(abs) { + if !f.includeIgnored && f.isIgnored(path) { return false } return true diff --git a/pkg/cli/permission_catalog.go b/pkg/cli/permission_catalog.go index 3d02967..d00ee09 100644 --- a/pkg/cli/permission_catalog.go +++ b/pkg/cli/permission_catalog.go @@ -45,6 +45,10 @@ func resolveCatalogDir(baseCwd, dir string) (string, error) { if dir == "" { return base, nil } + resolvedBase, err := filepath.EvalSymlinks(base) + if err != nil { + return "", fmt.Errorf("resolve workspace root: %w", err) + } // Reject traversal sequences in the raw input before any path is built; // the prefix check below is defense in depth. @@ -57,8 +61,12 @@ func resolveCatalogDir(baseCwd, dir string) (string, error) { target = filepath.Join(base, target) } target = filepath.Clean(target) + resolvedTarget, err := filepath.EvalSymlinks(target) + if err != nil { + return "", fmt.Errorf("resolve dir %q: %w", dir, err) + } - if target != base && !strings.HasPrefix(target, base+string(os.PathSeparator)) { + if resolvedTarget != resolvedBase && !strings.HasPrefix(resolvedTarget, resolvedBase+string(os.PathSeparator)) { return "", fmt.Errorf("dir %q escapes workspace root", dir) } return target, nil diff --git a/pkg/cli/plan.go b/pkg/cli/plan.go index 5a02138..a84a010 100644 --- a/pkg/cli/plan.go +++ b/pkg/cli/plan.go @@ -2,13 +2,14 @@ package cli import ( "context" + "errors" "fmt" "os" - "sort" "strings" "github.com/flanksource/captain/pkg/ai/history" "github.com/flanksource/captain/pkg/claude" + captaindb "github.com/flanksource/captain/pkg/database" captainsession "github.com/flanksource/captain/pkg/session" "github.com/flanksource/clicky" "github.com/flanksource/clicky/api" @@ -26,17 +27,21 @@ func (PlanOptions) GetName() string { return "plan [session-id]" } // PlanResult is the exit-plan-mode plan for a single session. type PlanResult struct { - SessionID string `json:"sessionId" pretty:"label=Session"` - Source string `json:"source,omitempty" pretty:"label=Source"` - Path string `json:"path,omitempty" pretty:"label=Plan"` - OnDisk bool `json:"onDisk" pretty:"label=On Disk"` - Slug string `json:"slug,omitempty" pretty:"label=Slug"` - Content string `json:"content,omitempty"` + SessionID string `json:"sessionId" pretty:"label=Session"` + PlanID string `json:"planId,omitempty" pretty:"label=Plan ID"` + RevisionID string `json:"revisionId,omitempty"` + Revision int `json:"revision,omitempty"` + Source string `json:"source,omitempty" pretty:"label=Source"` + Path string `json:"path,omitempty" pretty:"label=Plan"` + OnDisk bool `json:"onDisk" pretty:"label=On Disk"` + Slug string `json:"slug,omitempty" pretty:"label=Slug"` + Content string `json:"content,omitempty"` pathOnly bool } func RunPlan(opts PlanOptions) (PlanResult, error) { + ctx := context.Background() source, err := normalizeSessionSource(opts.Source) if err != nil { return PlanResult{}, err @@ -45,32 +50,28 @@ func RunPlan(opts PlanOptions) (PlanResult, error) { if err != nil { return PlanResult{}, err } + db, err := freshenSessionDB(ctx) + if err != nil { + return PlanResult{}, err + } id := strings.TrimSpace(opts.SessionID) if id != "" { - if candidate, ok, err := findSessionCandidateByID(id, source); err != nil { + persisted, ok, err := resolveNativePlan(ctx, db, id, source) + if err != nil { return PlanResult{}, err - } else if ok { - plan, err := resolveSessionPlan(candidate) - if err != nil { - return PlanResult{}, err - } - if plan == nil { - return PlanResult{}, fmt.Errorf("session %q has no plan", id) - } - plan.pathOnly = opts.PathOnly - return *plan, nil } - - // A hashed session key cannot be found from the transcript filename, so - // keep the older all-project scan as a fallback for that less-common form. - candidates, err := discoverSessionCandidates(context.Background(), "", true, source) + if ok { + persisted.pathOnly = opts.PathOnly + return *persisted, nil + } + overview, err := db.GetSessionOverviewByIdentity(ctx, id) if err != nil { return PlanResult{}, err } - candidate, ok := matchPlanCandidate(candidates, id) - if !ok { - return PlanResult{}, fmt.Errorf("session %q not found", id) + candidate := candidateFromOverview(*overview) + if candidate.path == "" { + return PlanResult{}, fmt.Errorf("session %q has no transcript recorded on this host", id) } plan, err := resolveSessionPlan(candidate) if err != nil { @@ -83,33 +84,119 @@ func RunPlan(opts PlanOptions) (PlanResult, error) { return *plan, nil } - candidates, err := discoverSessionCandidates(context.Background(), cwd, opts.All, source) + _, projectRoot, _ := resolveSessionScope(cwd, opts.All, "") + plan, err := resolveLatestTranscriptPlan(ctx, db, latestTranscriptPlanQuery{ + Source: source, ProjectRoot: projectRoot, + }) if err != nil { return PlanResult{}, err } - sort.SliceStable(candidates, func(i, j int) bool { - return sessionSortTime(candidates[i].record).After(sessionSortTime(candidates[j].record)) - }) - for _, candidate := range candidates { - plan, err := resolveSessionPlan(candidate) - if err != nil || plan == nil { - continue + if plan == nil { + return PlanResult{}, fmt.Errorf("no session with a plan found in %s", scopeLabel(cwd, opts.All)) + } + plan.pathOnly = opts.PathOnly + return *plan, nil +} + +type latestTranscriptPlanQuery struct { + Source string + ProjectRoot string +} + +const latestTranscriptPlanPageLimit = 100 + +func resolveLatestTranscriptPlan( + ctx context.Context, + db sessionListStore, + query latestTranscriptPlanQuery, +) (*PlanResult, error) { + filter := captaindb.SessionListFilter{ + ProjectRoot: query.ProjectRoot, + RootsOnly: true, + Limit: latestTranscriptPlanPageLimit, + } + if query.Source != "all" { + filter.Source = query.Source + } + seen := map[string]struct{}{} + for { + page, err := db.ListSessionSummaries(ctx, filter) + if err != nil { + return nil, fmt.Errorf("list Captain sessions while resolving latest plan: %w", err) } - plan.pathOnly = opts.PathOnly - return *plan, nil + for i := range page.Rows { + candidate := candidateFromOverview(overviewFromSummary(page.Rows[i])) + if candidate.path == "" { + continue + } + plan, err := resolveSessionPlan(candidate) + if err == nil && plan != nil { + return plan, nil + } + } + if page.NextCursor == "" { + return nil, nil + } + if _, ok := seen[page.NextCursor]; ok { + return nil, fmt.Errorf("captain plan search pagination repeated cursor %q", page.NextCursor) + } + seen[page.NextCursor] = struct{}{} + filter.Cursor = page.NextCursor } - return PlanResult{}, fmt.Errorf("no session with a plan found in %s", scopeLabel(cwd, opts.All)) } -// matchPlanCandidate finds the candidate whose record key or id matches id -// exactly, or whose id has id as a prefix (so the short ids printed elsewhere work). -func matchPlanCandidate(candidates []sessionCandidate, id string) (sessionCandidate, bool) { - for _, c := range candidates { - if c.record.Key == id || c.record.ID == id || (c.record.ID != "" && strings.HasPrefix(c.record.ID, id)) { - return c, true +// resolveNativePlan resolves persisted plan content without consulting the +// transcript or source plan path. Approved content wins; otherwise the latest +// immutable revision of the newest plan variant is returned. +func resolveNativePlan(ctx context.Context, db *captaindb.DB, identity, source string) (*PlanResult, bool, error) { + sourceFilter := source + if sourceFilter == "all" { + sourceFilter = "" + } + session, err := db.GetSessionByIdentity(ctx, identity, sourceFilter, "", "") + if err != nil { + if errors.Is(err, captaindb.ErrSessionNotFound) { + return nil, false, nil } + return nil, false, fmt.Errorf("resolve persisted Captain session %q: %w", identity, err) } - return sessionCandidate{}, false + plans, err := db.ListPlans(ctx, captaindb.PlanFilter{SourceSessionID: &session.ID}) + if err != nil { + return nil, false, fmt.Errorf("list persisted plans for session %s: %w", session.ID, err) + } + var selected *captaindb.Plan + for i := range plans { + if plans[i].ApprovedRevision != nil { + selected = &plans[i] + break + } + if selected == nil && plans[i].LatestRevision != nil { + selected = &plans[i] + } + } + if selected == nil { + return nil, false, nil + } + revision := selected.ApprovedRevision + if revision == nil { + revision = selected.LatestRevision + } + onDisk := false + if selected.Path != "" { + _, err := os.Stat(selected.Path) + onDisk = err == nil + } + return &PlanResult{ + SessionID: session.ID.String(), + PlanID: selected.ID.String(), + RevisionID: revision.ID.String(), + Revision: revision.Revision, + Source: session.Source, + Path: selected.Path, + OnDisk: onDisk, + Slug: selected.Slug, + Content: revision.PlanMarkdown, + }, true, nil } // resolveSessionPlan reads a session transcript and recovers its plan. It returns diff --git a/pkg/cli/plan_native_integration_test.go b/pkg/cli/plan_native_integration_test.go new file mode 100644 index 0000000..8f715db --- /dev/null +++ b/pkg/cli/plan_native_integration_test.go @@ -0,0 +1,66 @@ +package cli + +import ( + "os" + "path/filepath" + "testing" + + captaindb "github.com/flanksource/captain/pkg/database" + commonsdb "github.com/flanksource/commons-db/db" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResolveNativePlanUsesPersistedApprovedContentWithoutSourceFile(t *testing.T) { + if os.Getenv("CAPTAIN_DB_EMBEDDED_TEST") == "" { + t.Skip("set CAPTAIN_DB_EMBEDDED_TEST=1 to run embedded-postgres plan lookup tests") + } + + dsn, stop, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: filepath.Join(t.TempDir(), "postgres"), + Database: "captain_native_plan_lookup", + }) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, stop()) }) + db, err := captaindb.Open(t.Context(), captaindb.WithDSN(dsn), captaindb.WithMigrations()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, db.Close()) }) + + session, err := db.CreateOrGetSession(t.Context(), captaindb.CreateSessionInput{ + ProviderSessionID: "provider-plan-session", Source: "codex", Provider: "openai", + }) + require.NoError(t, err) + deletedPath := filepath.Join(t.TempDir(), "deleted-plan.md") + require.NoError(t, os.WriteFile(deletedPath, []byte("stale disk content"), 0o600)) + plan, err := db.CreateOrGetPlan(t.Context(), captaindb.CreatePlanInput{ + SourceSessionID: session.ID, Variant: "approved", Path: deletedPath, Slug: "durable", + }) + require.NoError(t, err) + first, err := db.AppendPlanRevision(t.Context(), captaindb.AppendPlanRevisionInput{ + PlanID: plan.ID, PlanMarkdown: "# Approved durable plan", + }) + require.NoError(t, err) + _, err = db.AppendPlanRevision(t.Context(), captaindb.AppendPlanRevisionInput{ + PlanID: plan.ID, PlanMarkdown: "# Newer but unapproved plan", + }) + require.NoError(t, err) + _, err = db.ApprovePlanRevision(t.Context(), captaindb.ApprovePlanRevisionInput{ + PlanID: plan.ID, RevisionID: first.ID, + }) + require.NoError(t, err) + require.NoError(t, os.Remove(deletedPath)) + + result, ok, err := resolveNativePlan(t.Context(), db, "provider-plan-session", "all") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, session.ID.String(), result.SessionID) + assert.Equal(t, plan.ID.String(), result.PlanID) + assert.Equal(t, first.ID.String(), result.RevisionID) + assert.Equal(t, "# Approved durable plan", result.Content) + assert.False(t, result.OnDisk) + + byUUID, ok, err := resolveNativePlan(t.Context(), db, session.ID.String(), "codex") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, result.Content, byUUID.Content) +} diff --git a/pkg/cli/plan_search_ginkgo_test.go b/pkg/cli/plan_search_ginkgo_test.go new file mode 100644 index 0000000..a994517 --- /dev/null +++ b/pkg/cli/plan_search_ginkgo_test.go @@ -0,0 +1,83 @@ +package cli + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + + "github.com/flanksource/captain/pkg/database" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + . "github.com/onsi/gomega/gstruct" +) + +var _ = Describe("latest transcript plan search", func() { + It("pages bounded session summaries until the newest transcript plan", func(ctx SpecContext) { + home := GinkgoT().TempDir() + project := filepath.Join(home, "work", "captain") + planPath := filepath.Join(home, ".claude", "plans", "paged-plan.md") + historyPath := filepath.Join(home, ".claude", "projects", "captain", "session-with-plan.jsonl") + entry, err := json.Marshal(map[string]any{ + "type": "assistant", "sessionId": "session-with-plan", "uuid": "assistant-1", + "timestamp": "2026-07-16T10:00:00Z", "cwd": project, "slug": "paged-plan", + "message": map[string]any{"role": "assistant", "content": []any{ + map[string]any{"type": "tool_use", "id": "tool-1", "name": "ExitPlanMode", + "input": map[string]any{"planFilePath": planPath, "plan": "# paged plan"}}, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(os.MkdirAll(filepath.Dir(historyPath), 0o755)).To(Succeed()) + Expect(os.WriteFile(historyPath, append(entry, '\n'), 0o644)).To(Succeed()) + missingPath := filepath.Join(home, ".claude", "projects", "missing.jsonl") + providerID := "session-with-plan" + store := &planListStore{pages: map[string]database.SessionListPage{ + "": { + Rows: []database.SessionListSummary{{ID: uuid.New(), Source: "claude", Path: &missingPath}}, + NextCursor: "second-page", + }, + "second-page": { + Rows: []database.SessionListSummary{{ + ID: uuid.New(), ProviderSessionID: &providerID, Source: "claude", Path: &historyPath, + }}, + NextCursor: "unused-third-page", + }, + }} + + plan, err := resolveLatestTranscriptPlan(ctx, store, latestTranscriptPlanQuery{ + Source: "claude", ProjectRoot: project, + }) + + Expect(err).NotTo(HaveOccurred()) + Expect(plan).NotTo(BeNil()) + Expect(*plan).To(MatchFields(IgnoreExtras, Fields{ + "SessionID": Equal(providerID), + "Source": Equal("claude"), + "Path": Equal(planPath), + "Content": Equal("# paged plan"), + })) + Expect(store.filters).To(HaveLen(2)) + Expect(store.filters[0]).To(MatchFields(IgnoreExtras, Fields{ + "Source": Equal("claude"), + "ProjectRoot": Equal(project), + "RootsOnly": BeTrue(), + "Limit": Equal(latestTranscriptPlanPageLimit), + "Cursor": BeEmpty(), + })) + Expect(store.filters[1].Cursor).To(Equal("second-page")) + }) +}) + +type planListStore struct { + pages map[string]database.SessionListPage + filters []database.SessionListFilter +} + +func (s *planListStore) ListSessionSummaries( + _ context.Context, + filter database.SessionListFilter, +) (database.SessionListPage, error) { + s.filters = append(s.filters, filter) + return s.pages[filter.Cursor], nil +} diff --git a/pkg/cli/plan_test.go b/pkg/cli/plan_test.go index 138d705..72d64c8 100644 --- a/pkg/cli/plan_test.go +++ b/pkg/cli/plan_test.go @@ -37,6 +37,7 @@ func claudePlanSession(t *testing.T, home, project, id, slug, planPath, inlineBo func TestRunPlanClaudePrefersDiskContent(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "proj") require.NoError(t, os.MkdirAll(project, 0o755)) t.Chdir(project) @@ -59,6 +60,7 @@ func TestRunPlanClaudePrefersDiskContent(t *testing.T) { func TestRunPlanClaudeDefaultSourceUsesDirectSessionLookup(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "proj") require.NoError(t, os.MkdirAll(project, 0o755)) t.Chdir(project) @@ -76,6 +78,7 @@ func TestRunPlanClaudeDefaultSourceUsesDirectSessionLookup(t *testing.T) { func TestRunPlanClaudeInlineWhenMissingOnDisk(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "proj") require.NoError(t, os.MkdirAll(project, 0o755)) t.Chdir(project) @@ -93,6 +96,7 @@ func TestRunPlanClaudeInlineWhenMissingOnDisk(t *testing.T) { func TestRunPlanClaudeNoPlanErrors(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "proj") require.NoError(t, os.MkdirAll(project, 0o755)) t.Chdir(project) @@ -118,6 +122,7 @@ func TestRunPlanClaudeNoPlanErrors(t *testing.T) { func TestRunPlanCodexChecklist(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "proj") require.NoError(t, os.MkdirAll(project, 0o755)) t.Chdir(project) diff --git a/pkg/cli/project_options_ginkgo_test.go b/pkg/cli/project_options_ginkgo_test.go new file mode 100644 index 0000000..6b44910 --- /dev/null +++ b/pkg/cli/project_options_ginkgo_test.go @@ -0,0 +1,41 @@ +package cli + +import ( + "os" + "path/filepath" + "time" + + "github.com/flanksource/captain/pkg/database" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("project options", func() { + It("groups focused session aggregates by detected project root", func() { + root := filepath.Join(GinkgoT().TempDir(), "work", "captain") + cliDir := filepath.Join(root, "pkg", "cli") + databaseDir := filepath.Join(root, "pkg", "database") + Expect(os.MkdirAll(cliDir, 0o755)).To(Succeed()) + Expect(os.MkdirAll(databaseDir, 0o755)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(root, "go.mod"), []byte("module example\n"), 0o644)).To(Succeed()) + + earlier := time.Date(2026, 7, 16, 10, 0, 0, 0, time.UTC) + latest := earlier.Add(time.Hour) + result := projectOptionsFromAggregates([]database.ProjectSessionAggregate{ + {CWD: cliDir, Source: "codex", SessionCount: 2, LastActivityAt: &earlier}, + {CWD: databaseDir, Source: "claude", SessionCount: 3, ProcessActive: true, LastActivityAt: &latest}, + }) + + Expect(result).To(Equal(ProjectOptionsResult{ + Total: 1, + Projects: []ProjectOption{{ + Value: root, + Label: "work/captain", + Path: root, + Sources: []string{"claude", "codex", "live"}, + Sessions: 5, + LastUsed: &latest, + }}, + })) + }) +}) diff --git a/pkg/cli/projects.go b/pkg/cli/projects.go index 0d0d64a..f2c823e 100644 --- a/pkg/cli/projects.go +++ b/pkg/cli/projects.go @@ -1,14 +1,17 @@ package cli import ( + "context" "fmt" "os" "path/filepath" "regexp" "sort" + "strings" "time" "github.com/flanksource/captain/pkg/claude" + "github.com/flanksource/captain/pkg/database" "github.com/timberio/go-datemath" ) @@ -26,6 +29,27 @@ type ProjectsListResult struct { Rows []ProjectRow `json:"rows"` } +type ProjectOption struct { + Value string `json:"value"` + Label string `json:"label"` + Path string `json:"path"` + Sources []string `json:"sources,omitempty"` + Sessions int `json:"sessions,omitempty"` + LastUsed *time.Time `json:"lastUsed,omitempty"` +} + +type ProjectOptionsResult struct { + Total int `json:"total"` + Projects []ProjectOption `json:"projects"` +} + +type projectOptionAccumulator struct { + path string + sources map[string]bool + sessions int + lastUsed time.Time +} + func RunProjectsList(_ ProjectsListOptions) (any, error) { projects, err := scanProjects() if err != nil { @@ -49,6 +73,110 @@ func RunProjectsList(_ ProjectsListOptions) (any, error) { return ProjectsListResult{Total: len(projects), Rows: rows}, nil } +func RunProjectOptions(ctx context.Context) (ProjectOptionsResult, error) { + db, err := freshenSessionDB(ctx) + if err != nil { + return ProjectOptionsResult{}, err + } + aggregates, err := db.ListProjectSessionAggregates(ctx) + if err != nil { + return ProjectOptionsResult{}, err + } + return projectOptionsFromAggregates(aggregates), nil +} + +func projectOptionsFromAggregates(aggregates []database.ProjectSessionAggregate) ProjectOptionsResult { + accs := map[string]*projectOptionAccumulator{} + for _, aggregate := range aggregates { + if strings.TrimSpace(aggregate.CWD) == "" { + continue + } + lastUsed := time.Time{} + if aggregate.LastActivityAt != nil { + lastUsed = *aggregate.LastActivityAt + } + root := sessionProjectRoot(aggregate.CWD) + addProjectOption(accs, root, aggregate.Source, aggregate.SessionCount, lastUsed) + if aggregate.ProcessActive { + addProjectOption(accs, root, "live", 0, lastUsed) + } + } + + projectsOut := make([]ProjectOption, 0, len(accs)) + for _, acc := range accs { + sources := make([]string, 0, len(acc.sources)) + for source := range acc.sources { + sources = append(sources, source) + } + sort.Strings(sources) + var lastUsed *time.Time + if !acc.lastUsed.IsZero() { + value := acc.lastUsed + lastUsed = &value + } + projectsOut = append(projectsOut, ProjectOption{ + Value: acc.path, + Label: projectOptionLabel(acc.path), + Path: acc.path, + Sources: sources, + Sessions: acc.sessions, + LastUsed: lastUsed, + }) + } + + sort.Slice(projectsOut, func(i, j int) bool { + left, right := projectsOut[i], projectsOut[j] + if left.LastUsed != nil && right.LastUsed != nil && !left.LastUsed.Equal(*right.LastUsed) { + return left.LastUsed.After(*right.LastUsed) + } + if left.LastUsed != nil && right.LastUsed == nil { + return true + } + if left.LastUsed == nil && right.LastUsed != nil { + return false + } + return left.Label < right.Label + }) + + return ProjectOptionsResult{Total: len(projectsOut), Projects: projectsOut} +} + +func addProjectOption(accs map[string]*projectOptionAccumulator, path, source string, sessions int, lastUsed time.Time) { + path = normalizeSessionProject(path) + if path == "" { + return + } + acc := accs[path] + if acc == nil { + acc = &projectOptionAccumulator{path: path, sources: map[string]bool{}} + accs[path] = acc + } + if source != "" { + acc.sources[source] = true + } + acc.sessions += sessions + if lastUsed.After(acc.lastUsed) { + acc.lastUsed = lastUsed + } +} + +func projectOptionLabel(path string) string { + parts := strings.Split(filepath.ToSlash(path), "/") + filtered := make([]string, 0, len(parts)) + for _, part := range parts { + if part != "" { + filtered = append(filtered, part) + } + } + if len(filtered) >= 2 { + return strings.Join(filtered[len(filtered)-2:], "/") + } + if len(filtered) == 1 { + return filtered[0] + } + return path +} + type ProjectsCleanOptions struct { Since string `flag:"since" help:"Delete sessions not accessed within this period (e.g. 30d, now-30d)" default:"30d" short:"s"` DryRun bool `flag:"dry-run" help:"Preview without deleting" short:"n"` diff --git a/pkg/cli/projects_test.go b/pkg/cli/projects_test.go index 319e2fc..7b440d3 100644 --- a/pkg/cli/projects_test.go +++ b/pkg/cli/projects_test.go @@ -1,8 +1,13 @@ package cli import ( + "os" + "path/filepath" + "slices" "testing" "time" + + "github.com/flanksource/captain/pkg/database" ) func TestFormatBytes(t *testing.T) { @@ -46,6 +51,56 @@ func TestUuidStem(t *testing.T) { } } +func TestProjectOptionsFromAggregatesMergesClaudeCodexAndLive(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + claudeProject := filepath.Join(home, "work", "claude-project") + codexProject := filepath.Join(home, "work", "codex-project") + liveProject := filepath.Join(home, "work", "live-project") + for _, dir := range []string{claudeProject, codexProject, liveProject} { + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + markProjectRoot(t, dir) + } + + started := time.Date(2026, 6, 1, 11, 0, 0, 0, time.UTC) + aggregates := []database.ProjectSessionAggregate{ + {Source: "claude", CWD: claudeProject, SessionCount: 1, LastActivityAt: &started}, + {Source: "codex", CWD: codexProject, SessionCount: 1, LastActivityAt: &started}, + {Source: "codex", CWD: liveProject, SessionCount: 1, LastActivityAt: &started, ProcessActive: true}, + } + + result := projectOptionsFromAggregates(aggregates) + if result.Total != 3 { + t.Fatalf("projects = %+v", result) + } + assertProjectOption(t, result.Projects, claudeProject, "claude") + assertProjectOption(t, result.Projects, codexProject, "codex") + assertProjectOption(t, result.Projects, liveProject, "live") + assertProjectOption(t, result.Projects, liveProject, "codex") +} + +func assertProjectOption(t *testing.T, projects []ProjectOption, path, source string) { + t.Helper() + for _, project := range projects { + if project.Value != path { + continue + } + if project.Path != path { + t.Fatalf("project path = %q, want %q", project.Path, path) + } + if project.Label == "" { + t.Fatalf("empty label for %+v", project) + } + if !slices.Contains(project.Sources, source) { + t.Fatalf("sources for %q = %v, want %q", path, project.Sources, source) + } + return + } + t.Fatalf("project %q not found in %+v", path, projects) +} + func TestShellQuote(t *testing.T) { tests := []struct { input string diff --git a/pkg/cli/prompt_batch_run.go b/pkg/cli/prompt_batch_run.go new file mode 100644 index 0000000..6fedab8 --- /dev/null +++ b/pkg/cli/prompt_batch_run.go @@ -0,0 +1,86 @@ +package cli + +import ( + "context" + "errors" + "fmt" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/clicky/task" + flanksourceContext "github.com/flanksource/commons/context" +) + +func launchAsyncBatch(ctx context.Context, id string, rendered PromptRenderResult, runtimes []api.Model, chat bool) (PromptRunResult, error) { + if rendered.Input.SessionID != "" { + return PromptRunResult{}, errors.New("a multi-model run cannot resume one provider session") + } + batch, err := createPromptBatchSessions(ctx, rendered, runtimes) + if err != nil { + return PromptRunResult{}, err + } + updatePromptSessionLifecycle(ctx, batch.ID, database.SessionLifecycleRunning, "") + group := task.StartGroup[PromptRunSummary]( + "prompt "+rendered.Name, + task.WithGroupID(batch.ID.String()), + task.WithKind("prompt"), + task.WithLabels(promptTaskLabelsWithID(rendered, id, "multi")), + task.WithConcurrency(len(batch.Runs)), + ) + handles := make([]task.TypedTask[PromptRunSummary], len(batch.Runs)) + result := PromptRunResult{ + BatchID: batch.ID.String(), Status: "running", Chat: chat, + Total: len(batch.Runs), Runs: make([]PromptRunItem, len(batch.Runs)), + } + for i := range batch.Runs { + i := i + run := batch.Runs[i] + binding := promptBinding(batch, i) + variant := renderVariant(rendered, run.Runtime, nil) + runID := run.SessionID.String() + stream := promptRuns.create(runID) + capabilities := chatCapabilitiesForBackend(variant.Backend) + stream.setRun(PromptRunFrame{ + RunID: runID, SessionID: runID, Status: "running", Chat: chat, + Model: variant.Model, Backend: variant.Backend, Capabilities: capabilities, + }) + updatePromptSessionLifecycle(ctx, run.SessionID, database.SessionLifecycleRunning, "") + result.Runs[i] = PromptRunItem{ + RunID: runID, SessionID: runID, Selector: runtimeSelector(run.Runtime), + Status: "running", Model: run.Runtime.Name, Backend: string(run.Runtime.Backend), + Effort: string(run.Runtime.Effort), Chat: chat, Capabilities: capabilities, + } + handles[i] = group.Add(runtimeSelector(run.Runtime), func(_ flanksourceContext.Context, t *task.Task) (PromptRunSummary, error) { + if chat { + chatSession := newChatSession(runID, variant, runtimeTimeout(variant.Input.Budget.Timeout), stream, binding) + promptChats.register(chatSession) + return chatSession.run(t) + } + summary, runErr := runPromptStream(t, variant, runtimeTimeout(variant.Input.Budget.Timeout), runID, stream, binding) + if runErr != nil { + persistPromptRun(context.WithoutCancel(t.Context()), promptRunRecordInput{ + Rendered: variant, RunID: runID, Binding: binding, + Model: run.Runtime.Name, Backend: string(run.Runtime.Backend), Error: runErr.Error(), + }) + } + return summary, runErr + }, task.WithModel(run.Runtime.Name), task.WithPrompt(rendered.Input.Prompt.User)) + } + + task.StartTask("finalize prompt batch", func(_ flanksourceContext.Context, t *task.Task) (PromptRunSummary, error) { + succeeded, failed := 0, 0 + for i := range handles { + summary, runErr := handles[i].GetResult() + if runErr != nil || !summary.Success { + failed++ + } else { + succeeded++ + } + } + reason := fmt.Sprintf("%d succeeded, %d failed", succeeded, failed) + updatePromptSessionLifecycle(context.WithoutCancel(t.Context()), batch.ID, batchLifecycle(succeeded, failed), reason) + t.Success() + return PromptRunSummary{RunID: batch.ID.String(), Success: failed == 0}, nil + }, task.WithIdentity("prompt-batch-finalize-"+batch.ID.String())) + return result, nil +} diff --git a/pkg/cli/prompt_batch_session.go b/pkg/cli/prompt_batch_session.go new file mode 100644 index 0000000..f01bfc9 --- /dev/null +++ b/pkg/cli/prompt_batch_session.go @@ -0,0 +1,122 @@ +package cli + +import ( + "context" + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/google/uuid" +) + +type promptSessionBinding struct { + BatchID uuid.UUID + SessionID uuid.UUID +} + +type promptBatchRun struct { + SessionID uuid.UUID + Runtime api.Model +} + +type promptBatchSession struct { + ID uuid.UUID + Runs []promptBatchRun +} + +func createPromptBatchSessions(ctx context.Context, rendered PromptRenderResult, runtimes []api.Model) (promptBatchSession, error) { + if err := validatePromptRuntimes(runtimes); err != nil { + return promptBatchSession{}, err + } + db, err := captainDB(ctx) + if err != nil { + return promptBatchSession{}, err + } + batch := promptBatchSession{ID: uuid.New(), Runs: make([]promptBatchRun, len(runtimes))} + children := make([]database.CreateSessionInput, len(runtimes)) + for i, runtime := range runtimes { + source := backendSource(runtime.Backend) + if source == "" { + source = "captain" + } + batch.Runs[i] = promptBatchRun{SessionID: uuid.New(), Runtime: runtime} + children[i] = database.CreateSessionInput{ + ID: batch.Runs[i].SessionID, Source: source, + Provider: ai.BackendToProvider(runtime.Backend), HostID: captainHostID(), + CWD: rendered.Input.Cwd(), Title: runtime.Name, InitialPrompt: rendered.Input.Prompt.User, + AgentType: "model", Description: runtimeSelector(runtime), + } + } + _, err = db.CreateSessionTree(ctx, database.CreateSessionTreeInput{ + Root: database.CreateSessionInput{ + ID: batch.ID, Source: "captain", Provider: "multi-model", HostID: captainHostID(), + CWD: rendered.Input.Cwd(), Title: rendered.Name, InitialPrompt: rendered.Input.Prompt.User, + AgentType: "batch", Description: fmt.Sprintf("%d model comparison", len(runtimes)), + }, + Children: children, + }) + if err != nil { + return promptBatchSession{}, err + } + return batch, nil +} + +func validatePromptRuntimes(runtimes []api.Model) error { + if len(runtimes) < 2 { + return fmt.Errorf("multi-model execution requires at least two runtimes") + } + seen := map[string]struct{}{} + for i := range runtimes { + runtime := runtimes[i] + if err := runtime.Validate(); err != nil { + return fmt.Errorf("runtime %d: %w", i+1, err) + } + backend, err := runtime.ResolveBackend() + if err != nil { + return fmt.Errorf("runtime %d: %w", i+1, err) + } + runtimes[i].Backend = backend + key := runtimeSelector(runtimes[i]) + if _, ok := seen[key]; ok { + return fmt.Errorf("runtime %d duplicates %s", i+1, key) + } + seen[key] = struct{}{} + } + return nil +} + +func runtimeSelector(runtime api.Model) string { + parts := []string{string(runtime.Backend), runtime.Name} + if runtime.Effort != api.EffortNone { + parts = append(parts, string(runtime.Effort)) + } + return strings.Join(parts, ":") +} + +func promptBinding(batch promptBatchSession, index int) *promptSessionBinding { + return &promptSessionBinding{BatchID: batch.ID, SessionID: batch.Runs[index].SessionID} +} + +func updatePromptSessionLifecycle(ctx context.Context, id uuid.UUID, lifecycle database.SessionLifecycleStatus, reason string) { + db, err := captainDB(ctx) + if err != nil { + log.Errorf("open database for session %s lifecycle: %v", id, err) + return + } + if _, err := db.UpdateSessionLifecycle(ctx, id, lifecycle, reason); err != nil { + log.Errorf("update session %s lifecycle to %s: %v", id, lifecycle, err) + } +} + +func batchLifecycle(succeeded, failed int) database.SessionLifecycleStatus { + switch { + case failed == 0: + return database.SessionLifecycleSucceeded + case succeeded == 0: + return database.SessionLifecycleFailed + default: + return database.SessionLifecyclePartial + } +} diff --git a/pkg/cli/prompt_batch_session_ginkgo_test.go b/pkg/cli/prompt_batch_session_ginkgo_test.go new file mode 100644 index 0000000..52fe56a --- /dev/null +++ b/pkg/cli/prompt_batch_session_ginkgo_test.go @@ -0,0 +1,84 @@ +package cli + +import ( + "os" + "path/filepath" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + commonsdb "github.com/flanksource/commons-db/db" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("prompt batch sessions", func() { + It("creates the batch ID as the canonical root with one child per runtime", func() { + if os.Getenv("CAPTAIN_DB_EMBEDDED_TEST") == "" { + Skip("set CAPTAIN_DB_EMBEDDED_TEST=1 to run embedded-postgres cli tests") + } + dsn, stop, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ + DataDir: filepath.Join(GinkgoT().TempDir(), "postgres"), Database: "captain_prompt_batch", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { Expect(stop()).To(Succeed()) }) + db, err := database.Open(GinkgoT().Context(), database.WithDSN(dsn), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { + setCaptainDBForTest(nil) + Expect(db.Close()).To(Succeed()) + }) + setCaptainDBForTest(db) + + rendered := PromptRenderResult{Name: "compare", Backend: "codex-cmux", Model: "gpt-5.6-sol"} + rendered.Input.Prompt.User = "Compare these approaches" + rendered.Input.SetCwd("/workspace/captain") + batch, err := createPromptBatchSessions(GinkgoT().Context(), rendered, []api.Model{ + {Name: "gpt-5.6-sol", Backend: api.BackendCodexCmux, Effort: api.EffortHigh}, + {Name: "gemini-2.5-flash", Backend: api.BackendGemini}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(batch.ID).NotTo(Equal(uuid.Nil)) + Expect(batch.Runs).To(HaveLen(2)) + + root, err := db.GetSession(GinkgoT().Context(), batch.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(root.ParentSessionID).To(BeNil()) + Expect(root.RootSessionID).To(BeNil()) + Expect(root.Source).To(Equal("captain")) + Expect(root.Provider).To(Equal("multi-model")) + Expect(root.AgentType).To(Equal("batch")) + Expect(root.InitialPrompt).To(Equal("Compare these approaches")) + + thread, err := db.ListThreadSessionOverviews(GinkgoT().Context(), batch.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(thread).To(HaveLen(3)) + updatedRoot, err := db.UpdateSessionLifecycle( + GinkgoT().Context(), + batch.ID, + database.SessionLifecyclePartial, + "1 succeeded, 1 failed", + ) + Expect(err).NotTo(HaveOccurred()) + Expect(updatedRoot.LifecycleStatus).To(Equal(database.SessionLifecyclePartial)) + Expect(updatedRoot.StateVersion).To(Equal(int64(1))) + Expect(updatedRoot.EndedAt).NotTo(BeNil()) + result, err := RunSessionGet(GinkgoT().Context(), SessionGetOptions{ID: batch.ID.String()}) + Expect(err).NotTo(HaveOccurred()) + Expect(result.RootSessionID).To(Equal(batch.ID.String())) + Expect(result.Sessions).To(HaveLen(3)) + Expect(result.Tree().GetChildren()).To(HaveLen(1)) + Expect(result.Tree().GetChildren()[0].GetChildren()).To(HaveLen(2)) + for i, run := range batch.Runs { + Expect(run.SessionID).NotTo(Equal(uuid.Nil)) + child, err := db.GetSession(GinkgoT().Context(), run.SessionID) + Expect(err).NotTo(HaveOccurred()) + Expect(child.ParentSessionID).NotTo(BeNil()) + Expect(*child.ParentSessionID).To(Equal(batch.ID)) + Expect(child.RootSessionID).NotTo(BeNil()) + Expect(*child.RootSessionID).To(Equal(batch.ID)) + Expect(child.ProviderSessionID).To(BeEmpty()) + Expect(child.Description).To(Equal(runtimeSelector(batch.Runs[i].Runtime))) + } + }) +}) diff --git a/pkg/cli/prompt_chat.go b/pkg/cli/prompt_chat.go new file mode 100644 index 0000000..0f16ae1 --- /dev/null +++ b/pkg/cli/prompt_chat.go @@ -0,0 +1,496 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/session" + "github.com/flanksource/clicky/task" + "github.com/google/uuid" +) + +const chatIdleTimeout = 10 * time.Minute + +type chatSession struct { + mu sync.Mutex + runID string + rendered PromptRenderResult + timeout time.Duration + stream *runStream + wake chan struct{} + + provider ai.Provider + streamer ai.StreamingProvider + baseCancel context.CancelFunc + turnCancel context.CancelFunc + turnDone chan struct{} + acc *promptEventAccumulator + + state ChatStateFrame + queue []ChatQueuedMessage + discarded []string + interruptedTurn int + terminal bool + startedAt time.Time + binding *promptSessionBinding +} + +func newChatSession(runID string, rendered PromptRenderResult, timeout time.Duration, stream *runStream, binding *promptSessionBinding) *chatSession { + capabilities := chatCapabilitiesForBackend(rendered.Backend) + chat := &chatSession{ + runID: runID, rendered: rendered, timeout: timeout, stream: stream, binding: binding, + wake: make(chan struct{}, 1), startedAt: time.Now(), + state: ChatStateFrame{ + RunID: runID, Status: "starting", Capabilities: capabilities, + }, + } + stream.setChatState(chat.state) + return chat +} + +func (c *chatSession) run(t *task.Task) (PromptRunSummary, error) { + baseCtx, cancel := context.WithCancel(t.Context()) + c.mu.Lock() + c.baseCancel = cancel + c.mu.Unlock() + c.stream.setCancel(cancel) + defer cancel() + + req := c.rendered.Input + if err := preparePromptAttachments(baseCtx, &req, c.rendered.Config); err != nil { + return c.fail(t, err) + } + provider, cleanup, err := buildProvider(baseCtx, &req, c.rendered.Config) + if err != nil { + return c.fail(t, err) + } + defer cleanup() + defer closeProvider(provider) + streamer, ok := provider.(ai.StreamingProvider) + if !ok { + return c.fail(t, fmt.Errorf("backend %s does not support streaming", c.rendered.Backend)) + } + c.mu.Lock() + c.provider = provider + c.streamer = streamer + c.acc = newPromptEventAccumulator(c.stream.publish, t, c.rendered.Model, c.rendered.Backend) + c.acc.cwd = req.Cwd() + c.acc.idPrefix = c.runID + c.mu.Unlock() + + next := req + for { + summary, interrupted, turnErr := c.runTurn(baseCtx, t, next) + if turnErr != nil { + if c.stream.wasStopped() { + turnErr = errors.New("stopped") + } + return c.fail(t, turnErr) + } + if c.stream.wasStopped() { + return c.fail(t, errors.New("stopped")) + } + if !interrupted { + c.persistTurn(next, summary) + } + if !summary.Success && !interrupted { + return c.fail(t, errors.New(summary.Error)) + } + nextMessage, waitErr := c.waitForMessage(baseCtx, summary) + if errors.Is(waitErr, errChatIdle) { + return c.complete(t, summary), nil + } + if waitErr != nil { + if c.stream.wasStopped() || errors.Is(waitErr, context.Canceled) { + return c.fail(t, errors.New("stopped")) + } + return c.fail(t, waitErr) + } + next = c.followUpRequest(req, nextMessage.Text) + } +} + +var errChatIdle = errors.New("chat idle timeout") + +func (c *chatSession) runTurn(baseCtx context.Context, t *task.Task, req ai.Request) (PromptRunSummary, bool, error) { + turnCtx, cancel := runContext(baseCtx, req, c.timeout) + turnDone := make(chan struct{}) + c.mu.Lock() + c.state.Turn++ + turn := c.state.Turn + c.state.Status = "starting" + c.turnCancel = cancel + c.turnDone = turnDone + state := c.stateCopyLocked() + streamer := c.streamer + c.mu.Unlock() + c.stream.setChatState(state) + defer func() { + cancel() + close(turnDone) + c.mu.Lock() + c.turnCancel = nil + c.turnDone = nil + c.mu.Unlock() + }() + + events, err := streamer.ExecuteStream(turnCtx, req) + if err != nil { + return PromptRunSummary{}, false, err + } + started := false + var eventErr string + for event := range events { + if !started { + started = true + c.markRunning(event.SessionID) + } + if event.SessionID != "" { + c.rememberSession(event.SessionID) + } + if event.Kind == ai.EventError { + eventErr = event.Error + } + c.acc.handle(turn, event) + } + sessionID, model, usage, cost := c.acc.snapshot() + if sessionID != "" { + c.rememberSession(sessionID) + } + c.mu.Lock() + interrupted := c.interruptedTurn == turn + if interrupted { + c.interruptedTurn = 0 + } + c.mu.Unlock() + if eventErr != "" && !interrupted { + return PromptRunSummary{}, false, errors.New(eventErr) + } + summary := PromptRunSummary{ + RunID: c.runID, SessionID: sessionID, Model: model, + Backend: string(c.provider.GetBackend()), InputTokens: usage.InputTokens, + OutputTokens: usage.OutputTokens, CostUSD: cost, + Duration: time.Since(c.startedAt).Round(time.Millisecond).String(), Success: true, + } + return summary, interrupted, nil +} + +func (c *chatSession) waitForMessage(ctx context.Context, summary PromptRunSummary) (ChatQueuedMessage, error) { + c.mu.Lock() + if len(c.queue) > 0 { + message := c.popQueueLocked() + state := c.stateCopyLocked() + c.mu.Unlock() + c.stream.setChatState(state) + return message, nil + } + if !c.state.Capabilities.FollowUp { + c.mu.Unlock() + return ChatQueuedMessage{}, errChatIdle + } + c.state.Status = "idle" + c.state.Summary = &summary + state := c.stateCopyLocked() + c.mu.Unlock() + c.stream.setChatState(state) + + timer := time.NewTimer(chatIdleTimeout) + defer timer.Stop() + select { + case <-ctx.Done(): + return ChatQueuedMessage{}, ctx.Err() + case <-timer.C: + return ChatQueuedMessage{}, errChatIdle + case <-c.wake: + c.mu.Lock() + if len(c.queue) == 0 { + c.mu.Unlock() + return ChatQueuedMessage{}, fmt.Errorf("chat wake without a queued message") + } + message := c.popQueueLocked() + state := c.stateCopyLocked() + c.mu.Unlock() + c.stream.setChatState(state) + return message, nil + } +} + +func (c *chatSession) send(ctx context.Context, request ChatMessageRequest) (ChatMessageResponse, error) { + text := strings.TrimSpace(request.Text) + if text == "" { + return ChatMessageResponse{}, newChatError(http.StatusBadRequest, "message text is required") + } + messageID := strings.TrimSpace(request.MessageID) + if messageID == "" { + messageID = uuid.NewString() + } + message := ChatQueuedMessage{MessageID: messageID, Text: text} + + c.mu.Lock() + if c.terminal { + c.mu.Unlock() + return ChatMessageResponse{}, newChatError(http.StatusConflict, "run is terminal") + } + status := c.state.Status + capabilities := c.state.Capabilities + provider := c.provider + if status == "starting" || status == "interrupting" || status == "stopping" { + c.mu.Unlock() + return ChatMessageResponse{}, newChatError(http.StatusConflict, "run is not ready for a message") + } + if status == "running" && capabilities.Steer { + c.mu.Unlock() + steerer, ok := api.ProviderAs[api.SteerableProvider](provider) + if !ok { + return ChatMessageResponse{}, newChatError(http.StatusConflict, "active provider cannot be steered") + } + req := c.followUpRequest(c.rendered.Input, text) + if err := steerer.Steer(ctx, req); err != nil { + return ChatMessageResponse{}, err + } + c.publishUser(message) + return ChatMessageResponse{RunID: c.runID, MessageID: messageID, Status: "steered", Capabilities: capabilities}, nil + } + if !capabilities.FollowUp { + c.mu.Unlock() + return ChatMessageResponse{}, newChatError(http.StatusConflict, "active provider cannot accept follow-up messages") + } + c.queue = append(c.queue, message) + c.state.Queued = append([]ChatQueuedMessage(nil), c.queue...) + state := c.stateCopyLocked() + c.mu.Unlock() + c.publishUser(message) + c.stream.setChatState(state) + if status == "idle" { + select { + case c.wake <- struct{}{}: + default: + } + } + responseStatus := "queued" + if status == "idle" { + responseStatus = "started" + } + return ChatMessageResponse{RunID: c.runID, MessageID: messageID, Status: responseStatus, Capabilities: capabilities}, nil +} + +func (c *chatSession) interrupt(ctx context.Context) (ChatInterruptResponse, error) { + c.mu.Lock() + if c.terminal || c.state.Status != "running" || !c.state.Capabilities.Interrupt { + c.mu.Unlock() + return ChatInterruptResponse{}, newChatError(http.StatusConflict, "run cannot be interrupted") + } + provider := c.provider + turn := c.state.Turn + queued := append([]ChatQueuedMessage(nil), c.queue...) + c.state.Status = "interrupting" + state := c.stateCopyLocked() + c.mu.Unlock() + c.stream.setChatState(state) + + interruptible, ok := api.ProviderAs[api.InterruptibleProvider](provider) + if !ok { + c.restoreRunning() + return ChatInterruptResponse{}, newChatError(http.StatusConflict, "active provider cannot be interrupted") + } + if err := interruptible.Interrupt(ctx); err != nil { + c.restoreRunning() + return ChatInterruptResponse{}, err + } + + c.mu.Lock() + c.interruptedTurn = turn + c.queue = nil + discarded := make([]string, 0, len(queued)) + for _, message := range queued { + discarded = append(discarded, message.MessageID) + } + c.discarded = append(c.discarded, discarded...) + c.state.Queued = nil + c.state.DiscardedMessageIDs = append([]string(nil), c.discarded...) + state = c.stateCopyLocked() + turnCancel := c.turnCancel + turnDone := c.turnDone + c.mu.Unlock() + c.stream.setChatState(state) + go cancelTurnBackstop(turnDone, turnCancel) + return ChatInterruptResponse{Status: "interrupting", DiscardedMessageIDs: discarded}, nil +} + +func cancelTurnBackstop(done <-chan struct{}, cancel context.CancelFunc) { + if cancel == nil { + return + } + timer := time.NewTimer(5 * time.Second) + defer timer.Stop() + select { + case <-done: + case <-timer.C: + cancel() + } +} + +func (c *chatSession) stop() bool { + c.mu.Lock() + if c.terminal { + c.mu.Unlock() + return false + } + c.state.Status = "stopping" + state := c.stateCopyLocked() + cancel := c.baseCancel + c.mu.Unlock() + c.stream.setChatState(state) + c.stream.requestStop() + if cancel != nil { + cancel() + } + return true +} + +func (c *chatSession) markRunning(sessionID string) { + c.mu.Lock() + c.state.Status = "running" + c.state.Capabilities = chatCapabilitiesForBackend(string(c.provider.GetBackend())) + if sessionID != "" { + c.state.SessionID = sessionID + } + state := c.stateCopyLocked() + c.mu.Unlock() + c.stream.setChatState(state) +} + +func (c *chatSession) rememberSession(sessionID string) { + c.mu.Lock() + c.state.SessionID = sessionID + state := c.stateCopyLocked() + c.mu.Unlock() + c.stream.setChatState(state) + promptChats.bindSession(c, sessionID) +} + +func (c *chatSession) persistTurn(req ai.Request, summary PromptRunSummary) { + turn := c.state.Turn + runID := c.runID + if turn > 1 { + runID = fmt.Sprintf("%s-turn-%d", c.runID, turn) + } + rendered := c.rendered + rendered.Input = req + persistPromptRun(context.Background(), promptRunRecordInput{ + Rendered: rendered, RunID: runID, SessionID: summary.SessionID, + Binding: c.binding, Model: summary.Model, Backend: summary.Backend, + }) +} + +func (c *chatSession) complete(t *task.Task, summary PromptRunSummary) PromptRunSummary { + c.mu.Lock() + c.terminal = true + c.state.Summary = &summary + c.mu.Unlock() + promptChats.finish(c) + c.stream.complete(summary) + t.Success() + return summary +} + +func (c *chatSession) fail(t *task.Task, err error) (PromptRunSummary, error) { + c.mu.Lock() + c.terminal = true + c.mu.Unlock() + promptChats.finish(c) + if c.binding != nil { + persistPromptRun(context.Background(), promptRunRecordInput{ + Rendered: c.rendered, RunID: c.runID, Binding: c.binding, + Model: c.rendered.Model, Backend: c.rendered.Backend, Error: err.Error(), + }) + } + c.stream.fail(err.Error()) + _, _ = t.FailedWithError(err) + return PromptRunSummary{RunID: c.runID, Error: err.Error()}, err +} + +func (c *chatSession) followUpRequest(base ai.Request, text string) ai.Request { + req := base + req.Prompt.User = text + req.Prompt.Attachments = nil + req.Prompt.Schema = nil + req.Prompt.SchemaJSON = nil + req.Workflow = nil + req.Setup = nil + return req +} + +func (c *chatSession) publishUser(message ChatQueuedMessage) { + c.stream.publish(session.Message{ + ID: message.MessageID, Role: "user", + Parts: []session.Part{{Type: session.PartText, Text: message.Text}}, + }) +} + +func (c *chatSession) popQueueLocked() ChatQueuedMessage { + message := c.queue[0] + c.queue = c.queue[1:] + c.state.Queued = append([]ChatQueuedMessage(nil), c.queue...) + return message +} + +func (c *chatSession) stateCopyLocked() ChatStateFrame { + state := c.state + state.Queued = append([]ChatQueuedMessage(nil), c.state.Queued...) + state.DiscardedMessageIDs = append([]string(nil), c.state.DiscardedMessageIDs...) + return state +} + +func (c *chatSession) restoreRunning() { + c.mu.Lock() + c.state.Status = "running" + state := c.stateCopyLocked() + c.mu.Unlock() + c.stream.setChatState(state) +} + +func (c *chatSession) terminalState() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.terminal +} + +func (c *chatSession) sessionID() string { + c.mu.Lock() + defer c.mu.Unlock() + return c.state.SessionID +} + +func (c *chatSession) projection() (string, ChatCapabilities, *ChatStateFrame) { + c.mu.Lock() + defer c.mu.Unlock() + state := c.stateCopyLocked() + return c.runID, state.Capabilities, &state +} + +func closeProvider(provider ai.Provider) { + if closer, ok := api.ProviderAs[api.CloseableProvider](provider); ok { + if err := closer.Close(); err != nil { + log.Errorf("close chat provider: %v", err) + } + } +} + +type chatHTTPError struct { + status int + msg string +} + +func (e chatHTTPError) Error() string { return e.msg } + +func newChatError(status int, message string) error { + return chatHTTPError{status: status, msg: message} +} diff --git a/pkg/cli/prompt_chat_broker.go b/pkg/cli/prompt_chat_broker.go new file mode 100644 index 0000000..93235fd --- /dev/null +++ b/pkg/cli/prompt_chat_broker.go @@ -0,0 +1,69 @@ +package cli + +import "sync" + +type chatBroker struct { + mu sync.Mutex + byRun map[string]*chatSession + bySession map[string]*chatSession +} + +func newChatBroker() *chatBroker { + return &chatBroker{byRun: map[string]*chatSession{}, bySession: map[string]*chatSession{}} +} + +var promptChats = newChatBroker() + +func (b *chatBroker) register(chat *chatSession) { + b.mu.Lock() + b.byRun[chat.runID] = chat + b.mu.Unlock() +} + +func (b *chatBroker) bindSession(chat *chatSession, sessionID string) { + if sessionID == "" { + return + } + b.mu.Lock() + b.bySession[sessionID] = chat + b.mu.Unlock() +} + +func (b *chatBroker) getRun(runID string) (*chatSession, bool) { + b.mu.Lock() + defer b.mu.Unlock() + chat, ok := b.byRun[runID] + return chat, ok +} + +func (b *chatBroker) getSession(sessionID string) (*chatSession, bool) { + b.mu.Lock() + defer b.mu.Unlock() + chat, ok := b.bySession[sessionID] + if ok && chat.terminalState() { + delete(b.bySession, sessionID) + return nil, false + } + return chat, ok +} + +func (b *chatBroker) finish(chat *chatSession) { + sessionID := chat.sessionID() + b.mu.Lock() + defer b.mu.Unlock() + if b.bySession[sessionID] == chat { + delete(b.bySession, sessionID) + } +} + +func (b *chatBroker) stopAll() { + b.mu.Lock() + chats := make([]*chatSession, 0, len(b.byRun)) + for _, chat := range b.byRun { + chats = append(chats, chat) + } + b.mu.Unlock() + for _, chat := range chats { + chat.stop() + } +} diff --git a/pkg/cli/prompt_chat_ginkgo_test.go b/pkg/cli/prompt_chat_ginkgo_test.go new file mode 100644 index 0000000..bac0282 --- /dev/null +++ b/pkg/cli/prompt_chat_ginkgo_test.go @@ -0,0 +1,104 @@ +package cli + +import ( + "context" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/clicky/rpc" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type interruptProviderStub struct { + interrupted bool +} + +func (p *interruptProviderStub) Execute(context.Context, api.Spec) (*api.Response, error) { + return &api.Response{}, nil +} +func (p *interruptProviderStub) GetModel() string { return "test-model" } +func (p *interruptProviderStub) GetBackend() api.Backend { return api.BackendCodexAgent } +func (p *interruptProviderStub) Interrupt(context.Context) error { + p.interrupted = true + return nil +} + +var _ = Describe("prompt chat lifecycle", func() { + It("queues an idle follow-up and publishes its exact message id", func() { + stream := newRunStream() + chat := newChatSession("run-1", PromptRenderResult{Backend: string(api.BackendCodexAgent)}, 0, stream, nil) + chat.state.Status = "idle" + + response, err := chat.send(context.Background(), ChatMessageRequest{Text: "next", MessageID: "message-1"}) + + Expect(err).NotTo(HaveOccurred()) + Expect(response.Status).To(Equal("started")) + Expect(response.MessageID).To(Equal("message-1")) + entries, _, _, _ := stream.snapshot() + Expect(entries).To(HaveLen(1)) + Expect(entries[0].ID).To(Equal("message-1")) + Expect(chat.queue).To(Equal([]ChatQueuedMessage{{MessageID: "message-1", Text: "next"}})) + }) + + It("interrupts the active turn and discards queued messages", func() { + stream := newRunStream() + provider := &interruptProviderStub{} + chat := newChatSession("run-1", PromptRenderResult{Backend: string(api.BackendCodexAgent)}, 0, stream, nil) + chat.provider = provider + chat.state.Status = "running" + chat.state.Turn = 2 + chat.turnDone = make(chan struct{}) + chat.queue = []ChatQueuedMessage{{MessageID: "queued-1", Text: "later"}} + + response, err := chat.interrupt(context.Background()) + close(chat.turnDone) + + Expect(err).NotTo(HaveOccurred()) + Expect(provider.interrupted).To(BeTrue()) + Expect(response.DiscardedMessageIDs).To(Equal([]string{"queued-1"})) + Expect(chat.state.Status).To(Equal("interrupting")) + Expect(chat.state.Queued).To(BeEmpty()) + Expect(chat.state.DiscardedMessageIDs).To(Equal([]string{"queued-1"})) + }) + + It("latches stop requests made before a cancel function is attached", func() { + stream := newRunStream() + Expect(stream.requestStop()).To(BeTrue()) + cancelled := false + stream.setCancel(func() { cancelled = true }) + Expect(cancelled).To(BeTrue()) + }) + + It("publishes every prompt chat control route in OpenAPI", func() { + spec := &rpc.OpenAPISpec{} + addCaptainPromptRunPaths(spec) + + Expect(spec.Paths).To(HaveKey("/api/captain/prompt/runs/{runId}")) + Expect(spec.Paths).To(HaveKey("/api/captain/prompt/runs/{runId}/stream")) + Expect(spec.Paths).To(HaveKey("/api/captain/prompt/runs/{runId}/message")) + Expect(spec.Paths).To(HaveKey("/api/captain/prompt/runs/{runId}/interrupt")) + Expect(spec.Paths).To(HaveKey("/api/captain/prompt/runs/{runId}/stop")) + Expect(spec.Paths).To(HaveKey("/api/captain/sessions/{id}/message")) + }) + + It("collapses duplicate records for the same provider session when resuming", func() { + rich := SessionGetItem{ + CaptainID: "captain-rich", ProviderSessionID: "provider-1", + Summary: SessionRecord{Model: "gpt-5.6-sol", Backend: "codex-agent", CWD: "/repo"}, + } + selected, err := selectResumeSession([]SessionGetItem{ + {CaptainID: "captain-sparse", ProviderSessionID: "provider-1"}, rich, + }) + + Expect(err).NotTo(HaveOccurred()) + Expect(selected).To(Equal(rich)) + }) + + It("rejects resume prefixes that identify different provider sessions", func() { + _, err := selectResumeSession([]SessionGetItem{ + {ProviderSessionID: "provider-1"}, {ProviderSessionID: "provider-2"}, + }) + + Expect(err).To(MatchError("session id is ambiguous")) + }) +}) diff --git a/pkg/cli/prompt_chat_http.go b/pkg/cli/prompt_chat_http.go new file mode 100644 index 0000000..2296c92 --- /dev/null +++ b/pkg/cli/prompt_chat_http.go @@ -0,0 +1,246 @@ +package cli + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + "os" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" +) + +func handlePromptRunMessage(chats *chatBroker) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + chat, ok := chats.getRun(r.PathValue("runId")) + if !ok { + http.Error(w, "unknown run", http.StatusNotFound) + return + } + var request ChatMessageRequest + if !decodeChatRequest(w, r, &request) { + return + } + response, err := chat.send(r.Context(), request) + if err != nil { + writeChatError(w, err) + return + } + writeChatJSON(w, http.StatusAccepted, response) + } +} + +func handlePromptRunInterrupt(chats *chatBroker) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + chat, ok := chats.getRun(r.PathValue("runId")) + if !ok { + http.Error(w, "unknown run", http.StatusNotFound) + return + } + response, err := chat.interrupt(r.Context()) + if err != nil { + writeChatError(w, err) + return + } + writeChatJSON(w, http.StatusAccepted, response) + } +} + +func handlePromptRunStop(runs *runBroker, chats *chatBroker) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + runID := r.PathValue("runId") + stream, ok := runs.get(runID) + if !ok { + http.Error(w, "unknown run", http.StatusNotFound) + return + } + if chat, found := chats.getRun(runID); found { + if !chat.stop() { + writeChatError(w, newChatError(http.StatusConflict, "run is terminal")) + return + } + } else if !stream.requestStop() { + writeChatError(w, newChatError(http.StatusConflict, "run is terminal")) + return + } + writeChatJSON(w, http.StatusAccepted, map[string]string{"status": "stopping"}) + } +} + +func handleSessionMessage(chats *chatBroker) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + var request ChatMessageRequest + if !decodeChatRequest(w, r, &request) { + return + } + result, err := RunSessionGet(r.Context(), SessionGetOptions{ID: r.PathValue("id")}) + if err != nil { + writeChatError(w, err) + return + } + if result.Total == 0 { + http.Error(w, "unknown session", http.StatusNotFound) + return + } + item, selectErr := selectResumeSession(result.Sessions) + if selectErr != nil { + writeChatError(w, selectErr) + return + } + if chat, ok := chats.getSession(item.ProviderSessionID); ok { + response, sendErr := chat.send(r.Context(), request) + if sendErr != nil { + writeChatError(w, sendErr) + return + } + writeChatJSON(w, http.StatusAccepted, response) + return + } + response, resumeErr := resumeSessionMessage(item, request) + if resumeErr != nil { + writeChatError(w, resumeErr) + return + } + writeChatJSON(w, http.StatusAccepted, response) + } +} + +func selectResumeSession(items []SessionGetItem) (SessionGetItem, error) { + if len(items) == 0 { + return SessionGetItem{}, newChatError(http.StatusNotFound, "unknown session") + } + providerID := strings.TrimSpace(items[0].ProviderSessionID) + selected := items[0] + for _, item := range items[1:] { + if providerID == "" || strings.TrimSpace(item.ProviderSessionID) != providerID { + return SessionGetItem{}, newChatError(http.StatusConflict, "session id is ambiguous") + } + if resumeSessionScore(item) > resumeSessionScore(selected) { + selected = item + } + } + return selected, nil +} + +func resumeSessionScore(item SessionGetItem) int { + score := 0 + if strings.TrimSpace(item.Summary.Model) != "" { + score += 4 + } + if strings.TrimSpace(item.Summary.Backend) != "" { + score += 2 + } + if strings.TrimSpace(item.Summary.CWD) != "" { + score += 2 + } + if item.Detail != nil { + score++ + } + return score +} + +func resumeSessionMessage(item SessionGetItem, request ChatMessageRequest) (ChatMessageResponse, error) { + if strings.TrimSpace(request.Text) == "" { + return ChatMessageResponse{}, newChatError(http.StatusBadRequest, "message text is required") + } + if strings.TrimSpace(item.ProviderSessionID) == "" { + return ChatMessageResponse{}, newChatError(http.StatusUnprocessableEntity, "session has no provider session id") + } + if strings.TrimSpace(item.Summary.CWD) == "" { + return ChatMessageResponse{}, newChatError(http.StatusUnprocessableEntity, "session has no working directory") + } + info, err := os.Stat(item.Summary.CWD) + if err != nil || !info.IsDir() { + return ChatMessageResponse{}, newChatError(http.StatusUnprocessableEntity, "session working directory is unavailable") + } + backend, err := resumeBackend(item.Summary.Source) + if err != nil { + return ChatMessageResponse{}, err + } + if request.Backend != "" && api.Backend(request.Backend) != backend { + return ChatMessageResponse{}, newChatError(http.StatusUnprocessableEntity, + fmt.Sprintf("session source %s must resume with backend %s", item.Summary.Source, backend)) + } + model := strings.TrimSpace(request.Model) + if model == "" { + model = strings.TrimSpace(item.Summary.Model) + } + if model == "" { + return ChatMessageResponse{}, newChatError(http.StatusUnprocessableEntity, "session has no model") + } + messageID := strings.TrimSpace(request.MessageID) + if messageID == "" { + messageID = newMessageID() + } + modelSpec := api.Model{ + Name: model, Backend: backend, Effort: api.Effort(item.Summary.ReasoningEffort), + } + req := api.Spec{ + Model: modelSpec, Prompt: api.Prompt{User: strings.TrimSpace(request.Text)}, + SessionID: item.ProviderSessionID, + } + req.SetCwd(item.Summary.CWD) + rendered := PromptRenderResult{ + Name: "resume " + item.CaptainID, Model: model, Backend: string(backend), + User: req.Prompt.User, Input: req, + Config: ai.Config{Model: modelSpec, SessionID: item.ProviderSessionID}, + } + run := launchAsyncRun(item.CaptainID, rendered, true) + if stream, ok := promptRuns.get(run.RunID); ok { + stream.publish(userChatMessage(messageID, req.Prompt.User)) + } + return ChatMessageResponse{ + RunID: run.RunID, MessageID: messageID, Status: "started", Capabilities: run.Capabilities, + }, nil +} + +func resumeBackend(source string) (api.Backend, error) { + switch strings.ToLower(strings.TrimSpace(source)) { + case "claude": + return api.BackendClaudeAgent, nil + case "codex": + return api.BackendCodexAgent, nil + default: + return "", newChatError(http.StatusUnprocessableEntity, "session source is not resumable") + } +} + +func userChatMessage(messageID, text string) session.Message { + return session.Message{ + ID: messageID, Role: "user", + Parts: []session.Part{{Type: session.PartText, Text: text}}, + } +} + +func newMessageID() string { return uuid.NewString() } + +func decodeChatRequest(w http.ResponseWriter, r *http.Request, target any) bool { + decoder := json.NewDecoder(r.Body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(target); err != nil { + http.Error(w, "invalid request: "+err.Error(), http.StatusBadRequest) + return false + } + return true +} + +func writeChatError(w http.ResponseWriter, err error) { + var target chatHTTPError + if errors.As(err, &target) { + http.Error(w, target.Error(), target.status) + return + } + http.Error(w, err.Error(), http.StatusInternalServerError) +} + +func writeChatJSON(w http.ResponseWriter, status int, value any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(value); err != nil { + log.Errorf("write chat response: %v", err) + } +} diff --git a/pkg/cli/prompt_chat_types.go b/pkg/cli/prompt_chat_types.go new file mode 100644 index 0000000..4f39167 --- /dev/null +++ b/pkg/cli/prompt_chat_types.go @@ -0,0 +1,68 @@ +package cli + +import "github.com/flanksource/captain/pkg/api" + +type ChatCapabilities struct { + Interrupt bool `json:"interrupt"` + Steer bool `json:"steer"` + FollowUp bool `json:"followUp"` + Resume bool `json:"resume"` +} + +type ChatQueuedMessage struct { + MessageID string `json:"messageId"` + Text string `json:"text"` +} + +type ChatStateFrame struct { + RunID string `json:"runId"` + SessionID string `json:"sessionId,omitempty"` + Status string `json:"status"` + Turn int `json:"turn"` + Capabilities ChatCapabilities `json:"capabilities"` + Queued []ChatQueuedMessage `json:"queued,omitempty"` + DiscardedMessageIDs []string `json:"discardedMessageIds,omitempty"` + Summary *PromptRunSummary `json:"summary,omitempty"` +} + +type PromptRunFrame struct { + RunID string `json:"runId"` + SessionID string `json:"sessionId,omitempty"` + Status string `json:"status"` + Chat bool `json:"chat"` + Model string `json:"model,omitempty"` + Backend string `json:"backend,omitempty"` + Capabilities ChatCapabilities `json:"capabilities"` +} + +type ChatMessageRequest struct { + Text string `json:"text"` + MessageID string `json:"messageId,omitempty"` + Model string `json:"model,omitempty"` + Backend string `json:"backend,omitempty"` +} + +type ChatMessageResponse struct { + RunID string `json:"runId"` + MessageID string `json:"messageId"` + Status string `json:"status"` + Capabilities ChatCapabilities `json:"capabilities"` +} + +type ChatInterruptResponse struct { + Status string `json:"status"` + DiscardedMessageIDs []string `json:"discardedMessageIds,omitempty"` +} + +func chatCapabilitiesForBackend(backend string) ChatCapabilities { + switch api.Backend(backend) { + case api.BackendClaudeAgent: + return ChatCapabilities{Interrupt: true, Steer: true, FollowUp: true, Resume: true} + case api.BackendCodexAgent: + return ChatCapabilities{Interrupt: true, FollowUp: true, Resume: true} + case api.BackendClaudeCLI, api.BackendCodexCLI, api.BackendClaudeCmux, api.BackendCodexCmux: + return ChatCapabilities{Resume: true} + default: + return ChatCapabilities{} + } +} diff --git a/pkg/cli/prompt_entity.go b/pkg/cli/prompt_entity.go index b712b04..135629c 100644 --- a/pkg/cli/prompt_entity.go +++ b/pkg/cli/prompt_entity.go @@ -2,32 +2,21 @@ package cli import ( "context" - "crypto/sha256" - "encoding/base64" - "encoding/hex" - "encoding/json" "errors" "fmt" - "io" "io/fs" "net/http" "os" "path/filepath" "sort" - "strconv" "strings" "sync" - "time" "github.com/flanksource/captain/pkg/ai" - promptlib "github.com/flanksource/captain/pkg/ai/prompt" "github.com/flanksource/captain/pkg/api" - "github.com/flanksource/captain/pkg/captainconfig" - "github.com/flanksource/captain/pkg/collections" "github.com/flanksource/clicky" clickyapi "github.com/flanksource/clicky/api" clickyrpc "github.com/flanksource/clicky/rpc" - dp "github.com/google/dotprompt/go/dotprompt" ) type promptDirsContextKey struct{} @@ -58,6 +47,7 @@ type PromptSummary struct { Writable bool `json:"writable"` Model string `json:"model,omitempty"` Backend string `json:"backend,omitempty"` + Runtimes []api.Model `json:"runtimes,omitempty"` Variables []PromptVariable `json:"variables,omitempty"` ParseError string `json:"parseError,omitempty"` UpdatedAt string `json:"updatedAt,omitempty"` @@ -105,6 +95,8 @@ type PromptWriteRequest struct { type PromptRenderRequest struct { Variables map[string]any `json:"variables,omitempty"` Spec *api.Spec `json:"spec,omitempty"` + Runtimes []api.Model `json:"runtimes,omitempty"` + Chat bool `json:"chat,omitempty"` } type PromptRenderResult struct { @@ -119,13 +111,15 @@ type PromptRenderResult struct { InputSchema map[string]any `json:"inputSchema,omitempty"` InputDefault map[string]any `json:"inputDefault,omitempty"` OutputSchema map[string]any `json:"outputSchema,omitempty"` + Runtimes []api.Model `json:"runtimes,omitempty"` ValidationError string `json:"validationError,omitempty"` } // PromptActionFlags is the full flag surface for `captain prompt run|render` — // the same knobs as `captain ai prompt` (AIRuntimeOptions) plus the prompt-body -// fields — so the two commands are one. The positional (a .prompt filepath or a -// registry id) is the prompt source; --prompt/-p and stdin are alternatives. +// fields — so the two commands are one. The positional (a discovered name, +// .prompt filepath, or registry id) is the prompt source; --prompt/-p and stdin +// are alternatives. type PromptActionFlags struct { AIRuntimeOptions @@ -133,7 +127,9 @@ type PromptActionFlags struct { System string `flag:"system" help:"System prompt" short:"s"` AppendSystem string `flag:"append-system" help:"Append text to the default system prompt"` Var []string `flag:"var" help:"Template variable key=value (repeatable)" short:"V"` + Attach []string `flag:"attach" help:"Attach a local path or URL (repeatable; RFC 4180 comma-separated values allowed)" short:"A"` Vars string `flag:"vars" help:"JSON object of template variables (HTTP callers)"` + MultiModels []string `flag:"multi-models" help:"Run prompt once per runtime selector in parallel, e.g. cli:sonnet-5,cmux:opus (repeatable; comma-separated allowed)" short:"M"` Timeout string `flag:"timeout" help:"Request timeout" default:"120s"` NoStream bool `flag:"no-stream" help:"Disable streaming; print only the final text (CLI)"` } @@ -169,6 +165,7 @@ type promptInspection struct { InputSchema map[string]any InputDefault map[string]any OutputSchema map[string]any + Runtimes []api.Model Variables []PromptVariable } @@ -183,10 +180,10 @@ func RegisterPromptEntity() { UpdateWithContext(updatePrompt). DeleteWithContext(deletePrompt). WithAction(clicky.ActionWithFlagsAndContext("render", PromptActionFlags{}, renderPromptAction). - WithShort("Render a prompt (id, .prompt file, --prompt/-p, or stdin) without calling a model"). + WithShort("Render a prompt (id, name, .prompt file, --prompt/-p, or stdin) without calling a model"). WithOptionalID()). WithAction(clicky.ActionWithFlagsAndContext("run", PromptActionFlags{}, runPromptAction). - WithShort("Run a prompt (id, .prompt file, --prompt/-p, or stdin)"). + WithShort("Run a prompt (id, name, .prompt file, --prompt/-p, or stdin)"). WithOptionalID()). Register() }) @@ -359,1019 +356,3 @@ func renderPromptAction(ctx context.Context, id string, flags map[string]string) } return renderPromptCLI(ctx, id, opts, flags["vars"], readStdinIfCLI(ctx)) } - -// renderPrompt is the HTTP/Spec render path: overlay a structured api.Spec (the -// web UI's rich runtime overrides) onto the rendered template. -func renderPrompt(ctx context.Context, id string, renderReq PromptRenderRequest) (PromptRenderResult, error) { - record, err := resolvePromptRecord(ctx, id) - if err != nil { - return PromptRenderResult{}, err - } - content, err := readPromptContent(record) - if err != nil { - return PromptRenderResult{}, err - } - vars := renderReq.Variables - if vars == nil { - vars = map[string]any{} - } - req, cfg, err := promptlib.Load(content).Render(vars, nil) - if err != nil { - return PromptRenderResult{}, err - } - req.Prompt.Source = record.Rel - if renderReq.Spec != nil { - overlayRuntimeSpec(&req, &cfg, *renderReq.Spec) - } - applyPromptDefaults(&req, &cfg) - cwd, err := os.Getwd() - if err != nil { - return PromptRenderResult{}, fmt.Errorf("get working directory: %w", err) - } - if err := normalizePromptContextDir(&req, cwd); err != nil { - return PromptRenderResult{}, err - } - return finalizeRenderResult(record, content, req, cfg) -} - -// renderPromptCLI is the CLI render path: load from id | .prompt filepath | -p | -// stdin and overlay the flat CLI flags (overlayCLI). -func renderPromptCLI(ctx context.Context, id string, opts AIPromptOptions, varsJSON, stdin string) (PromptRenderResult, error) { - content, source, usedStdin, record, err := loadPromptContent(ctx, id, opts, stdin) - if err != nil { - return PromptRenderResult{}, err - } - vars, err := promptVars(opts, varsJSON, stdin, usedStdin) - if err != nil { - return PromptRenderResult{}, err - } - req, cfg, err := renderLoadedContent(content, source, vars, opts) - if err != nil { - return PromptRenderResult{}, err - } - return finalizeRenderResult(record, content, req, cfg) -} - -// finalizeRenderResult packages the rendered request/config + prompt detail into -// a PromptRenderResult and sets the validation error (shared by both paths). -func finalizeRenderResult(record promptRecord, content string, req ai.Request, cfg ai.Config) (PromptRenderResult, error) { - // Normalize a comma-separated model into a clean primary + fallbacks so the - // displayed Model is a single name, then catch a mistyped model (primary or any - // fallback) at render time, not just on run. - req.Model = req.ExpandCSV() - cfg.Model = cfg.Model.ExpandCSV() - for _, c := range cfg.Model.Candidates() { - warnIfLikelyModelTypo(c.Name) - } - detail, err := promptDetailFromContent(record, content) - if err != nil { - return PromptRenderResult{}, err - } - result := PromptRenderResult{ - ID: detail.ID, - Name: detail.Name, - Model: cfg.Model.Name, - Backend: string(cfg.Model.Backend), - User: req.Prompt.User, - System: req.Prompt.System, - Input: req, - Config: cfg, - InputSchema: detail.InputSchema, - InputDefault: detail.InputDefault, - OutputSchema: detail.OutputSchema, - } - switch { - case req.Prompt.User == "" && !req.IsVerifyOnly(): - result.ValidationError = "prompt text required" - case cfg.Model.Name == "": - result.ValidationError = "no model: set prompt frontmatter, pass a model override, or run 'captain configure'" - default: - if err := req.Validate(); err != nil { - result.ValidationError = err.Error() - } - } - return result, nil -} - -func overlayRuntimeSpec(req *ai.Request, cfg *ai.Config, spec api.Spec) { - if spec.Name != "" { - req.Name = spec.Name - cfg.Model.Name = spec.Name - } - if spec.ID != "" { - req.ID = spec.ID - cfg.Model.ID = spec.ID - } - if spec.Backend != "" { - req.Backend = spec.Backend - cfg.Model.Backend = spec.Backend - } - if spec.Temperature != nil { - req.Temperature = spec.Temperature - cfg.Model.Temperature = spec.Temperature - } - if spec.Effort != "" { - req.Effort = spec.Effort - cfg.Model.Effort = spec.Effort - } - if len(spec.Fallbacks) > 0 { - req.Fallbacks = spec.Fallbacks - cfg.Model.Fallbacks = spec.Fallbacks - } - req.NoCache = req.NoCache || spec.NoCache - cfg.NoCache = cfg.NoCache || spec.NoCache - if spec.Budget.Cost > 0 { - req.Budget.Cost = spec.Budget.Cost - cfg.Budget.Cost = spec.Budget.Cost - } - if spec.Budget.MaxTokens > 0 { - req.Budget.MaxTokens = spec.Budget.MaxTokens - cfg.Budget.MaxTokens = spec.Budget.MaxTokens - } - if spec.Budget.MaxTurns > 0 { - req.Budget.MaxTurns = spec.Budget.MaxTurns - cfg.Budget.MaxTurns = spec.Budget.MaxTurns - } - if spec.Budget.Timeout != "" { - req.Budget.Timeout = spec.Budget.Timeout - cfg.Budget.Timeout = spec.Budget.Timeout - } - - if spec.Prompt.User != "" { - req.Prompt.User = spec.Prompt.User - } - if spec.Prompt.System != "" { - req.Prompt.System = spec.Prompt.System - } - if spec.Prompt.AppendSystem != "" { - req.Prompt.AppendSystem = spec.Prompt.AppendSystem - } - if spec.Prompt.Source != "" { - req.Prompt.Source = spec.Prompt.Source - } - req.Prompt.Metadata = mergeStringMaps(req.Prompt.Metadata, spec.Prompt.Metadata) - if len(spec.Prompt.SchemaJSON) > 0 { - req.Prompt.SchemaJSON = spec.Prompt.SchemaJSON - } - if spec.Prompt.SchemaStrictness != "" { - req.Prompt.SchemaStrictness = spec.Prompt.SchemaStrictness - } - if spec.Workflow != nil { - req.Workflow = spec.Workflow - } - - if spec.Permissions.Mode != "" { - req.Permissions.Mode = spec.Permissions.Mode - } - req.Permissions.Presets = mergePresets(req.Permissions.Presets, spec.Permissions.Presets) - toolPolicies := spec.Permissions.Tools.Policies() - if len(toolPolicies) > 0 { - req.Permissions.Tools.Allow = nil - req.Permissions.Tools.Deny = nil - req.Permissions.Tools.Modes = nil - for _, tool := range sortedStringKeys(toolPolicies) { - switch toolPolicies[tool] { - case api.ToolPolicyAllow: - req.Permissions.Tools.Allow = append(req.Permissions.Tools.Allow, tool) - case api.ToolPolicyDeny: - req.Permissions.Tools.Deny = append(req.Permissions.Tools.Deny, tool) - case api.ToolPolicyAsk: - if req.Permissions.Tools.Modes == nil { - req.Permissions.Tools.Modes = map[string]api.ToolMode{} - } - req.Permissions.Tools.Modes[tool] = api.ToolModeAsk - case api.ToolPolicyAuto: - if req.Permissions.Tools.Modes == nil { - req.Permissions.Tools.Modes = map[string]api.ToolMode{} - } - req.Permissions.Tools.Modes[tool] = api.ToolModeEnabled - } - } - } - req.Permissions.Tools.Modes = mergeToolModes(req.Permissions.Tools.Modes, spec.Permissions.Tools.Modes) - req.Permissions.MCP.Disabled = req.Permissions.MCP.Disabled || spec.Permissions.MCP.Disabled - if servers := spec.Permissions.MCP.EnabledServers(); len(servers) > 0 { - req.Permissions.MCP.Servers = servers - } - if len(spec.Permissions.Plugins) > 0 { - req.Permissions.Plugins = enabledResourcePolicies(spec.Permissions.Plugins) - } - - skills := append([]string(nil), spec.Memory.Skills...) - skills = append(skills, spec.Permissions.Skills.Enabled()...) - if len(skills) > 0 { - req.Memory.Skills = dedupeStrings(skills) - } - req.Memory.SkipProject = req.Memory.SkipProject || spec.Memory.SkipProject - req.Memory.SkipUser = req.Memory.SkipUser || spec.Memory.SkipUser - req.Memory.SkipSkills = req.Memory.SkipSkills || spec.Memory.SkipSkills - req.Memory.SkipHooks = req.Memory.SkipHooks || spec.Memory.SkipHooks - req.Memory.SkipMemory = req.Memory.SkipMemory || spec.Memory.SkipMemory - req.Memory.Bare = req.Memory.Bare || spec.Memory.Bare - - if spec.Setup != nil { - req.Setup = spec.Setup - } - - if spec.SessionID != "" { - req.SessionID = spec.SessionID - cfg.SessionID = spec.SessionID - } -} - -func applyPromptDefaults(req *ai.Request, cfg *ai.Config) { - saved := loadSavedAI() - if req.Name == "" { - req.Name = firstNonEmpty(cfg.Model.Name, saved.Model) - } - if req.Backend == "" { - req.Backend = api.Backend(firstNonEmpty(string(cfg.Model.Backend), saved.Backend)) - } - if req.Effort == "" { - req.Effort = api.Effort(firstNonEmpty(string(cfg.Model.Effort), saved.ReasoningEffort)) - } - req.NoCache = req.NoCache || saved.NoCache - if req.Budget.MaxTokens == 0 { - req.Budget.MaxTokens = firstPositive(cfg.Budget.MaxTokens, saved.MaxTokens, 4096) - } - if req.Budget.Cost == 0 { - req.Budget.Cost = firstPositiveFloat(cfg.Budget.Cost, saved.BudgetUSD) - } - - cfg.Model = req.Model - cfg.Budget = req.Budget - cfg.NoCache = req.NoCache -} - -func mergeStringMaps(base, overlay map[string]string) map[string]string { - if len(overlay) == 0 { - return base - } - out := make(map[string]string, collections.SafeAdd(len(base), len(overlay))) - for k, v := range base { - out[k] = v - } - for k, v := range overlay { - out[k] = v - } - return out -} - -func mergeToolModes(base, overlay map[string]api.ToolMode) map[string]api.ToolMode { - if len(overlay) == 0 { - return base - } - out := make(map[string]api.ToolMode, collections.SafeAdd(len(base), len(overlay))) - for k, v := range base { - out[k] = v - } - for k, v := range overlay { - out[k] = v - } - return out -} - -func mergePresets(base, overlay []api.Preset) []api.Preset { - if len(overlay) == 0 { - return base - } - seen := make(map[api.Preset]bool, collections.SafeAdd(len(base), len(overlay))) - out := make([]api.Preset, 0, collections.SafeAdd(len(base), len(overlay))) - for _, preset := range base { - if seen[preset] { - continue - } - seen[preset] = true - out = append(out, preset) - } - for _, preset := range overlay { - if seen[preset] { - continue - } - seen[preset] = true - out = append(out, preset) - } - return out -} - -func sortedStringKeys[V any](m map[string]V) []string { - keys := make([]string, 0, len(m)) - for key := range m { - if key != "" { - keys = append(keys, key) - } - } - sort.Strings(keys) - return keys -} - -func enabledResourcePolicies(in api.ResourcePolicies) api.ResourcePolicies { - out := api.ResourcePolicies{} - for _, key := range sortedStringKeys(in) { - if in[key] == api.ResourceEnabled { - out[key] = api.ResourceEnabled - } - } - return out -} - -func dedupeStrings(in []string) []string { - seen := map[string]bool{} - var out []string - for _, item := range in { - if item == "" || seen[item] { - continue - } - seen[item] = true - out = append(out, item) - } - return out -} - -func readRenderRequest(ctx context.Context, flags map[string]string) (PromptRenderRequest, error) { - var req PromptRenderRequest - if err := decodePromptBody(ctx, map[string]any{}, &req); err != nil { - return PromptRenderRequest{}, err - } - if req.Variables == nil { - req.Variables = map[string]any{} - } - if err := mergePromptActionFlags(&req, flags); err != nil { - return PromptRenderRequest{}, err - } - return req, nil -} - -func mergePromptActionFlags(req *PromptRenderRequest, flags map[string]string) error { - if len(flags) == 0 { - return nil - } - if raw := strings.TrimSpace(flags["vars"]); raw != "" { - var vars map[string]any - if err := json.Unmarshal([]byte(raw), &vars); err != nil { - return fmt.Errorf("parse --vars JSON: %w", err) - } - req.Variables = vars - } - if v := strings.TrimSpace(flags["model"]); v != "" { - ensureRenderSpec(req).Name = v - } - if v := strings.TrimSpace(flags["fallback"]); v != "" { - ensureRenderSpec(req).Fallbacks = fallbackModelsFromFlags([]string{v}) - } - if v := strings.TrimSpace(flags["backend"]); v != "" { - ensureRenderSpec(req).Backend = api.Backend(v) - } - if v := strings.TrimSpace(flags["timeout"]); v != "" { - ensureRenderSpec(req).Budget.Timeout = v - } - if v := strings.TrimSpace(flags["max-tokens"]); v != "" { - n, err := strconv.Atoi(v) - if err != nil { - return fmt.Errorf("invalid --max-tokens %q: %w", v, err) - } - ensureRenderSpec(req).Budget.MaxTokens = n - } - return nil -} - -func ensureRenderSpec(req *PromptRenderRequest) *api.Spec { - if req.Spec == nil { - req.Spec = &api.Spec{} - } - return req.Spec -} - -func decodePromptBody(ctx context.Context, flat map[string]any, dst any) error { - if r, ok := clickyrpc.RequestFromContext(ctx); ok && r.Body != nil { - body, err := io.ReadAll(r.Body) - if err != nil { - return fmt.Errorf("read request body: %w", err) - } - r.Body = io.NopCloser(strings.NewReader(string(body))) - if len(strings.TrimSpace(string(body))) > 0 { - decoder := json.NewDecoder(strings.NewReader(string(body))) - decoder.DisallowUnknownFields() - if err := decoder.Decode(dst); err != nil { - return fmt.Errorf("decode request body: %w", err) - } - return nil - } - } - if len(flat) == 0 { - return nil - } - data, err := json.Marshal(flat) - if err != nil { - return err - } - decoder := json.NewDecoder(strings.NewReader(string(data))) - decoder.DisallowUnknownFields() - if err := decoder.Decode(dst); err != nil { - return fmt.Errorf("decode command body: %w", err) - } - return nil -} - -func runtimeTimeout(raw string) time.Duration { - timeout, _ := time.ParseDuration(raw) - if timeout <= 0 { - return 120 * time.Second - } - return timeout -} - -func listPromptRecords(ctx context.Context) ([]promptRecord, error) { - sources, err := buildPromptSources(ctx) - if err != nil { - return nil, err - } - var records []promptRecord - for _, source := range sources { - recs, err := listPromptRecordsFromSource(source) - if err != nil { - return nil, err - } - records = append(records, recs...) - } - return records, nil -} - -func listPromptRecordsFromSource(source promptSource) ([]promptRecord, error) { - var records []promptRecord - add := func(rel string, info fs.FileInfo) { - rel = filepath.ToSlash(rel) - path := rel - if source.Root != "" { - path = filepath.Join(source.Root, filepath.FromSlash(rel)) - } - updatedAt := "" - if info != nil && !info.ModTime().IsZero() { - updatedAt = info.ModTime().Format(time.RFC3339) - } - records = append(records, promptRecord{ - Source: source, - ID: encodePromptID(source.Kind, source.ID, rel), - Path: path + "\x00" + updatedAt, - Rel: rel, - }) - } - - if source.FS != nil { - err := fs.WalkDir(source.FS, source.WalkRoot, func(path string, d fs.DirEntry, err error) error { - if err != nil { - return err - } - if d.IsDir() { - return nil - } - if !strings.HasSuffix(path, ".prompt") { - return nil - } - info, _ := d.Info() - add(path, info) - return nil - }) - return records, err - } - - err := filepath.WalkDir(source.Root, func(path string, d fs.DirEntry, err error) error { - if err != nil { - return err - } - name := d.Name() - if d.IsDir() { - if name == ".git" || name == "node_modules" || name == "vendor" || name == "dist" { - return filepath.SkipDir - } - return nil - } - if d.Type()&fs.ModeSymlink != 0 || !strings.HasSuffix(name, ".prompt") { - return nil - } - rel, err := filepath.Rel(source.Root, path) - if err != nil { - return err - } - info, _ := d.Info() - add(rel, info) - return nil - }) - if errors.Is(err, fs.ErrNotExist) && source.Implicit { - return records, nil - } - return records, err -} - -func resolvePromptRecord(ctx context.Context, id string) (promptRecord, error) { - if looksLikePromptPath(id) { - return filePromptRecord(id) - } - ref, err := decodePromptID(id) - if err != nil { - return promptRecord{}, err - } - sources, err := buildPromptSources(ctx) - if err != nil { - return promptRecord{}, err - } - for _, source := range sources { - if source.Kind != ref.Kind || source.ID != ref.SourceID { - continue - } - path := ref.RelPath - if source.Root != "" { - path = filepath.Join(source.Root, filepath.FromSlash(ref.RelPath)) - } - return promptRecord{Source: source, ID: id, Path: path, Rel: ref.RelPath}, nil - } - return promptRecord{}, fmt.Errorf("prompt source %q not found", ref.SourceID) -} - -// looksLikePromptPath reports whether id is a filesystem path rather than a -// base64 registry id. Registry ids are base64-raw-url (no ".", "/", or leading -// "."), so a .prompt suffix, a path separator, or a leading "." marks a path. -func looksLikePromptPath(id string) bool { - return strings.HasSuffix(id, ".prompt") || - strings.ContainsRune(id, os.PathSeparator) || - strings.HasPrefix(id, ".") -} - -// filePromptRecord resolves an ad-hoc .prompt file path (not a registered id) -// into a record readable via readPromptContent/safeLocalPromptPath. Mirrors the -// captain-ai-prompt file loader so `captain prompt run|render ./x.prompt` works. -func filePromptRecord(id string) (promptRecord, error) { - abs, err := filepath.Abs(id) - if err != nil { - return promptRecord{}, err - } - info, err := os.Stat(abs) - if err != nil { - return promptRecord{}, fmt.Errorf("prompt file %s: %w", id, err) - } - if info.IsDir() { - return promptRecord{}, fmt.Errorf("%s is a directory, not a .prompt file", id) - } - return promptRecord{ - Source: promptSource{Kind: "file", ID: "file", Label: "File", Root: filepath.Dir(abs)}, - ID: id, - Path: abs, - Rel: filepath.Base(abs), - }, nil -} - -func promptSummary(record promptRecord) (PromptSummary, error) { - content, err := readPromptContent(record) - if err != nil { - return PromptSummary{}, err - } - summary, err := promptSummaryFromContent(record, content) - if err != nil { - summary = basePromptSummary(record) - summary.ParseError = err.Error() - } - if idx := strings.LastIndex(record.Path, "\x00"); idx >= 0 { - summary.UpdatedAt = strings.TrimPrefix(record.Path[idx+1:], "\x00") - } - return summary, nil -} - -func promptDetail(record promptRecord) (PromptDetail, error) { - content, err := readPromptContent(record) - if err != nil { - return PromptDetail{}, err - } - return promptDetailFromContent(record, content) -} - -func promptDetailFromContent(record promptRecord, content string) (PromptDetail, error) { - summary, err := promptSummaryFromContent(record, content) - if err != nil { - return PromptDetail{}, err - } - inspection, err := inspectPrompt(content, nil) - if err != nil { - return PromptDetail{}, err - } - return PromptDetail{ - PromptSummary: summary, - Content: content, - InputSchema: inspection.InputSchema, - InputDefault: inspection.InputDefault, - OutputSchema: inspection.OutputSchema, - Metadata: inspection.Metadata, - }, nil -} - -func promptSummaryFromContent(record promptRecord, content string) (PromptSummary, error) { - tmpl := promptlib.Load(content) - req, cfg, err := tmpl.Render(map[string]any{}, nil) - if err != nil { - return PromptSummary{}, err - } - inspection, err := inspectPrompt(content, nil) - if err != nil { - return PromptSummary{}, err - } - summary := basePromptSummary(record) - if v, ok := inspection.Metadata["name"].(string); ok && strings.TrimSpace(v) != "" { - summary.Name = strings.TrimSpace(v) - } - if v, ok := inspection.Metadata["description"].(string); ok { - summary.Description = strings.TrimSpace(v) - } - summary.Model = firstNonEmpty(cfg.Model.Name, req.Name) - summary.Backend = firstNonEmpty(string(cfg.Model.Backend), string(req.Backend)) - summary.Variables = inspection.Variables - return summary, nil -} - -func basePromptSummary(record promptRecord) PromptSummary { - name := strings.TrimSuffix(filepath.Base(record.Rel), ".prompt") - return PromptSummary{ - ID: record.ID, - Name: name, - SourceKind: record.Source.Kind, - SourceID: record.Source.ID, - Source: record.Source.Label, - Path: displayPromptPath(record), - RelPath: record.Rel, - Writable: record.Source.Writable, - } -} - -func displayPromptPath(record promptRecord) string { - if idx := strings.LastIndex(record.Path, "\x00"); idx >= 0 { - return record.Path[:idx] - } - return record.Path -} - -func readPromptContent(record promptRecord) (string, error) { - if record.Source.FS != nil { - data, err := fs.ReadFile(record.Source.FS, record.Rel) - if err != nil { - return "", fmt.Errorf("read embedded prompt %s: %w", record.Rel, err) - } - return string(data), nil - } - full, err := safeLocalPromptPath(record.Source, record.Rel) - if err != nil { - return "", err - } - data, err := os.ReadFile(full) - if err != nil { - return "", fmt.Errorf("read prompt %s: %w", full, err) - } - return string(data), nil -} - -func inspectPrompt(content string, data map[string]any) (promptInspection, error) { - if data == nil { - data = map[string]any{} - } - rendered, err := dp.NewDotprompt(nil).Render(content, &dp.DataArgument{Input: data}, nil) - if err != nil { - return promptInspection{}, err - } - metadata := map[string]any{} - if rendered.Raw != nil { - for k, v := range rendered.Raw { - metadata[k] = v - } - } - if rendered.Name != "" { - metadata["name"] = rendered.Name - } - if rendered.Description != "" { - metadata["description"] = rendered.Description - } - if rendered.Model != "" { - metadata["model"] = rendered.Model - } - inputSchema := anyToMap(rendered.Input.Schema) - inputDefault := map[string]any{} - for k, v := range rendered.Input.Default { - inputDefault[k] = v - } - return promptInspection{ - Metadata: metadata, - InputSchema: inputSchema, - InputDefault: inputDefault, - OutputSchema: anyToMap(rendered.Output.Schema), - Variables: variablesFromSchema(inputSchema), - }, nil -} - -func anyToMap(v any) map[string]any { - if v == nil { - return nil - } - data, err := json.Marshal(v) - if err != nil { - return nil - } - var out map[string]any - if err := json.Unmarshal(data, &out); err != nil { - return nil - } - return out -} - -func variablesFromSchema(schema map[string]any) []PromptVariable { - props, _ := schema["properties"].(map[string]any) - if len(props) == 0 { - return nil - } - required := map[string]bool{} - if raw, ok := schema["required"].([]any); ok { - for _, item := range raw { - if s, ok := item.(string); ok { - required[s] = true - } - } - } - keys := make([]string, 0, len(props)) - for k := range props { - keys = append(keys, k) - } - sort.Strings(keys) - var vars []PromptVariable - for _, name := range keys { - prop, _ := props[name].(map[string]any) - item := PromptVariable{Name: name, Required: required[name]} - if v, ok := prop["type"].(string); ok { - item.Type = v - } - if v, ok := prop["description"].(string); ok { - item.Description = v - } - vars = append(vars, item) - } - return vars -} - -func promptSourceMatches(summary PromptSummary, source string) bool { - switch source { - case "", "all": - return true - case "embedded": - return summary.SourceKind == "embedded" - case "local": - return summary.SourceKind == "local" - default: - return summary.SourceID == source || strings.EqualFold(summary.Source, source) - } -} - -func promptMatches(summary PromptSummary, filter string) bool { - if filter == "" { - return true - } - haystack := strings.ToLower(strings.Join([]string{ - summary.Name, - summary.Description, - summary.Source, - summary.Path, - summary.RelPath, - summary.Model, - summary.Backend, - }, "\n")) - return strings.Contains(haystack, filter) -} - -func buildPromptSources(ctx context.Context) ([]promptSource, error) { - sources := []promptSource{{ - Kind: "embedded", - ID: "embedded", - Label: "Embedded examples", - WalkRoot: "testdata", - FS: promptlib.Examples, - Writable: false, - }} - - seen := map[string]bool{} - addLocal := func(raw, base string) error { - dir, err := resolvePromptDir(raw, base) - if err != nil { - return err - } - if seen[dir] { - return nil - } - seen[dir] = true - sources = append(sources, promptSource{ - Kind: "local", - ID: hashPromptDir(dir), - Label: dir, - Root: dir, - Writable: true, - }) - return nil - } - - cfg, exists, err := captainconfig.Load() - if err != nil { - return nil, err - } - if exists { - configPath, err := captainconfig.Path() - if err != nil { - return nil, err - } - base := filepath.Dir(configPath) - for _, dir := range cfg.Prompts.Dirs { - if err := addLocal(dir, base); err != nil { - return nil, err - } - } - } - cwd, err := os.Getwd() - if err != nil { - return nil, err - } - for _, dir := range promptDirsFromContext(ctx) { - if err := addLocal(dir, cwd); err != nil { - return nil, err - } - } - if _, ok := firstWritableSource(sources); !ok { - dir := filepath.Join(cwd, ".captain", "prompts") - sources = append(sources, promptSource{ - Kind: "local", - ID: hashPromptDir(dir), - Label: dir, - Root: dir, - Writable: true, - Implicit: true, - }) - } - return sources, nil -} - -func promptDirsFromContext(ctx context.Context) []string { - if ctx == nil { - return nil - } - if dirs, ok := ctx.Value(promptDirsContextKey{}).([]string); ok { - return dirs - } - return nil -} - -func resolvePromptDir(raw, base string) (string, error) { - dir := strings.TrimSpace(raw) - if dir == "" { - return "", fmt.Errorf("prompt dir cannot be empty") - } - if strings.HasPrefix(dir, "~") { - home, err := os.UserHomeDir() - if err != nil { - return "", err - } - switch { - case dir == "~": - dir = home - case strings.HasPrefix(dir, "~/"): - dir = filepath.Join(home, dir[2:]) - default: - return "", fmt.Errorf("unsupported home-relative prompt dir %q", raw) - } - } - if !filepath.IsAbs(dir) { - dir = filepath.Join(base, dir) - } - abs, err := filepath.Abs(dir) - if err != nil { - return "", err - } - info, err := os.Stat(abs) - if err != nil { - return "", fmt.Errorf("prompt dir %s: %w", abs, err) - } - if !info.IsDir() { - return "", fmt.Errorf("prompt dir %s is not a directory", abs) - } - if resolved, err := filepath.EvalSymlinks(abs); err == nil { - abs = resolved - } - return filepath.Clean(abs), nil -} - -func writableSourceByID(sources []promptSource, id string) (promptSource, bool) { - for _, source := range sources { - if source.ID == id && source.Writable { - return source, true - } - } - return promptSource{}, false -} - -func firstWritableSource(sources []promptSource) (promptSource, bool) { - for _, source := range sources { - if source.Writable { - return source, true - } - } - return promptSource{}, false -} - -func safeLocalPromptPath(source promptSource, rel string) (string, error) { - cleanRel := strings.TrimPrefix(filepath.Clean(filepath.FromSlash(rel)), string(filepath.Separator)) - if cleanRel == "." || cleanRel == "" || filepath.IsAbs(rel) || strings.HasPrefix(cleanRel, ".."+string(filepath.Separator)) || cleanRel == ".." { - return "", fmt.Errorf("invalid prompt path %q", rel) - } - if filepath.Ext(cleanRel) != ".prompt" { - return "", fmt.Errorf("prompt path must end with .prompt") - } - full := filepath.Join(source.Root, cleanRel) - abs, err := filepath.Abs(full) - if err != nil { - return "", err - } - relToRoot, err := filepath.Rel(source.Root, abs) - if err != nil { - return "", err - } - if strings.HasPrefix(relToRoot, ".."+string(filepath.Separator)) || relToRoot == ".." || filepath.IsAbs(relToRoot) { - return "", fmt.Errorf("prompt path escapes source root") - } - if info, err := os.Lstat(abs); err == nil && info.Mode()&fs.ModeSymlink != 0 { - return "", fmt.Errorf("prompt symlinks are not supported") - } - return abs, nil -} - -func normalizeWriteRelPath(relPath, name string) (string, error) { - rel := strings.TrimSpace(relPath) - if rel == "" { - rel = slugPromptName(name) - } - if rel == "" { - return "", fmt.Errorf("prompt name or path required") - } - if !strings.HasSuffix(rel, ".prompt") { - rel += ".prompt" - } - rel = filepath.ToSlash(filepath.Clean(filepath.FromSlash(rel))) - if strings.HasPrefix(rel, "../") || rel == ".." || strings.HasPrefix(rel, "/") { - return "", fmt.Errorf("invalid prompt path %q", relPath) - } - return rel, nil -} - -func slugPromptName(name string) string { - name = strings.ToLower(strings.TrimSpace(name)) - var b strings.Builder - lastDash := false - for _, r := range name { - ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') - if ok { - b.WriteRune(r) - lastDash = false - continue - } - if !lastDash { - b.WriteRune('-') - lastDash = true - } - } - return strings.Trim(b.String(), "-") -} - -func encodePromptID(kind, sourceID, rel string) string { - return base64.RawURLEncoding.EncodeToString([]byte(kind + "\x00" + sourceID + "\x00" + filepath.ToSlash(rel))) -} - -func decodePromptID(id string) (promptRef, error) { - data, err := base64.RawURLEncoding.DecodeString(id) - if err != nil { - return promptRef{}, fmt.Errorf("invalid prompt id: %w", err) - } - parts := strings.SplitN(string(data), "\x00", 3) - if len(parts) != 3 || parts[0] == "" || parts[1] == "" || parts[2] == "" { - return promptRef{}, fmt.Errorf("invalid prompt id") - } - return promptRef{Kind: parts[0], SourceID: parts[1], RelPath: filepath.ToSlash(parts[2])}, nil -} - -func hashPromptDir(dir string) string { - sum := sha256.Sum256([]byte(dir)) - return hex.EncodeToString(sum[:])[:12] -} - -func ValidatePromptDirs(dirs []string) error { - cwd, err := os.Getwd() - if err != nil { - return err - } - for _, dir := range dirs { - if _, err := resolvePromptDir(dir, cwd); err != nil { - return err - } - } - return nil -} - -var _ clicky.EntityItem = PromptSummary{} -var _ clickyapi.TableProvider = PromptSummary{} diff --git a/pkg/cli/prompt_entity_test.go b/pkg/cli/prompt_entity_test.go index b576fce..601e09f 100644 --- a/pkg/cli/prompt_entity_test.go +++ b/pkg/cli/prompt_entity_test.go @@ -5,12 +5,11 @@ import ( "fmt" "os" "path/filepath" + "reflect" "strings" "testing" - "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/captainconfig" - "github.com/flanksource/commons-db/shell" ) func TestPromptEntityListsEmbeddedExamples(t *testing.T) { @@ -37,8 +36,34 @@ func TestPromptEntityListsEmbeddedExamples(t *testing.T) { if commit.Model != "claude-sonnet-4-6" { t.Fatalf("embedded prompt model = %q, want claude-sonnet-4-6", commit.Model) } - if len(commit.Variables) != 1 || commit.Variables[0].Name != "diff" { - t.Fatalf("embedded prompt variables = %+v, want diff", commit.Variables) + wantVariables := []PromptVariable{ + {Name: "maxBodyLines", Type: "integer", Description: "Maximum commit-message body lines; zero omits the cap", Required: true}, + {Name: "patch", Type: "string", Description: "Git patch to summarize", Required: true}, + } + if !reflect.DeepEqual(commit.Variables, wantVariables) { + t.Fatalf("embedded prompt variables = %+v, want %+v", commit.Variables, wantVariables) + } + detail, err := getPrompt(context.Background(), commit.ID) + if err != nil { + t.Fatalf("getPrompt(commit) err = %v", err) + } + wantInputSchema := map[string]any{ + "type": "object", + "additionalProperties": false, + "required": []any{"patch", "maxBodyLines"}, + "properties": map[string]any{ + "patch": map[string]any{ + "type": "string", + "description": "Git patch to summarize", + }, + "maxBodyLines": map[string]any{ + "type": "integer", + "description": "Maximum commit-message body lines; zero omits the cap", + }, + }, + } + if !reflect.DeepEqual(detail.InputSchema, wantInputSchema) { + t.Fatalf("embedded prompt input schema = %#v, want %#v", detail.InputSchema, wantInputSchema) } } @@ -296,146 +321,3 @@ func findEmbeddedPrompt(ctx context.Context, relPath string) (PromptSummary, err } return PromptSummary{}, fmt.Errorf("embedded prompt %q not found", relPath) } - -func TestRenderPromptAppliesRuntimeSpec(t *testing.T) { - isolateCaptainConfig(t) - - dir := t.TempDir() - cwd := t.TempDir() - t.Chdir(cwd) - ctx := ContextWithPromptDirs(context.Background(), []string{dir}) - content := `--- -name: Runtime Spec -model: claude-sonnet-4-6 ---- -{{role "user"}} -Hello {{name}} -` - created, err := createPrompt(ctx, map[string]any{ - "name": "Runtime Spec", - "content": content, - }) - if err != nil { - t.Fatalf("createPrompt() err = %v", err) - } - - temp := 0.2 - rendered, err := renderPrompt(ctx, created.ID, PromptRenderRequest{ - Variables: map[string]any{"name": "Ada"}, - Spec: &api.Spec{ - Model: api.Model{ - Name: "gpt-4o", - ID: "openai/gpt-4o", - Backend: api.BackendOpenAI, - Temperature: &temp, - Effort: api.EffortLow, - NoCache: true, - }, - Prompt: api.Prompt{ - System: "runtime system", - AppendSystem: "runtime append", - Source: "runtime-source", - Metadata: map[string]string{"surface": "prompt-ui"}, - }, - Budget: api.Budget{Cost: 0.5, MaxTokens: 1234, MaxTurns: 4, Timeout: "90s"}, - Permissions: api.Permissions{ - Mode: api.PermissionAcceptEdits, - Presets: []api.Preset{api.PresetEdit}, - Tools: api.Tools{ - Allow: []string{"Read"}, - Deny: []string{"Bash"}, - Modes: map[string]api.ToolMode{"Bash": api.ToolModeDisabled}, - }, - MCP: api.MCP{ - Disabled: true, - Servers: []string{"filesystem"}, - Modes: api.ResourcePolicies{"gavel": api.ResourceDisabled}, - }, - Plugins: api.ResourcePolicies{"/plugins": api.ResourceEnabled}, - Skills: api.ResourcePolicies{"/permission-skills": api.ResourceEnabled}, - }, - Memory: api.Memory{ - Skills: []string{"/skills"}, - SkipUser: true, - SkipMemory: true, - Bare: true, - }, - Setup: &shell.Setup{ - Cwd: "workspace", - DotEnv: []string{".env"}, - Checkout: &shell.Checkout{ - Mode: shell.CheckoutLocal, - Path: "/repo", - Ref: "abc123", - Worktree: &shell.Worktree{ - Mode: shell.WorktreeNew, - Prefix: "runtime-branch", - Keep: true, - }, - }, - }, - SessionID: "sess-runtime", - }, - }) - if err != nil { - t.Fatalf("renderPrompt() err = %v", err) - } - if rendered.ValidationError != "" { - t.Fatalf("render validation error = %q", rendered.ValidationError) - } - if rendered.Model != "gpt-4o" || rendered.Backend != "openai" { - t.Fatalf("rendered model/backend = %s/%s, want gpt-4o/openai", rendered.Model, rendered.Backend) - } - if rendered.Config.Model.ID != "openai/gpt-4o" { - t.Fatalf("config model ID = %q, want openai/gpt-4o", rendered.Config.Model.ID) - } - if rendered.Input.Temperature == nil || *rendered.Input.Temperature != temp { - t.Fatalf("temperature = %v, want %v", rendered.Input.Temperature, temp) - } - if rendered.Input.Budget.Cost != 0.5 || rendered.Input.Budget.MaxTokens != 1234 || - rendered.Input.Budget.MaxTurns != 4 || rendered.Input.Budget.Timeout != "90s" { - t.Fatalf("budget = %+v, want cost/maxTokens override", rendered.Input.Budget) - } - if !rendered.Input.Model.NoCache || !rendered.Config.NoCache { - t.Fatalf("noCache = input %v config %v, want true", rendered.Input.Model.NoCache, rendered.Config.NoCache) - } - if rendered.Input.Prompt.System != "runtime system" || rendered.Input.Prompt.AppendSystem != "runtime append" { - t.Fatalf("prompt system fields = %+v, want runtime overrides", rendered.Input.Prompt) - } - if rendered.Input.Prompt.Source != "runtime-source" || rendered.Input.Prompt.Metadata["surface"] != "prompt-ui" { - t.Fatalf("prompt source/metadata = %+v, want runtime overrides", rendered.Input.Prompt) - } - if rendered.Input.Permissions.Mode != api.PermissionAcceptEdits || - rendered.Input.Permissions.Tools.Modes["Bash"] != api.ToolModeDisabled || - !rendered.Input.Permissions.MCP.Disabled { - t.Fatalf("permissions = %+v, want runtime overrides", rendered.Input.Permissions) - } - if rendered.Input.Permissions.Plugins["/plugins"] != api.ResourceEnabled { - t.Fatalf("plugins = %+v, want enabled runtime plugin", rendered.Input.Permissions.Plugins) - } - if !strings.Contains(strings.Join(rendered.Input.Memory.Skills, ","), "/permission-skills") { - t.Fatalf("skills = %+v, want permission skills merged into memory skills", rendered.Input.Memory.Skills) - } - if !rendered.Input.Memory.SkipUser || !rendered.Input.Memory.SkipMemory || !rendered.Input.Memory.Bare { - t.Fatalf("memory = %+v, want runtime overrides", rendered.Input.Memory) - } - if rendered.Input.Cwd() != filepath.Join(cwd, "workspace") { - t.Fatalf("setup cwd = %q, want cwd-relative runtime dir", rendered.Input.Cwd()) - } - if rendered.Input.Setup == nil || rendered.Input.Setup.Checkout == nil || rendered.Input.Setup.Checkout.Ref != "abc123" { - t.Fatalf("setup checkout = %+v, want runtime git checkout overlay", rendered.Input.Setup) - } - if rendered.Input.Setup.Checkout.Worktree == nil || !rendered.Input.Setup.Checkout.Worktree.Keep { - t.Fatalf("worktree setup = %+v, want runtime worktree overlay", rendered.Input.Setup.Checkout.Worktree) - } - if rendered.Input.SessionID != "sess-runtime" { - t.Fatalf("runtime session = input=%+v", rendered.Input) - } -} - -func isolateCaptainConfig(t *testing.T) { - t.Helper() - path := filepath.Join(t.TempDir(), ".captain.yaml") - captainconfig.SetPathForTesting(path) - t.Cleanup(func() { captainconfig.SetPathForTesting("") }) -} diff --git a/pkg/cli/prompt_records.go b/pkg/cli/prompt_records.go new file mode 100644 index 0000000..664a530 --- /dev/null +++ b/pkg/cli/prompt_records.go @@ -0,0 +1,425 @@ +package cli + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" + "sort" + "strings" + "time" + + promptlib "github.com/flanksource/captain/pkg/ai/prompt" + dp "github.com/google/dotprompt/go/dotprompt" +) + +func listPromptRecords(ctx context.Context) ([]promptRecord, error) { + sources, err := buildPromptSources(ctx) + if err != nil { + return nil, err + } + var records []promptRecord + for _, source := range sources { + recs, err := listPromptRecordsFromSource(source) + if err != nil { + return nil, err + } + records = append(records, recs...) + } + return records, nil +} + +func listPromptRecordsFromSource(source promptSource) ([]promptRecord, error) { + var records []promptRecord + add := func(rel string, info fs.FileInfo) { + rel = filepath.ToSlash(rel) + path := rel + if source.Root != "" { + path = filepath.Join(source.Root, filepath.FromSlash(rel)) + } + updatedAt := "" + if info != nil && !info.ModTime().IsZero() { + updatedAt = info.ModTime().Format(time.RFC3339) + } + records = append(records, promptRecord{ + Source: source, + ID: encodePromptID(source.Kind, source.ID, rel), + Path: path + "\x00" + updatedAt, + Rel: rel, + }) + } + + if source.FS != nil { + err := fs.WalkDir(source.FS, source.WalkRoot, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() { + return nil + } + if !strings.HasSuffix(path, ".prompt") { + return nil + } + info, _ := d.Info() + add(path, info) + return nil + }) + return records, err + } + + err := filepath.WalkDir(source.Root, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + name := d.Name() + if d.IsDir() { + if name == ".git" || name == "node_modules" || name == "vendor" || name == "dist" { + return filepath.SkipDir + } + return nil + } + if d.Type()&fs.ModeSymlink != 0 || !strings.HasSuffix(name, ".prompt") { + return nil + } + rel, err := filepath.Rel(source.Root, path) + if err != nil { + return err + } + info, _ := d.Info() + add(rel, info) + return nil + }) + if errors.Is(err, fs.ErrNotExist) && source.Implicit { + return records, nil + } + return records, err +} + +func resolvePromptRecord(ctx context.Context, id string) (promptRecord, error) { + id = strings.TrimSpace(id) + if looksLikePromptPath(id) { + record, err := filePromptRecord(id) + if err == nil { + return record, nil + } + if !isBarePromptFilename(id) || !errors.Is(err, fs.ErrNotExist) { + return promptRecord{}, err + } + } + sources, err := buildPromptSources(ctx) + if err != nil { + return promptRecord{}, err + } + ref, decodeErr := decodePromptID(id) + if decodeErr == nil { + for _, source := range sources { + if source.Kind != ref.Kind || source.ID != ref.SourceID { + continue + } + path := ref.RelPath + if source.Root != "" { + path = filepath.Join(source.Root, filepath.FromSlash(ref.RelPath)) + } + return promptRecord{Source: source, ID: id, Path: path, Rel: ref.RelPath}, nil + } + return promptRecord{}, fmt.Errorf("prompt source %q not found", ref.SourceID) + } + return resolvePromptRecordByName(sources, id) +} + +func resolvePromptRecordByName(sources []promptSource, name string) (promptRecord, error) { + bareName := strings.TrimSuffix(name, ".prompt") + var matches []promptRecord + for _, source := range sources { + records, err := listPromptRecordsFromSource(source) + if err != nil { + return promptRecord{}, err + } + for _, record := range records { + recordName := strings.TrimSuffix(filepath.Base(record.Rel), ".prompt") + if recordName == bareName { + matches = append(matches, record) + } + } + } + switch len(matches) { + case 0: + return promptRecord{}, fmt.Errorf("prompt %q not found", name) + case 1: + return matches[0], nil + default: + paths := make([]string, len(matches)) + for i, match := range matches { + paths[i] = match.Source.Label + ":" + match.Rel + } + return promptRecord{}, fmt.Errorf("prompt name %q is ambiguous (%s); use a prompt id or path", name, strings.Join(paths, ", ")) + } +} + +func isBarePromptFilename(id string) bool { + return filepath.Base(id) == id && !strings.HasPrefix(id, ".") +} + +// looksLikePromptPath reports whether id is a filesystem path rather than a +// base64 registry id. Registry ids are base64-raw-url (no ".", "/", or leading +// "."), so a .prompt suffix, a path separator, or a leading "." marks a path. +func looksLikePromptPath(id string) bool { + return strings.HasSuffix(id, ".prompt") || + strings.ContainsRune(id, os.PathSeparator) || + strings.HasPrefix(id, ".") +} + +// filePromptRecord resolves an ad-hoc .prompt file path (not a registered id) +// into a record readable via readPromptContent/safeLocalPromptPath. Mirrors the +// captain-ai-prompt file loader so `captain prompt run|render ./x.prompt` works. +func filePromptRecord(id string) (promptRecord, error) { + abs, err := filepath.Abs(id) + if err != nil { + return promptRecord{}, err + } + info, err := os.Stat(abs) + if err != nil { + return promptRecord{}, fmt.Errorf("prompt file %s: %w", id, err) + } + if info.IsDir() { + return promptRecord{}, fmt.Errorf("%s is a directory, not a .prompt file", id) + } + return promptRecord{ + Source: promptSource{Kind: "file", ID: "file", Label: "File", Root: filepath.Dir(abs)}, + ID: id, + Path: abs, + Rel: filepath.Base(abs), + }, nil +} + +func promptSummary(record promptRecord) (PromptSummary, error) { + content, err := readPromptContent(record) + if err != nil { + return PromptSummary{}, err + } + summary, err := promptSummaryFromContent(record, content) + if err != nil { + summary = basePromptSummary(record) + summary.ParseError = err.Error() + } + if idx := strings.LastIndex(record.Path, "\x00"); idx >= 0 { + summary.UpdatedAt = strings.TrimPrefix(record.Path[idx+1:], "\x00") + } + return summary, nil +} + +func promptDetail(record promptRecord) (PromptDetail, error) { + content, err := readPromptContent(record) + if err != nil { + return PromptDetail{}, err + } + return promptDetailFromContent(record, content) +} + +func promptDetailFromContent(record promptRecord, content string) (PromptDetail, error) { + summary, err := promptSummaryFromContent(record, content) + if err != nil { + return PromptDetail{}, err + } + inspection, err := inspectPrompt(content, nil) + if err != nil { + return PromptDetail{}, err + } + return PromptDetail{ + PromptSummary: summary, + Content: content, + InputSchema: inspection.InputSchema, + InputDefault: inspection.InputDefault, + OutputSchema: inspection.OutputSchema, + Metadata: inspection.Metadata, + }, nil +} + +func promptSummaryFromContent(record promptRecord, content string) (PromptSummary, error) { + tmpl := promptlib.Load(content) + req, cfg, err := tmpl.Render(map[string]any{}, nil) + if err != nil { + return PromptSummary{}, err + } + inspection, err := inspectPrompt(content, nil) + if err != nil { + return PromptSummary{}, err + } + summary := basePromptSummary(record) + if v, ok := inspection.Metadata["name"].(string); ok && strings.TrimSpace(v) != "" { + summary.Name = strings.TrimSpace(v) + } + if v, ok := inspection.Metadata["description"].(string); ok { + summary.Description = strings.TrimSpace(v) + } + summary.Model = firstNonEmpty(cfg.Model.Name, req.Name) + summary.Backend = firstNonEmpty(string(cfg.Model.Backend), string(req.Backend)) + summary.Runtimes, err = resolvePromptRuntimes(inspection.Runtimes, cfg.Model) + if err != nil { + return PromptSummary{}, err + } + summary.Variables = inspection.Variables + return summary, nil +} + +func basePromptSummary(record promptRecord) PromptSummary { + name := strings.TrimSuffix(filepath.Base(record.Rel), ".prompt") + return PromptSummary{ + ID: record.ID, + Name: name, + SourceKind: record.Source.Kind, + SourceID: record.Source.ID, + Source: record.Source.Label, + Path: displayPromptPath(record), + RelPath: record.Rel, + Writable: record.Source.Writable, + } +} + +func displayPromptPath(record promptRecord) string { + if idx := strings.LastIndex(record.Path, "\x00"); idx >= 0 { + return record.Path[:idx] + } + return record.Path +} + +func readPromptContent(record promptRecord) (string, error) { + if record.Source.FS != nil { + data, err := fs.ReadFile(record.Source.FS, record.Rel) + if err != nil { + return "", fmt.Errorf("read embedded prompt %s: %w", record.Rel, err) + } + return string(data), nil + } + full, err := safeLocalPromptPath(record.Source, record.Rel) + if err != nil { + return "", err + } + data, err := os.ReadFile(full) + if err != nil { + return "", fmt.Errorf("read prompt %s: %w", full, err) + } + return string(data), nil +} + +func inspectPrompt(content string, data map[string]any) (promptInspection, error) { + if data == nil { + data = map[string]any{} + } + rendered, err := dp.NewDotprompt(nil).Render(content, &dp.DataArgument{Input: data}, nil) + if err != nil { + return promptInspection{}, err + } + doc, err := promptlib.Parse(content) + if err != nil { + return promptInspection{}, err + } + metadata := map[string]any{} + if rendered.Raw != nil { + for k, v := range rendered.Raw { + metadata[k] = v + } + } + if rendered.Name != "" { + metadata["name"] = rendered.Name + } + if rendered.Description != "" { + metadata["description"] = rendered.Description + } + if rendered.Model != "" { + metadata["model"] = rendered.Model + } + inputSchema := anyToMap(rendered.Input.Schema) + inputDefault := map[string]any{} + for k, v := range rendered.Input.Default { + inputDefault[k] = v + } + return promptInspection{ + Metadata: metadata, + InputSchema: inputSchema, + InputDefault: inputDefault, + OutputSchema: anyToMap(rendered.Output.Schema), + Runtimes: doc.Runtimes, + Variables: variablesFromSchema(inputSchema), + }, nil +} + +func anyToMap(v any) map[string]any { + if v == nil { + return nil + } + data, err := json.Marshal(v) + if err != nil { + return nil + } + var out map[string]any + if err := json.Unmarshal(data, &out); err != nil { + return nil + } + return out +} + +func variablesFromSchema(schema map[string]any) []PromptVariable { + props, _ := schema["properties"].(map[string]any) + if len(props) == 0 { + return nil + } + required := map[string]bool{} + if raw, ok := schema["required"].([]any); ok { + for _, item := range raw { + if s, ok := item.(string); ok { + required[s] = true + } + } + } + keys := make([]string, 0, len(props)) + for k := range props { + keys = append(keys, k) + } + sort.Strings(keys) + var vars []PromptVariable + for _, name := range keys { + prop, _ := props[name].(map[string]any) + item := PromptVariable{Name: name, Required: required[name]} + if v, ok := prop["type"].(string); ok { + item.Type = v + } + if v, ok := prop["description"].(string); ok { + item.Description = v + } + vars = append(vars, item) + } + return vars +} + +func promptSourceMatches(summary PromptSummary, source string) bool { + switch source { + case "", "all": + return true + case "embedded": + return summary.SourceKind == "embedded" + case "local": + return summary.SourceKind == "local" + default: + return summary.SourceID == source || strings.EqualFold(summary.Source, source) + } +} + +func promptMatches(summary PromptSummary, filter string) bool { + if filter == "" { + return true + } + haystack := strings.ToLower(strings.Join([]string{ + summary.Name, + summary.Description, + summary.Source, + summary.Path, + summary.RelPath, + summary.Model, + summary.Backend, + }, "\n")) + return strings.Contains(haystack, filter) +} diff --git a/pkg/cli/prompt_render.go b/pkg/cli/prompt_render.go new file mode 100644 index 0000000..0e81ba7 --- /dev/null +++ b/pkg/cli/prompt_render.go @@ -0,0 +1,333 @@ +package cli + +import ( + "context" + "fmt" + "os" + "strings" + + "github.com/flanksource/captain/pkg/ai" + promptlib "github.com/flanksource/captain/pkg/ai/prompt" + "github.com/flanksource/captain/pkg/api" +) + +// renderPrompt is the HTTP/Spec render path: overlay a structured api.Spec (the +// web UI's rich runtime overrides) onto the rendered template. +func renderPrompt(ctx context.Context, id string, renderReq PromptRenderRequest) (PromptRenderResult, error) { + if strings.TrimSpace(id) == "" { + return renderEphemeralPrompt(renderReq) + } + record, err := resolvePromptRecord(ctx, id) + if err != nil { + return PromptRenderResult{}, err + } + content, err := readPromptContent(record) + if err != nil { + return PromptRenderResult{}, err + } + vars := renderReq.Variables + if vars == nil { + vars = map[string]any{} + } + req, cfg, err := promptlib.Load(content).Render(vars, nil) + if err != nil { + return PromptRenderResult{}, err + } + req.Prompt.Source = record.Rel + if renderReq.Spec != nil { + overlayRuntimeSpec(&req, &cfg, *renderReq.Spec) + } + if err := applyPromptDefaults(&req, &cfg); err != nil { + return PromptRenderResult{}, err + } + cwd, err := os.Getwd() + if err != nil { + return PromptRenderResult{}, fmt.Errorf("get working directory: %w", err) + } + if err := normalizePromptContextDir(&req, cwd); err != nil { + return PromptRenderResult{}, err + } + return finalizeRenderResult(record, content, req, cfg, renderReq.Runtimes) +} + +func renderEphemeralPrompt(renderReq PromptRenderRequest) (PromptRenderResult, error) { + record := promptRecord{ + Source: promptSource{Kind: "ephemeral", ID: "ephemeral", Label: "Ephemeral"}, + ID: "", + Path: "", + Rel: "scratch.prompt", + } + content := ephemeralPromptContent() + var req ai.Request + var cfg ai.Config + if renderReq.Spec != nil { + overlayRuntimeSpec(&req, &cfg, *renderReq.Spec) + } + if req.Prompt.Source == "" { + req.Prompt.Source = "" + } + if err := applyPromptDefaults(&req, &cfg); err != nil { + return PromptRenderResult{}, err + } + cwd, err := os.Getwd() + if err != nil { + return PromptRenderResult{}, fmt.Errorf("get working directory: %w", err) + } + if err := normalizePromptContextDir(&req, cwd); err != nil { + return PromptRenderResult{}, err + } + return finalizeRenderResult(record, content, req, cfg, renderReq.Runtimes) +} + +func ephemeralPromptContent() string { + return `--- +name: Scratch Prompt +description: Ephemeral prompt +--- +{{role "user"}} +` +} + +// renderPromptCLI is the CLI render path: load from id | discovered name | +// .prompt filepath | -p | stdin and overlay the flat CLI flags (overlayCLI). +func renderPromptCLI(ctx context.Context, id string, opts AIPromptOptions, varsJSON, stdin string) (PromptRenderResult, error) { + content, source, usedStdin, record, err := loadPromptContent(ctx, id, opts, stdin) + if err != nil { + return PromptRenderResult{}, err + } + vars, err := promptVars(opts, varsJSON, stdin, usedStdin) + if err != nil { + return PromptRenderResult{}, err + } + req, cfg, err := renderLoadedContent(content, source, vars, opts) + if err != nil { + return PromptRenderResult{}, err + } + result, err := finalizeRenderResult(record, content, req, cfg, nil) + if err != nil { + return PromptRenderResult{}, err + } + if len(opts.MultiModels) > 0 { + result.Runtimes, err = ai.ResolveRuntimeSelectors(opts.MultiModels, result.Config.Model) + if err != nil { + return PromptRenderResult{}, err + } + } + return result, nil +} + +// finalizeRenderResult packages the rendered request/config + prompt detail into +// a PromptRenderResult and sets the validation error (shared by both paths). +func finalizeRenderResult(record promptRecord, content string, req ai.Request, cfg ai.Config, runtimeOverride []api.Model) (PromptRenderResult, error) { + // Normalize a comma-separated model into a clean primary + fallbacks so the + // displayed Model is a single name, then catch a mistyped model (primary or any + // fallback) at render time, not just on run. + req.Model = req.ExpandCSV() + cfg.Model = cfg.Model.ExpandCSV() + var err error + req.Model, err = ai.ResolveModelSelectors(req.Model) + if err != nil { + return PromptRenderResult{}, err + } + cfg.Model, err = ai.ResolveModelSelectors(cfg.Model) + if err != nil { + return PromptRenderResult{}, err + } + for _, c := range cfg.Model.Candidates() { + warnIfLikelyModelTypo(c.Name) + } + detail, err := promptDetailFromContent(record, content) + if err != nil { + return PromptRenderResult{}, err + } + runtimes := detail.Runtimes + if len(runtimeOverride) > 0 { + runtimes = runtimeOverride + } + runtimes, err = resolvePromptRuntimes(runtimes, cfg.Model) + if err != nil { + return PromptRenderResult{}, err + } + result := PromptRenderResult{ + ID: detail.ID, + Name: detail.Name, + Model: cfg.Model.Name, + Backend: string(cfg.Model.Backend), + User: req.Prompt.User, + System: req.Prompt.System, + Input: req, + Config: cfg, + InputSchema: detail.InputSchema, + InputDefault: detail.InputDefault, + OutputSchema: detail.OutputSchema, + Runtimes: runtimes, + } + switch { + case req.Prompt.User == "" && len(req.Prompt.Attachments) == 0 && !req.IsVerifyOnly(): + result.ValidationError = "prompt text or attachment required" + case cfg.Model.Name == "" && len(result.Runtimes) == 0: + result.ValidationError = "no model: set prompt frontmatter, pass a model override, or run 'captain configure'" + default: + if err := req.Validate(); err != nil { + result.ValidationError = err.Error() + } + } + return result, nil +} + +func resolvePromptRuntimes(runtimes []api.Model, base api.Model) ([]api.Model, error) { + if len(runtimes) == 0 { + return nil, nil + } + resolved := make([]api.Model, len(runtimes)) + for i, runtime := range runtimes { + if runtime.Temperature == nil { + runtime.Temperature = base.Temperature + } + if runtime.Effort == api.EffortNone { + runtime.Effort = base.Effort + } + runtime.NoCache = runtime.NoCache || base.NoCache + var err error + resolved[i], err = ai.ResolveModelSelectors(runtime) + if err != nil { + return nil, fmt.Errorf("runtime %d: %w", i+1, err) + } + } + if err := validatePromptRuntimes(resolved); err != nil { + return nil, err + } + return resolved, nil +} + +func overlayRuntimeSpec(req *ai.Request, cfg *ai.Config, spec api.Spec) { + if spec.Name != "" { + req.Name = spec.Name + cfg.Model.Name = spec.Name + req.ID = spec.ID + cfg.Model.ID = spec.ID + req.Backend = spec.Backend + cfg.Model.Backend = spec.Backend + } else { + if spec.ID != "" { + req.ID = spec.ID + cfg.Model.ID = spec.ID + } + if spec.Backend != "" { + req.Backend = spec.Backend + cfg.Model.Backend = spec.Backend + } + } + if spec.Temperature != nil { + req.Temperature = spec.Temperature + cfg.Model.Temperature = spec.Temperature + } + if spec.Effort != "" { + req.Effort = spec.Effort + cfg.Model.Effort = spec.Effort + } + if len(spec.Fallbacks) > 0 { + req.Fallbacks = spec.Fallbacks + cfg.Model.Fallbacks = spec.Fallbacks + } + req.NoCache = req.NoCache || spec.NoCache + cfg.NoCache = cfg.NoCache || spec.NoCache + if spec.Budget.Cost > 0 { + req.Budget.Cost = spec.Budget.Cost + cfg.Budget.Cost = spec.Budget.Cost + } + if spec.Budget.MaxTokens > 0 { + req.Budget.MaxTokens = spec.Budget.MaxTokens + cfg.Budget.MaxTokens = spec.Budget.MaxTokens + } + if spec.Budget.MaxTurns > 0 { + req.Budget.MaxTurns = spec.Budget.MaxTurns + cfg.Budget.MaxTurns = spec.Budget.MaxTurns + } + if spec.Budget.Timeout != "" { + req.Budget.Timeout = spec.Budget.Timeout + cfg.Budget.Timeout = spec.Budget.Timeout + } + + if spec.Prompt.User != "" { + req.Prompt.User = spec.Prompt.User + } + if spec.Prompt.System != "" { + req.Prompt.System = spec.Prompt.System + } + if spec.Prompt.AppendSystem != "" { + req.Prompt.AppendSystem = spec.Prompt.AppendSystem + } + if spec.Prompt.Source != "" { + req.Prompt.Source = spec.Prompt.Source + } + req.Prompt.Attachments = append(req.Prompt.Attachments, spec.Prompt.Attachments...) + req.Prompt.Metadata = mergeStringMaps(req.Prompt.Metadata, spec.Prompt.Metadata) + if len(spec.Prompt.SchemaJSON) > 0 { + req.Prompt.SchemaJSON = spec.Prompt.SchemaJSON + } + if spec.Prompt.SchemaStrictness != "" { + req.Prompt.SchemaStrictness = spec.Prompt.SchemaStrictness + } + if spec.Workflow != nil { + req.Workflow = spec.Workflow + } + + if spec.Permissions.Mode != "" { + req.Permissions.Mode = spec.Permissions.Mode + } + req.Permissions.Presets = mergePresets(req.Permissions.Presets, spec.Permissions.Presets) + toolPolicies := spec.Permissions.Tools.Policies() + if len(toolPolicies) > 0 { + req.Permissions.Tools.Allow = nil + req.Permissions.Tools.Deny = nil + req.Permissions.Tools.Modes = nil + for _, tool := range sortedStringKeys(toolPolicies) { + switch toolPolicies[tool] { + case api.ToolPolicyAllow: + req.Permissions.Tools.Allow = append(req.Permissions.Tools.Allow, tool) + case api.ToolPolicyDeny: + req.Permissions.Tools.Deny = append(req.Permissions.Tools.Deny, tool) + case api.ToolPolicyAsk: + if req.Permissions.Tools.Modes == nil { + req.Permissions.Tools.Modes = map[string]api.ToolMode{} + } + req.Permissions.Tools.Modes[tool] = api.ToolModeAsk + case api.ToolPolicyAuto: + if req.Permissions.Tools.Modes == nil { + req.Permissions.Tools.Modes = map[string]api.ToolMode{} + } + req.Permissions.Tools.Modes[tool] = api.ToolModeOn + } + } + } + req.Permissions.Tools.Modes = mergeToolModes(req.Permissions.Tools.Modes, spec.Permissions.Tools.Modes) + req.Permissions.MCP.Disabled = req.Permissions.MCP.Disabled || spec.Permissions.MCP.Disabled + if servers := spec.Permissions.MCP.EnabledServers(); len(servers) > 0 { + req.Permissions.MCP.Servers = servers + } + if len(spec.Permissions.Plugins) > 0 { + req.Permissions.Plugins = enabledResourcePolicies(spec.Permissions.Plugins) + } + + skills := append([]string(nil), spec.Memory.Skills...) + skills = append(skills, spec.Permissions.Skills.Enabled()...) + if len(skills) > 0 { + req.Memory.Skills = dedupeStrings(skills) + } + req.Memory.SkipProject = req.Memory.SkipProject || spec.Memory.SkipProject + req.Memory.SkipUser = req.Memory.SkipUser || spec.Memory.SkipUser + req.Memory.SkipSkills = req.Memory.SkipSkills || spec.Memory.SkipSkills + req.Memory.SkipHooks = req.Memory.SkipHooks || spec.Memory.SkipHooks + req.Memory.SkipMemory = req.Memory.SkipMemory || spec.Memory.SkipMemory + req.Memory.Bare = req.Memory.Bare || spec.Memory.Bare + + if spec.Setup != nil { + req.Setup = spec.Setup + } + + if spec.SessionID != "" { + req.SessionID = spec.SessionID + cfg.SessionID = spec.SessionID + } +} diff --git a/pkg/cli/prompt_render_test.go b/pkg/cli/prompt_render_test.go new file mode 100644 index 0000000..7599c21 --- /dev/null +++ b/pkg/cli/prompt_render_test.go @@ -0,0 +1,219 @@ +package cli + +import ( + "context" + "path/filepath" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" + "github.com/flanksource/commons-db/shell" +) + +func TestRenderPromptAppliesRuntimeSpec(t *testing.T) { + isolateCaptainConfig(t) + + dir := t.TempDir() + cwd := t.TempDir() + t.Chdir(cwd) + ctx := ContextWithPromptDirs(context.Background(), []string{dir}) + content := `--- +name: Runtime Spec +model: claude-sonnet-4-6 +--- +{{role "user"}} +Hello {{name}} +` + created, err := createPrompt(ctx, map[string]any{ + "name": "Runtime Spec", + "content": content, + }) + if err != nil { + t.Fatalf("createPrompt() err = %v", err) + } + + temp := 0.2 + rendered, err := renderPrompt(ctx, created.ID, PromptRenderRequest{ + Variables: map[string]any{"name": "Ada"}, + Spec: &api.Spec{ + Model: api.Model{ + Name: "gpt-4o", + ID: "openai/gpt-4o", + Backend: api.BackendOpenAI, + Temperature: &temp, + Effort: api.EffortLow, + NoCache: true, + }, + Prompt: api.Prompt{ + System: "runtime system", + AppendSystem: "runtime append", + Source: "runtime-source", + Metadata: map[string]string{"surface": "prompt-ui"}, + }, + Budget: api.Budget{Cost: 0.5, MaxTokens: 1234, MaxTurns: 4, Timeout: "90s"}, + Permissions: api.Permissions{ + Mode: api.PermissionAcceptEdits, + Presets: []api.Preset{api.PresetEdit}, + Tools: api.Tools{ + Allow: []string{"Read"}, + Deny: []string{"Bash"}, + Modes: map[string]api.ToolMode{"Bash": api.ToolModeOff}, + }, + MCP: api.MCP{ + Disabled: true, + Servers: []string{"filesystem"}, + Modes: api.ResourcePolicies{"gavel": api.ResourceDisabled}, + }, + Plugins: api.ResourcePolicies{"/plugins": api.ResourceEnabled}, + Skills: api.ResourcePolicies{"/permission-skills": api.ResourceEnabled}, + }, + Memory: api.Memory{ + Skills: []string{"/skills"}, + SkipUser: true, + SkipMemory: true, + Bare: true, + }, + Setup: &shell.Setup{ + Cwd: "workspace", + DotEnv: []string{".env"}, + Checkout: &shell.Checkout{ + Mode: shell.CheckoutLocal, + Path: "/repo", + Ref: "abc123", + Worktree: &shell.Worktree{ + Mode: shell.WorktreeNew, + Prefix: "runtime-branch", + Keep: true, + }, + }, + }, + SessionID: "sess-runtime", + }, + }) + if err != nil { + t.Fatalf("renderPrompt() err = %v", err) + } + if rendered.ValidationError != "" { + t.Fatalf("render validation error = %q", rendered.ValidationError) + } + if rendered.Model != "gpt-4o" || rendered.Backend != "openai" { + t.Fatalf("rendered model/backend = %s/%s, want gpt-4o/openai", rendered.Model, rendered.Backend) + } + if rendered.Config.Model.ID != "openai/gpt-4o" { + t.Fatalf("config model ID = %q, want openai/gpt-4o", rendered.Config.Model.ID) + } + if rendered.Input.Temperature == nil || *rendered.Input.Temperature != temp { + t.Fatalf("temperature = %v, want %v", rendered.Input.Temperature, temp) + } + if rendered.Input.Budget.Cost != 0.5 || rendered.Input.Budget.MaxTokens != 1234 || + rendered.Input.Budget.MaxTurns != 4 || rendered.Input.Budget.Timeout != "90s" { + t.Fatalf("budget = %+v, want cost/maxTokens override", rendered.Input.Budget) + } + if !rendered.Input.Model.NoCache || !rendered.Config.NoCache { + t.Fatalf("noCache = input %v config %v, want true", rendered.Input.Model.NoCache, rendered.Config.NoCache) + } + if rendered.Input.Prompt.System != "runtime system" || rendered.Input.Prompt.AppendSystem != "runtime append" { + t.Fatalf("prompt system fields = %+v, want runtime overrides", rendered.Input.Prompt) + } + if rendered.Input.Prompt.Source != "runtime-source" || rendered.Input.Prompt.Metadata["surface"] != "prompt-ui" { + t.Fatalf("prompt source/metadata = %+v, want runtime overrides", rendered.Input.Prompt) + } + if rendered.Input.Permissions.Mode != api.PermissionAcceptEdits || + rendered.Input.Permissions.Tools.Modes["Bash"] != api.ToolModeOff || + !rendered.Input.Permissions.MCP.Disabled { + t.Fatalf("permissions = %+v, want runtime overrides", rendered.Input.Permissions) + } + if rendered.Input.Permissions.Plugins["/plugins"] != api.ResourceEnabled { + t.Fatalf("plugins = %+v, want enabled runtime plugin", rendered.Input.Permissions.Plugins) + } + if !strings.Contains(strings.Join(rendered.Input.Memory.Skills, ","), "/permission-skills") { + t.Fatalf("skills = %+v, want permission skills merged into memory skills", rendered.Input.Memory.Skills) + } + if !rendered.Input.Memory.SkipUser || !rendered.Input.Memory.SkipMemory || !rendered.Input.Memory.Bare { + t.Fatalf("memory = %+v, want runtime overrides", rendered.Input.Memory) + } + if rendered.Input.Cwd() != filepath.Join(cwd, "workspace") { + t.Fatalf("setup cwd = %q, want cwd-relative runtime dir", rendered.Input.Cwd()) + } + if rendered.Input.Setup == nil || rendered.Input.Setup.Checkout == nil || rendered.Input.Setup.Checkout.Ref != "abc123" { + t.Fatalf("setup checkout = %+v, want runtime git checkout overlay", rendered.Input.Setup) + } + if rendered.Input.Setup.Checkout.Worktree == nil || !rendered.Input.Setup.Checkout.Worktree.Keep { + t.Fatalf("worktree setup = %+v, want runtime worktree overlay", rendered.Input.Setup.Checkout.Worktree) + } + if rendered.Input.SessionID != "sess-runtime" { + t.Fatalf("runtime session = input=%+v", rendered.Input) + } +} + +func TestRenderPromptEphemeralSpec(t *testing.T) { + isolateCaptainConfig(t) + + cwd := t.TempDir() + t.Chdir(cwd) + temp := 0.1 + rendered, err := renderPrompt(context.Background(), "", PromptRenderRequest{ + Spec: &api.Spec{ + Model: api.Model{ + Name: "gpt-5.5", + Backend: api.BackendCodexAgent, + Temperature: &temp, + Effort: api.EffortHigh, + }, + Prompt: api.Prompt{ + System: "scratch system", + User: "Draft a deployment plan", + }, + Budget: api.Budget{Timeout: "2h"}, + }, + }) + if err != nil { + t.Fatalf("renderPrompt() err = %v", err) + } + if rendered.ValidationError != "" { + t.Fatalf("render validation error = %q", rendered.ValidationError) + } + if rendered.ID != "" || rendered.Name != "Scratch Prompt" { + t.Fatalf("rendered prompt identity = id %q name %q, want scratch prompt", rendered.ID, rendered.Name) + } + if rendered.User != "Draft a deployment plan" || rendered.System != "scratch system" { + t.Fatalf("rendered prompt = user %q system %q", rendered.User, rendered.System) + } + if rendered.Model != "gpt-5.5" || rendered.Backend != "codex-agent" { + t.Fatalf("rendered model/backend = %s/%s, want gpt-5.5/codex-agent", rendered.Model, rendered.Backend) + } + if rendered.Input.Prompt.Source != "" { + t.Fatalf("prompt source = %q, want ", rendered.Input.Prompt.Source) + } + if rendered.Input.Cwd() != cwd { + t.Fatalf("setup cwd = %q, want %q", rendered.Input.Cwd(), cwd) + } +} + +func TestApplyPromptDefaultsSelectorEffortWins(t *testing.T) { + isolateCaptainConfig(t) + req := ai.Request{Model: api.Model{Effort: api.EffortLow}} + cfg := ai.Config{Model: api.Model{ + Name: "gpt-5.6-sol", + Backend: api.BackendCodexAgent, + Effort: api.EffortHigh, + }} + if err := applyPromptDefaults(&req, &cfg); err != nil { + t.Fatalf("applyPromptDefaults: %v", err) + } + if req.Name != "gpt-5.6-sol" || req.Backend != api.BackendCodexAgent || req.Effort != api.EffortHigh { + t.Fatalf("request = %+v, want selector model/effort", req.Model) + } + if cfg.Model.Effort != api.EffortHigh { + t.Fatalf("config effort = %q, want high", cfg.Model.Effort) + } +} + +func isolateCaptainConfig(t *testing.T) { + t.Helper() + path := filepath.Join(t.TempDir(), ".captain.yaml") + captainconfig.SetPathForTesting(path) + t.Cleanup(func() { captainconfig.SetPathForTesting("") }) +} diff --git a/pkg/cli/prompt_run.go b/pkg/cli/prompt_run.go index a553786..85e6499 100644 --- a/pkg/cli/prompt_run.go +++ b/pkg/cli/prompt_run.go @@ -4,35 +4,18 @@ import ( "context" "errors" "fmt" + "strings" "time" "github.com/flanksource/captain/pkg/ai" - "github.com/flanksource/captain/pkg/ai/agent" - "github.com/flanksource/captain/pkg/ai/agent/verify" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" clickyrpc "github.com/flanksource/clicky/rpc" "github.com/flanksource/clicky/task" flanksourceContext "github.com/flanksource/commons/context" "github.com/google/uuid" ) -// PromptRunResult is the unified result of the "run" action. Over HTTP (serve) -// it carries the async handle (RunID + Status "running") — the web UI then -// streams from /api/captain/prompt/runs/{runId}/stream. On the CLI it carries the -// synchronous result (Text + tokens/cost). One type serves both transports. -type PromptRunResult struct { - RunID string `json:"runId,omitempty"` - Status string `json:"status,omitempty" pretty:"label=Status"` - Model string `json:"model,omitempty" pretty:"label=Model"` - Backend string `json:"backend,omitempty" pretty:"label=Backend"` - - Text string `json:"text,omitempty" pretty:"label=Response"` - SessionID string `json:"sessionId,omitempty" pretty:"label=Session"` - InputTokens int `json:"inputTokens,omitempty" pretty:"label=Input Tokens"` - OutputTokens int `json:"outputTokens,omitempty" pretty:"label=Output Tokens"` - CostUSD float64 `json:"costUSD,omitempty" pretty:"label=Cost USD"` - Duration string `json:"duration,omitempty" pretty:"label=Duration"` -} - // PromptRunSummary is the terminal payload of a run: the SSE stream ends with it // and the task's typed result carries it. type PromptRunSummary struct { @@ -48,15 +31,17 @@ type PromptRunSummary struct { Error string `json:"error,omitempty"` } -// runPromptAction renders the prompt (from an id | .prompt filepath | --prompt | -// stdin), then executes it — synchronously on the CLI (returns the text + cost) -// or asynchronously over HTTP (returns a run handle to stream from). This is the -// single prompt-run implementation; `captain ai prompt` is a deprecated alias. +// runPromptAction renders the prompt (from an id | discovered name | .prompt +// filepath | --prompt | stdin), then executes it — synchronously on the CLI +// (returns the text + cost) or asynchronously over HTTP (returns a run handle to +// stream from). This is the single prompt-run implementation; `captain ai +// prompt` is a deprecated alias. func runPromptAction(ctx context.Context, id string, flags map[string]string) (PromptRunResult, error) { _, isHTTP := clickyrpc.RequestFromContext(ctx) var rendered PromptRenderResult var opts AIPromptOptions + var chatRequested bool if isHTTP { req, err := readRenderRequest(ctx, flags) if err != nil { @@ -65,6 +50,7 @@ func runPromptAction(ctx context.Context, id string, flags map[string]string) (P if rendered, err = renderPrompt(ctx, id, req); err != nil { return PromptRunResult{}, err } + chatRequested = req.Chat } else { var err error if opts, err = actionFlagsToOptions(flags); err != nil { @@ -77,154 +63,302 @@ func runPromptAction(ctx context.Context, id string, flags map[string]string) (P if rendered.ValidationError != "" { return PromptRunResult{}, errors.New(rendered.ValidationError) } + if chatRequested { + if rendered.Input.Prompt.HasSchema() { + return PromptRunResult{}, errors.New("chat mode does not support structured-output prompts") + } + if workflowConfigured(rendered.Input.Workflow) { + return PromptRunResult{}, errors.New("chat mode does not support workflow-backed prompts") + } + } if isHTTP { - return launchAsyncRun(id, rendered), nil + if len(rendered.Runtimes) > 0 { + return launchAsyncBatch(ctx, id, rendered, rendered.Runtimes, chatRequested) + } + return launchAsyncRun(id, rendered, chatRequested), nil } return executeSyncRun(ctx, rendered, opts) } +var executePromptRequestFunc = executePromptRequest + // launchAsyncRun starts the background clicky task + SSE stream and returns the // run handle (the serve/web-UI contract). -func launchAsyncRun(id string, rendered PromptRenderResult) PromptRunResult { +func launchAsyncRun(id string, rendered PromptRenderResult, chat bool) PromptRunResult { runID := uuid.NewString() stream := promptRuns.create(runID) timeout := runtimeTimeout(rendered.Input.Budget.Timeout) + capabilities := chatCapabilitiesForBackend(rendered.Backend) + stream.setRun(PromptRunFrame{ + RunID: runID, Status: "running", Chat: chat, Model: rendered.Model, + Backend: rendered.Backend, Capabilities: capabilities, + }) group := task.StartGroup[PromptRunSummary]( "prompt "+rendered.Name, task.WithGroupID(runID), task.WithKind("prompt"), - task.WithLabels(map[string]string{ - "promptId": id, - "model": rendered.Model, - "backend": rendered.Backend, - }), + task.WithLabels(promptTaskLabelsWithID(rendered, id, "")), ) - group.Add("execute", func(_ flanksourceContext.Context, t *task.Task) (PromptRunSummary, error) { - return runPromptStream(t, rendered, timeout, runID, stream) - }) - return PromptRunResult{RunID: runID, Status: "running", Model: rendered.Model, Backend: rendered.Backend} + if chat { + chatSession := newChatSession(runID, rendered, timeout, stream, nil) + promptChats.register(chatSession) + group.Add("execute", func(_ flanksourceContext.Context, t *task.Task) (PromptRunSummary, error) { + return chatSession.run(t) + }) + } else { + group.Add("execute", func(_ flanksourceContext.Context, t *task.Task) (PromptRunSummary, error) { + return runPromptStream(t, rendered, timeout, runID, stream, nil) + }) + } + return PromptRunResult{ + RunID: runID, Status: "running", Model: rendered.Model, Backend: rendered.Backend, + Chat: chat, Capabilities: capabilities, + } +} + +func workflowConfigured(workflow *api.Workflow) bool { + return workflow != nil && (workflow.Verify != nil || workflow.PostRun != nil || workflow.AutoVerifyWithoutFixture) } // executeSyncRun runs the prompt in-process (CLI) — live output to stderr, final // result returned — and persists the realized prompt for the launched session. func executeSyncRun(ctx context.Context, rendered PromptRenderResult, opts AIPromptOptions) (PromptRunResult, error) { - out, err := executePromptRequest(ctx, rendered.Input, rendered.Config, runtimeTimeout(rendered.Input.Budget.Timeout), opts.NoStream) + if len(rendered.Runtimes) > 0 || len(opts.MultiModels) > 0 { + return executeSyncBatch(ctx, rendered, opts) + } + return executeSyncRunSingle(ctx, rendered, opts) +} + +func executeSyncRunSingle(ctx context.Context, rendered PromptRenderResult, opts AIPromptOptions) (PromptRunResult, error) { + group := task.StartGroup[PromptRunResult]( + "prompt "+rendered.Name, + task.WithKind("prompt"), + task.WithLabels(promptTaskLabels(rendered, "")), + ) + run := group.Add("execute "+rendered.Backend+":"+rendered.Model, func(_ flanksourceContext.Context, t *task.Task) (PromptRunResult, error) { + taskCtx := ai.ContextWithLogger(t.Context(), t) + return executeSyncRunSingleDirect(taskCtx, rendered, opts, nil) + }, task.WithModel(rendered.Model), task.WithPrompt(rendered.Input.Prompt.User)) + return run.GetResult() +} + +func executeSyncRunSingleDirect(ctx context.Context, rendered PromptRenderResult, opts AIPromptOptions, binding *promptSessionBinding) (PromptRunResult, error) { + out, err := executePromptRequestFunc(ctx, rendered.Input, rendered.Config, runtimeTimeout(rendered.Input.Budget.Timeout), opts.NoStream) if err != nil { return PromptRunResult{}, err } r, _ := out.(AIPromptResult) - if r.SessionID != "" { - if st := sessionStore(); st != nil { - st.upsertPrompt(StoredPrompt{ - SessionID: r.SessionID, - Model: r.Model, - Backend: r.Backend, - Realized: rendered, - }) - } - } + persistPromptRun(context.WithoutCancel(ctx), promptRunRecordInput{ + Rendered: rendered, SessionID: r.SessionID, Model: r.Model, Backend: r.Backend, + Binding: binding, ResultText: r.Text, ResultJSON: r.StructuredOutput, + }) return PromptRunResult{ - Status: "completed", - Model: r.Model, - Backend: r.Backend, - Text: r.Text, - SessionID: r.SessionID, - InputTokens: r.InputTokens, - OutputTokens: r.Output, - CostUSD: r.CostUSD, - Duration: r.Duration, + Status: "completed", + Model: r.Model, + Backend: r.Backend, + Text: r.Text, + StructuredOutput: r.StructuredOutput, + SessionID: r.SessionID, + Dir: r.Dir, + HistoryFile: r.HistoryFile, + InputTokens: r.InputTokens, + OutputTokens: r.Output, + CostUSD: r.CostUSD, + Duration: r.Duration, }, nil } -// runPromptStream drives a single streaming iteration, converting each ai.Event -// into a SessionEntry frame (published to stream) while driving the task's live -// status. It runs on clicky's worker goroutine and derives its context from the -// task (t.Context()), not the HTTP request, so the run outlives the POST. -func runPromptStream(t *task.Task, rendered PromptRenderResult, timeout time.Duration, runID string, stream *runStream) (PromptRunSummary, error) { - ctx, cancel := runContext(t.Context(), rendered.Input, timeout) - defer cancel() - ctx = ai.ContextWithLogger(ctx, t) - - req := rendered.Input - cfg := rendered.Config - p, cleanup, err := buildProvider(ctx, &req, cfg) - if err != nil { - return failRun(t, stream, err) +func executeSyncBatch(ctx context.Context, rendered PromptRenderResult, opts AIPromptOptions) (PromptRunResult, error) { + models := rendered.Runtimes + if len(models) == 0 { + var err error + models, err = ai.ResolveRuntimeSelectors(opts.MultiModels, rendered.Config.Model) + if err != nil { + return PromptRunResult{}, err + } } - defer cleanup() - - streamer, ok := p.(ai.StreamingProvider) - if !ok && !req.IsVerifyOnly() { - return failRun(t, stream, fmt.Errorf("backend %s does not support streaming", rendered.Backend)) + if len(models) == 0 { + return executeSyncRunSingle(ctx, rendered, opts) + } + if len(models) == 1 { + variant := renderVariant(rendered, models[0], fallbackModelsFromFlags(opts.Fallback)) + opts.MultiModels = nil + start := time.Now() + single, runErr := executeSyncRunSingle(ctx, variant, opts) + item := PromptRunItem{ + Selector: runtimeSelector(models[0]), Model: firstNonEmpty(single.Model, models[0].Name), + Backend: firstNonEmpty(single.Backend, string(models[0].Backend)), Effort: string(models[0].Effort), + Text: single.Text, SessionID: single.SessionID, Dir: single.Dir, HistoryFile: single.HistoryFile, + InputTokens: single.InputTokens, OutputTokens: single.OutputTokens, CostUSD: single.CostUSD, Duration: single.Duration, + } + result := PromptRunResult{Total: 1, Runs: []PromptRunItem{item}, Duration: time.Since(start).Round(time.Millisecond).String()} + if runErr != nil { + result.Status, result.Failed, result.Runs[0].Status, result.Runs[0].Error = "failed", 1, "failed", runErr.Error() + return result, runErr + } + result.Status, result.Succeeded, result.Runs[0].Status = "completed", 1, single.Status + return result, nil + } + prepared := rendered.Input + if err := resolvePromptAttachments(ctx, &prepared); err != nil { + return PromptRunResult{}, err + } + rendered.Input.Prompt.Attachments = prepared.Prompt.Attachments + if rendered.Input.SessionID != "" && len(models) > 1 { + return PromptRunResult{}, errors.New("--resume cannot be used with multiple --multi-models variants") + } + if forced := api.Backend(strings.TrimSpace(opts.Backend)); forced != "" { + for _, model := range models { + if model.Backend != forced { + return PromptRunResult{}, fmt.Errorf("--multi-models selector %s:%s conflicts with --backend %s", model.Backend, model.Name, forced) + } + } } start := time.Now() - acc := newPromptEventAccumulator(stream.publish, t, rendered.Model, rendered.Backend) - acc.cwd = req.Cwd() - - // Drive the generate→verify loop declared by spec.Workflow via the hook - // runner: no verify ⇒ a single generation; verify commands re-run on failure - // (feedback appended) up to maxIterations; a body-less spec verifies only. - runner := &agent.Runner[string]{ - Provider: streamer, - Request: req, - Hooks: verify.HooksForWorkflow(req.Workflow), - MaxIterations: verify.MaxIterationsForWorkflow(req.Workflow), - Repo: req.Cwd(), - Cwd: req.Cwd(), - Scope: verify.ScopeForWorkflow(req.Workflow), - OnEvent: acc.handle, - } - runResult, err := runner.Run(ctx) + batch, err := createPromptBatchSessions(ctx, rendered, models) if err != nil { - return failRun(t, stream, err) - } - loop := runResult.Loop - - session := acc.sessionID - if session == "" && runResult.Response.Workspace != nil { - session = runResult.Response.Workspace.SessionID - } - if session == "" && loop != nil && len(loop.Iterations) > 0 { - session = loop.Iterations[0].SessionID - } - // Persist the realized prompt for this launched session so `sessions get` can - // show what produced it. External (non-captain) sessions have no such record. - if session != "" { - if st := sessionStore(); st != nil { - st.upsertPrompt(StoredPrompt{ - SessionID: session, - RunID: runID, - Model: acc.model, - Backend: rendered.Backend, - Realized: rendered, - }) + return PromptRunResult{}, err + } + updatePromptSessionLifecycle(ctx, batch.ID, database.SessionLifecycleRunning, "") + runs := make([]PromptRunItem, len(models)) + group := task.StartGroup[PromptRunItem]( + "prompt "+rendered.Name, + task.WithKind("prompt"), + task.WithLabels(promptTaskLabels(rendered, "multi")), + task.WithConcurrency(len(models)), + ) + tasks := make([]task.TypedTask[PromptRunItem], len(models)) + for i, model := range models { + i, model := i, model + selector := string(model.Backend) + ":" + model.Name + if model.Effort != api.EffortNone { + selector += ":" + string(model.Effort) + } + tasks[i] = group.Add(selector, func(_ flanksourceContext.Context, t *task.Task) (PromptRunItem, error) { + variant := renderVariant(rendered, model, fallbackModelsFromFlags(opts.Fallback)) + variantOpts := opts + variantOpts.MultiModels = nil + taskCtx := ai.ContextWithLogger(t.Context(), t) + binding := promptBinding(batch, i) + updatePromptSessionLifecycle(context.WithoutCancel(taskCtx), binding.SessionID, database.SessionLifecycleRunning, "") + result, err := executeSyncRunSingleDirect(taskCtx, variant, variantOpts, binding) + item := PromptRunItem{ + RunID: binding.SessionID.String(), + Selector: selector, + Model: model.Name, + Backend: string(model.Backend), + Dir: actualRunDir(variant.Input), + } + if err != nil { + item.Status = "failed" + item.HistoryFile = historyFileForRun(model.Backend, item.SessionID, item.Dir) + item.Error = err.Error() + persistPromptRun(context.WithoutCancel(taskCtx), promptRunRecordInput{ + Rendered: variant, RunID: binding.SessionID.String(), Binding: binding, + Model: model.Name, Backend: string(model.Backend), Error: err.Error(), + }) + return item, err + } + item.Status = result.Status + item.Model = firstNonEmpty(result.Model, item.Model) + item.Backend = firstNonEmpty(result.Backend, item.Backend) + item.Text = result.Text + item.StructuredOutput = result.StructuredOutput + providerSessionID := result.SessionID + item.SessionID = binding.SessionID.String() + item.Dir = firstNonEmpty(result.Dir, item.Dir) + item.HistoryFile = firstNonEmpty(result.HistoryFile, historyFileForRun(api.Backend(item.Backend), providerSessionID, item.Dir)) + item.InputTokens = result.InputTokens + item.OutputTokens = result.OutputTokens + item.CostUSD = result.CostUSD + item.Duration = result.Duration + return item, nil + }, task.WithModel(model.Name), task.WithPrompt(rendered.Input.Prompt.User)) + } + for i, task := range tasks { + item, err := task.GetResult() + if err != nil && item.Error == "" { + item.Status = "failed" + item.Error = err.Error() + } + runs[i] = item + } + + result := PromptRunResult{ + BatchID: batch.ID.String(), + Status: "completed", + Total: len(runs), + Runs: runs, + Duration: time.Since(start).Round(time.Millisecond).String(), + } + for _, run := range runs { + if run.Error != "" || run.Status == "failed" { + result.Failed++ + continue } + result.Succeeded++ + result.InputTokens += run.InputTokens + result.OutputTokens += run.OutputTokens + result.CostUSD += run.CostUSD + } + switch { + case result.Failed == 0: + result.Status = "completed" + case result.Succeeded == 0: + result.Status = "failed" + default: + result.Status = "partial" + } + updatePromptSessionLifecycle(context.WithoutCancel(ctx), batch.ID, + batchLifecycle(result.Succeeded, result.Failed), + fmt.Sprintf("%d succeeded, %d failed", result.Succeeded, result.Failed)) + if result.Succeeded == 0 && result.Failed > 0 { + return result, fmt.Errorf("all %d prompt variants failed", result.Failed) } - passed := verifyPassed(runResult.Verdicts) - summary := PromptRunSummary{ - RunID: runID, - SessionID: session, - Model: acc.model, - Backend: rendered.Backend, - InputTokens: acc.usage.InputTokens, - OutputTokens: acc.usage.OutputTokens, - CostUSD: acc.cost, - Duration: time.Since(start).Round(time.Millisecond).String(), - Success: passed, - } - if !passed { - summary.Error = verifyReason(runResult.Verdicts) - } - stream.complete(summary) - t.Success() - return summary, nil + return result, nil } -func failRun(t *task.Task, stream *runStream, err error) (PromptRunSummary, error) { - stream.fail(err.Error()) - _, _ = t.FailedWithError(err) - return PromptRunSummary{Error: err.Error()}, err +func promptTaskLabels(rendered PromptRenderResult, mode string) map[string]string { + labels := map[string]string{} + if mode != "" { + labels["mode"] = mode + } + if mode != "multi" { + labels["model"] = rendered.Model + labels["backend"] = rendered.Backend + } + if rendered.Name != "" { + labels["prompt"] = rendered.Name + } + if source := rendered.Input.Prompt.Source; source != "" { + labels["source"] = source + } + if cwd := rendered.Input.Cwd(); cwd != "" { + labels["cwd"] = cwd + } + return labels +} + +func promptTaskLabelsWithID(rendered PromptRenderResult, id, mode string) map[string]string { + labels := promptTaskLabels(rendered, mode) + if id != "" { + labels["promptId"] = id + } + return labels +} + +func renderVariant(rendered PromptRenderResult, model api.Model, fallbacks []api.Model) PromptRenderResult { + out := rendered + req := rendered.Input + cfg := rendered.Config + req.Model = variantModel(model, fallbacks) + cfg.Model = variantModel(model, fallbacks) + out.Input = req + out.Config = cfg + out.Model = cfg.Model.Name + out.Backend = string(cfg.Model.Backend) + return out } diff --git a/pkg/cli/prompt_run_events.go b/pkg/cli/prompt_run_events.go index cafb6f3..8a1b696 100644 --- a/pkg/cli/prompt_run_events.go +++ b/pkg/cli/prompt_run_events.go @@ -3,6 +3,7 @@ package cli import ( "fmt" "strings" + "sync" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/session" @@ -26,6 +27,7 @@ type taskSink interface { // into a single message keyed by a stable id so the viewer replaces-in-place, // and correlates each tool call with its later result via ToolCallID. type promptEventAccumulator struct { + mu sync.Mutex emit func(session.Message) task taskSink @@ -33,6 +35,7 @@ type promptEventAccumulator struct { model string backend string cwd string + idPrefix string toolByID map[string]*session.Message @@ -60,6 +63,8 @@ func newPromptEventAccumulator(emit func(session.Message), t taskSink, model, ba // handle maps a single ai.Event. It matches the ai.LoopOptions.OnEvent // signature; the iteration index is unused (prompt runs are single-iteration). func (a *promptEventAccumulator) handle(_ int, ev ai.Event) { + a.mu.Lock() + defer a.mu.Unlock() if ev.Model != "" { a.model = ev.Model } @@ -90,6 +95,10 @@ func (a *promptEventAccumulator) handle(_ int, ev ai.Event) { a.emitError(ev) case ai.EventResult: a.flush() + if len(ev.StructuredData) > 0 { + a.appendText(string(ev.StructuredData)) + a.flush() + } if ev.Usage != nil { a.usage = *ev.Usage } @@ -212,7 +221,16 @@ func (a *promptEventAccumulator) emitError(ev ai.Event) { func (a *promptEventAccumulator) nextID(kind string) string { a.seq++ - return fmt.Sprintf("%s-%d", kind, a.seq) + if a.idPrefix == "" { + return fmt.Sprintf("%s-%d", kind, a.seq) + } + return fmt.Sprintf("%s-%s-%d", a.idPrefix, kind, a.seq) +} + +func (a *promptEventAccumulator) snapshot() (string, string, ai.Usage, float64) { + a.mu.Lock() + defer a.mu.Unlock() + return a.sessionID, a.model, a.usage, a.cost } // toolID keys a tool message by its ToolCallID so the call and its later result diff --git a/pkg/cli/prompt_run_events_test.go b/pkg/cli/prompt_run_events_test.go index 86360f0..6c21c7b 100644 --- a/pkg/cli/prompt_run_events_test.go +++ b/pkg/cli/prompt_run_events_test.go @@ -49,6 +49,22 @@ func TestPromptRunAccumulator_CoalescesTextDeltas(t *testing.T) { } } +func TestPromptRunAccumulator_EmitsStructuredResultAsFreshMessage(t *testing.T) { + msgs := collectEntries("m", "b", + ai.Event{Kind: ai.EventText, Text: "narrative"}, + ai.Event{Kind: ai.EventResult, StructuredData: json.RawMessage(`{"answer":"42"}`)}, + ) + if len(msgs) != 2 { + t.Fatalf("want narrative and one structured result frame, got %d", len(msgs)) + } + if msgs[0].ID == msgs[1].ID { + t.Fatalf("structured result must start a fresh message") + } + if got := msgs[1].Parts[0].Text; got != `{"answer":"42"}` { + t.Fatalf("structured result text = %q, want coherent JSON", got) + } +} + func TestPromptRunAccumulator_ThinkingAndTextAreSeparateTurns(t *testing.T) { msgs := collectEntries("m", "b", ai.Event{Kind: ai.EventThinking, Text: "hmm"}, diff --git a/pkg/cli/prompt_run_history.go b/pkg/cli/prompt_run_history.go new file mode 100644 index 0000000..9962c5d --- /dev/null +++ b/pkg/cli/prompt_run_history.go @@ -0,0 +1,66 @@ +package cli + +import ( + "path/filepath" + "sort" + "strings" + + "github.com/flanksource/captain/pkg/ai/history" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" +) + +func historyFileForRun(backend api.Backend, sessionID, cwd string) string { + if strings.TrimSpace(sessionID) == "" { + return "" + } + switch backend { + case api.BackendClaudeAgent, api.BackendClaudeCLI, api.BackendClaudeCmux: + return claudeHistoryFile(sessionID, cwd) + case api.BackendCodexAgent, api.BackendCodexCLI, api.BackendCodexCmux: + return codexHistoryFile(sessionID) + default: + return "" + } +} + +func claudeHistoryFile(sessionID, cwd string) string { + if cwd == "" { + return "" + } + abs, err := filepath.Abs(cwd) + if err != nil { + return "" + } + return filepath.Join(claude.GetProjectsDir(), claude.NormalizePath(abs), sessionID+".jsonl") +} + +func codexHistoryFile(sessionID string) string { + files, err := history.FindCodexSessionFiles() + if err != nil || len(files) == 0 { + return "" + } + sort.Strings(files) + for _, file := range files { + if strings.TrimSuffix(filepath.Base(file), filepath.Ext(file)) == sessionID { + return file + } + meta, err := history.ReadCodexSessionMeta(file) + if err == nil && meta != nil && meta.ID == sessionID { + return file + } + } + for _, file := range files { + if strings.HasPrefix(strings.TrimSuffix(filepath.Base(file), filepath.Ext(file)), sessionID) { + return file + } + } + return "" +} + +func variantModel(model api.Model, fallbacks []api.Model) api.Model { + if len(fallbacks) > 0 { + model.Fallbacks = fallbacks + } + return model +} diff --git a/pkg/cli/prompt_run_live.go b/pkg/cli/prompt_run_live.go new file mode 100644 index 0000000..0d05f26 --- /dev/null +++ b/pkg/cli/prompt_run_live.go @@ -0,0 +1,115 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/agent" + "github.com/flanksource/captain/pkg/ai/agent/verify" + "github.com/flanksource/clicky/task" +) + +func runPromptStream(t *task.Task, rendered PromptRenderResult, timeout time.Duration, runID string, stream *runStream, binding *promptSessionBinding) (PromptRunSummary, error) { + ctx, cancel := runContext(t.Context(), rendered.Input, timeout) + stream.setCancel(cancel) + defer cancel() + ctx = ai.ContextWithLogger(ctx, t) + + req := rendered.Input + cfg := rendered.Config + if err := preparePromptAttachments(ctx, &req, cfg); err != nil { + return failRun(t, stream, err) + } + p, cleanup, err := buildProvider(ctx, &req, cfg) + if err != nil { + return failRun(t, stream, err) + } + defer cleanup() + defer closeProvider(p) + + streamer, ok := p.(ai.StreamingProvider) + if !ok && !req.IsVerifyOnly() { + return failRun(t, stream, fmt.Errorf("backend %s does not support streaming", rendered.Backend)) + } + + start := time.Now() + acc := newPromptEventAccumulator(stream.publish, t, rendered.Model, rendered.Backend) + acc.cwd = req.Cwd() + acc.idPrefix = runID + runner := &agent.Runner[string]{ + Provider: streamer, + Request: req, + Hooks: verify.HooksForWorkflow(req.Workflow), + MaxIterations: verify.MaxIterationsForWorkflow(req.Workflow), + Repo: req.Cwd(), + Cwd: req.Cwd(), + Scope: verify.ScopeForWorkflow(req.Workflow), + OnEvent: acc.handle, + } + runResult, err := runner.Run(ctx) + if err != nil { + if stream.wasStopped() { + err = errors.New("stopped") + } + return failRun(t, stream, err) + } + loop := runResult.Loop + + session := acc.sessionID + if session == "" && runResult.Response.Workspace != nil { + session = runResult.Response.Workspace.SessionID + } + if session == "" && loop != nil && len(loop.Iterations) > 0 { + session = loop.Iterations[0].SessionID + } + passed := verifyPassed(runResult.Verdicts) + structuredOutput, err := structuredOutputMap(runResult.Response.StructuredData) + if err != nil { + return failRun(t, stream, err) + } + resultText, err := structuredOutputText(runResult.Response.Text, structuredOutput) + if err != nil { + return failRun(t, stream, err) + } + record := promptRunRecordInput{ + Rendered: rendered, RunID: runID, Binding: binding, SessionID: session, + Model: acc.model, Backend: rendered.Backend, ResultText: resultText, ResultJSON: structuredOutput, + } + if !passed { + record.Error = verifyReason(runResult.Verdicts) + } + persistPromptRun(context.WithoutCancel(ctx), record) + summarySessionID := session + if binding != nil { + summarySessionID = binding.SessionID.String() + } + summary := PromptRunSummary{ + RunID: runID, + SessionID: summarySessionID, + Model: acc.model, + Backend: rendered.Backend, + InputTokens: acc.usage.InputTokens, + OutputTokens: acc.usage.OutputTokens, + CostUSD: acc.cost, + Duration: time.Since(start).Round(time.Millisecond).String(), + Success: passed, + } + if !passed { + summary.Error = verifyReason(runResult.Verdicts) + } + stream.complete(summary) + t.Success() + return summary, nil +} + +func failRun(t *task.Task, stream *runStream, err error) (PromptRunSummary, error) { + if stream.wasStopped() { + err = errors.New("stopped") + } + stream.fail(err.Error()) + _, _ = t.FailedWithError(err) + return PromptRunSummary{Error: err.Error()}, err +} diff --git a/pkg/cli/prompt_run_persist.go b/pkg/cli/prompt_run_persist.go new file mode 100644 index 0000000..5e25ea4 --- /dev/null +++ b/pkg/cli/prompt_run_persist.go @@ -0,0 +1,163 @@ +package cli + +import ( + "context" + "encoding/json" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/google/uuid" +) + +// promptRunRecordInput is a completed captain-launched prompt run to persist. +type promptRunRecordInput struct { + Rendered PromptRenderResult + RunID string + Binding *promptSessionBinding + SessionID string + Model string + Backend string + BatchID *uuid.UUID + ResultText string + ResultJSON map[string]any + Error string +} + +// persistPromptRun records a captain-launched run against its session in the +// native store and registers the transcript for live tailing. Persistence +// failures are reported loudly but never fail the completed run itself. +func persistPromptRun(ctx context.Context, input promptRunRecordInput) { + if input.Binding == nil && strings.TrimSpace(input.SessionID) == "" { + return + } + source := backendSource(api.Backend(input.Backend)) + if source == "" { + source = "claude" + } + db, err := captainDB(ctx) + if err != nil { + log.Errorf("persist prompt run for session %s: %v", input.SessionID, err) + return + } + var session *database.Session + if input.Binding != nil { + session, err = db.GetSession(ctx, input.Binding.SessionID) + if err == nil && strings.TrimSpace(input.SessionID) != "" { + providerSessionID := strings.TrimSpace(input.SessionID) + session, err = db.UpdateSessionState(ctx, database.UpdateSessionStateInput{ + ID: session.ID, ExpectedVersion: session.StateVersion, ProviderSessionID: &providerSessionID, + }) + } + } else { + session, err = db.CreateOrGetSession(ctx, database.CreateSessionInput{ + ProviderSessionID: input.SessionID, Source: source, HostID: captainHostID(), + Provider: ai.BackendToProvider(api.Backend(input.Backend)), CWD: input.Rendered.Input.Cwd(), + }) + } + if err != nil { + log.Errorf("persist prompt run for session %s: %v", firstNonEmpty(input.SessionID, bindingSessionID(input.Binding)), err) + return + } + batchID := input.BatchID + if input.Binding != nil { + batchID = &input.Binding.BatchID + } + err = db.Transaction(ctx, func(tx *database.DB) error { + run, createErr := tx.CreatePromptRun(ctx, database.CreatePromptRunInput{ + SessionID: session.ID, + BatchID: batchID, + Origin: "captain", + AdmissionKey: input.RunID, + RenderedSpec: renderedSpecMap(input.Rendered), + Runtime: database.PromptRunRuntime{ + Mode: "run", + Resolved: database.PromptRunRuntimeSelection{ + Provider: ai.BackendToProvider(api.Backend(input.Backend)), Backend: input.Backend, + Model: input.Model, Effort: string(input.Rendered.Config.Model.Effort), + }, + }, + PromptMarkdown: input.Rendered.Input.Prompt.User, + }) + if createErr != nil { + return createErr + } + finished := database.PromptRunPhaseFinished + state := database.PromptRunStateSucceeded + if input.Error != "" { + state = database.PromptRunStateFailed + } + update := database.UpdatePromptRunInput{ + ID: run.ID, ExpectedVersion: run.Version, Phase: &finished, State: &state, + } + if input.ResultText != "" { + update.ResultText = &input.ResultText + } + if input.ResultJSON != nil { + update.ResultJSON = &input.ResultJSON + } + if input.Error != "" { + update.Error = &input.Error + } + _, updateErr := tx.UpdatePromptRun(ctx, update) + return updateErr + }) + if err != nil { + log.Errorf("persist prompt run for session %s: %v", firstNonEmpty(input.SessionID, bindingSessionID(input.Binding)), err) + return + } + lifecycle := database.SessionLifecycleSucceeded + if input.Error != "" { + lifecycle = database.SessionLifecycleFailed + } + updatePromptSessionLifecycle(ctx, session.ID, lifecycle, input.Error) + trackLaunchedTranscript(input, source) +} + +func bindingSessionID(binding *promptSessionBinding) string { + if binding == nil { + return "" + } + return binding.SessionID.String() +} + +// trackLaunchedTranscript arms the serve monitor's fsnotify tail on the +// freshly launched session's transcript so it ingests immediately instead of +// waiting for the next process poll or backfill. +func trackLaunchedTranscript(input promptRunRecordInput, source string) { + mon := serveMonitor() + if mon == nil { + return + } + path := historyFileForRun(api.Backend(input.Backend), input.SessionID, input.Rendered.Input.Cwd()) + if path != "" { + mon.TrackTranscript(path, source) + } +} + +// backendSource maps a prompt backend to the transcript source it produces. +func backendSource(backend api.Backend) string { + switch backend { + case api.BackendClaudeAgent, api.BackendClaudeCLI, api.BackendClaudeCmux: + return "claude" + case api.BackendCodexAgent, api.BackendCodexCLI, api.BackendCodexCmux: + return "codex" + default: + return "" + } +} + +// renderedSpecMap round-trips the realized prompt render into the jsonb shape +// stored on the prompt run. +func renderedSpecMap(rendered PromptRenderResult) map[string]any { + raw, err := json.Marshal(rendered) + if err != nil { + return map[string]any{} + } + var out map[string]any + if err := json.Unmarshal(raw, &out); err != nil { + return map[string]any{} + } + return out +} diff --git a/pkg/cli/prompt_run_persist_test.go b/pkg/cli/prompt_run_persist_test.go new file mode 100644 index 0000000..cb142ac --- /dev/null +++ b/pkg/cli/prompt_run_persist_test.go @@ -0,0 +1,69 @@ +package cli + +import ( + "testing" + + "github.com/flanksource/captain/pkg/database" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPersistPromptRunRecordsNativeRun(t *testing.T) { + db := withTestCaptainDB(t) + rendered := PromptRenderResult{Name: "fix-bug", Model: "claude-sonnet-5", Backend: "claude-agent"} + rendered.Input.Prompt.User = "fix the failing test" + batchID := uuid.New() + + persistPromptRun(t.Context(), promptRunRecordInput{ + Rendered: rendered, RunID: "run-1", SessionID: "0195c1de-4ab8-7000-8000-00000000abcd", + Model: "claude-sonnet-5", Backend: "claude-agent", BatchID: &batchID, ResultText: "done", + ResultJSON: map[string]any{"answer": "42"}, + }) + + session, err := db.GetSessionByIdentity(t.Context(), "0195c1de-4ab8-7000-8000-00000000abcd", "claude", "", "") + require.NoError(t, err) + runs, err := db.ListPromptRuns(t.Context(), database.PromptRunFilter{SessionID: &session.ID}) + require.NoError(t, err) + require.Len(t, runs, 1) + assert.Equal(t, database.PromptRunStateSucceeded, runs[0].State) + assert.Equal(t, database.PromptRunPhaseFinished, runs[0].Phase) + assert.Equal(t, "done", runs[0].ResultText) + assert.Equal(t, map[string]any{"answer": "42"}, runs[0].ResultJSON) + assert.Equal(t, "captain", runs[0].Origin) + assert.Equal(t, "run-1", runs[0].AdmissionKey) + assert.Equal(t, "fix the failing test", runs[0].PromptMarkdown) + assert.NotEmpty(t, runs[0].RenderedSpec) + require.NotNil(t, runs[0].BatchID) + assert.Equal(t, batchID, *runs[0].BatchID) + assert.Equal(t, "run", runs[0].Runtime.Mode) + assert.Equal(t, "claude-agent", runs[0].Runtime.Resolved.Provider) + assert.Equal(t, "claude-agent", runs[0].Runtime.Resolved.Backend) + assert.Equal(t, "claude-sonnet-5", runs[0].Runtime.Resolved.Model) + + t.Run("replay with the same run id is idempotent", func(t *testing.T) { + persistPromptRun(t.Context(), promptRunRecordInput{ + Rendered: rendered, RunID: "run-1", SessionID: "0195c1de-4ab8-7000-8000-00000000abcd", + Model: "claude-sonnet-5", Backend: "claude-agent", ResultText: "done", + }) + runs, err := db.ListPromptRuns(t.Context(), database.PromptRunFilter{SessionID: &session.ID}) + require.NoError(t, err) + assert.Len(t, runs, 1) + }) + + t.Run("failed run records error and failed state", func(t *testing.T) { + persistPromptRun(t.Context(), promptRunRecordInput{ + Rendered: rendered, RunID: "run-2", SessionID: "0195c1de-4ab8-7000-8000-00000000abcd", + Model: "claude-sonnet-5", Backend: "claude-agent", Error: "verify failed: tests red", + }) + runs, err := db.ListPromptRuns(t.Context(), database.PromptRunFilter{SessionID: &session.ID}) + require.NoError(t, err) + require.Len(t, runs, 2) + failed := runs[0] + if failed.AdmissionKey != "run-2" { + failed = runs[1] + } + assert.Equal(t, database.PromptRunStateFailed, failed.State) + assert.Equal(t, "verify failed: tests red", failed.Error) + }) +} diff --git a/pkg/cli/prompt_run_pretty_ginkgo_test.go b/pkg/cli/prompt_run_pretty_ginkgo_test.go new file mode 100644 index 0000000..72046f8 --- /dev/null +++ b/pkg/cli/prompt_run_pretty_ginkgo_test.go @@ -0,0 +1,48 @@ +package cli + +import ( + "strings" + + clickyapi "github.com/flanksource/clicky/api" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("multi-model prompt output", func() { + It("renders metadata in the table and full responses after it", func() { + pretty := (PromptRunResult{ + Status: "completed", Total: 2, Succeeded: 2, + Runs: []PromptRunItem{ + {Selector: "api:sonnet-5", Status: "completed", Text: "first full response"}, + {Selector: "cmux:opus", Status: "completed", Text: "second full response"}, + }, + }).Pretty() + + tableIndex := -1 + var table clickyapi.TextTable + for i, child := range pretty.Children { + if candidate, ok := child.(clickyapi.TextTable); ok { + tableIndex = i + table = candidate + break + } + } + Expect(tableIndex).To(BeNumerically(">=", 0)) + metrics := make([]string, 0, len(table.Rows)) + for _, row := range table.Rows { + metrics = append(metrics, row["metric"].String()) + } + Expect(metrics).NotTo(ContainElement("Response")) + + var afterTable strings.Builder + for _, child := range pretty.Children[tableIndex+1:] { + afterTable.WriteString(child.String()) + } + output := afterTable.String() + Expect(output).To(ContainSubstring("api:sonnet-5")) + Expect(output).To(ContainSubstring("first full response")) + Expect(output).To(ContainSubstring("cmux:opus")) + Expect(output).To(ContainSubstring("second full response")) + Expect(strings.Index(output, "first full response")).To(BeNumerically("<", strings.Index(output, "second full response"))) + }) +}) diff --git a/pkg/cli/prompt_run_result.go b/pkg/cli/prompt_run_result.go new file mode 100644 index 0000000..ac8b026 --- /dev/null +++ b/pkg/cli/prompt_run_result.go @@ -0,0 +1,175 @@ +package cli + +import ( + "fmt" + "strings" + + clickyapi "github.com/flanksource/clicky/api" +) + +// PromptRunResult is the unified result of the "run" action. Over HTTP (serve) +// it carries the async handle; on the CLI it carries the synchronous result. +type PromptRunResult struct { + RunID string `json:"runId,omitempty"` + BatchID string `json:"batchId,omitempty"` + Status string `json:"status,omitempty" pretty:"label=Status"` + Model string `json:"model,omitempty" pretty:"label=Model"` + Backend string `json:"backend,omitempty" pretty:"label=Backend"` + Chat bool `json:"chat,omitempty"` + Capabilities ChatCapabilities `json:"capabilities,omitempty"` + + Text string `json:"text,omitempty" pretty:"label=Response"` + StructuredOutput map[string]any `json:"structuredOutput,omitempty" pretty:"-"` + SessionID string `json:"sessionId,omitempty" pretty:"label=Session"` + Dir string `json:"dir,omitempty" pretty:"label=Dir"` + HistoryFile string `json:"historyFile,omitempty" pretty:"label=History"` + InputTokens int `json:"inputTokens,omitempty" pretty:"label=Input Tokens"` + OutputTokens int `json:"outputTokens,omitempty" pretty:"label=Output Tokens"` + CostUSD float64 `json:"costUSD,omitempty" pretty:"label=Cost USD"` + Duration string `json:"duration,omitempty" pretty:"label=Duration"` + + Total int `json:"total,omitempty" pretty:"label=Total"` + Succeeded int `json:"succeeded,omitempty" pretty:"label=Succeeded"` + Failed int `json:"failed,omitempty" pretty:"label=Failed"` + Runs []PromptRunItem `json:"runs,omitempty" pretty:"label=Runs"` +} + +type PromptRunItem struct { + RunID string `json:"runId,omitempty" pretty:"label=Run"` + Selector string `json:"selector,omitempty" pretty:"label=Selector"` + Status string `json:"status,omitempty" pretty:"label=Status"` + Model string `json:"model,omitempty" pretty:"label=Model"` + Backend string `json:"backend,omitempty" pretty:"label=Backend"` + Effort string `json:"effort,omitempty" pretty:"label=Effort"` + Chat bool `json:"chat,omitempty"` + Capabilities ChatCapabilities `json:"capabilities,omitempty"` + Text string `json:"text,omitempty" pretty:"label=Response"` + StructuredOutput map[string]any `json:"structuredOutput,omitempty" pretty:"-"` + SessionID string `json:"sessionId,omitempty" pretty:"label=Session"` + Dir string `json:"dir,omitempty" pretty:"label=Dir"` + HistoryFile string `json:"historyFile,omitempty" pretty:"label=History"` + InputTokens int `json:"inputTokens,omitempty" pretty:"label=Input Tokens"` + OutputTokens int `json:"outputTokens,omitempty" pretty:"label=Output Tokens"` + CostUSD float64 `json:"costUSD,omitempty" pretty:"label=Cost USD"` + Duration string `json:"duration,omitempty" pretty:"label=Duration"` + Error string `json:"error,omitempty" pretty:"label=Error"` +} + +func (r PromptRunResult) Pretty() clickyapi.Text { + if len(r.Runs) == 0 { + if r.Text != "" { + return clickyapi.Text{Content: r.Text} + } + return clickyapi.Text{Content: r.Status} + } + t := clickyapi.Text{}. + Append(fmt.Sprintf("Status: %s Total: %d Succeeded: %d Failed: %d Duration: %s", + r.Status, r.Total, r.Succeeded, r.Failed, r.Duration), "font-medium") + + t = t.NewLine().Add(promptRunComparisonTable(r.Runs)) + for _, run := range r.Runs { + if strings.TrimSpace(run.Text) == "" { + continue + } + t = t.NewLine().NewLine(). + Append("Response — ", "text-gray-500"). + Append(runColumnHeader(run), "font-bold"). + NewLine(). + Append(run.Text) + } + return t +} + +func promptRunComparisonTable(runs []PromptRunItem) clickyapi.TextTable { + table := clickyapi.TextTable{ + Headers: clickyapi.TextList{textCell("Metric")}, + FieldNames: []string{"metric"}, + } + for i, run := range runs { + field := runColumnField(i) + table.FieldNames = append(table.FieldNames, field) + table.Headers = append(table.Headers, textCell(runColumnHeader(run))) + } + + add := func(metric string, values func(PromptRunItem) string) { + row := clickyapi.TableRow{"metric": cell(metric)} + for i, run := range runs { + row[runColumnField(i)] = cell(values(run)) + } + table.Rows = append(table.Rows, row) + } + add("Status", func(run PromptRunItem) string { return run.Status }) + add("Backend", func(run PromptRunItem) string { return run.Backend }) + add("Model", func(run PromptRunItem) string { return run.Model }) + add("Error", func(run PromptRunItem) string { return truncateCell(run.Error, 160) }) + add("Duration", func(run PromptRunItem) string { return run.Duration }) + add("Tokens", func(run PromptRunItem) string { return tokenCell(run.InputTokens, run.OutputTokens) }) + add("Cost", func(run PromptRunItem) string { return costCell(run.CostUSD) }) + add("Session", func(run PromptRunItem) string { return shortSessionCell(run.SessionID) }) + add("History", func(run PromptRunItem) string { return truncatePathCell(run.HistoryFile, 72) }) + add("Dir", func(run PromptRunItem) string { return truncatePathCell(run.Dir, 56) }) + return table +} + +func runColumnField(index int) string { + return fmt.Sprintf("run%d", index+1) +} + +func runColumnHeader(run PromptRunItem) string { + if strings.TrimSpace(run.Selector) != "" { + return run.Selector + } + if run.Backend != "" && run.Model != "" { + return run.Backend + ":" + run.Model + } + return firstNonEmpty(run.Model, run.Backend, "run") +} + +func textCell(s string) clickyapi.Textable { + return clickyapi.Text{Content: s} +} + +func cell(s string) clickyapi.TypedValue { + return clickyapi.TypedValue{Textable: clickyapi.Text{Content: s}} +} + +func truncateCell(s string, max int) string { + s = strings.TrimSpace(s) + if len(s) <= max { + return s + } + return s[:max] + "..." +} + +func truncatePathCell(s string, max int) string { + s = strings.TrimSpace(s) + if len(s) <= max { + return s + } + if max <= 3 { + return s[:max] + } + return "..." + s[len(s)-max+3:] +} + +func shortSessionCell(s string) string { + s = strings.TrimSpace(s) + if len(s) <= 12 { + return s + } + return s[:12] +} + +func tokenCell(input, output int) string { + if input == 0 && output == 0 { + return "" + } + return fmt.Sprintf("%d/%d", input, output) +} + +func costCell(cost float64) string { + if cost <= 0 { + return "" + } + return fmt.Sprintf("$%.4f", cost) +} diff --git a/pkg/cli/prompt_run_stream.go b/pkg/cli/prompt_run_stream.go index 9981c85..1ab8b44 100644 --- a/pkg/cli/prompt_run_stream.go +++ b/pkg/cli/prompt_run_stream.go @@ -19,17 +19,30 @@ const runSubBuffer = 64 // frames. Every frame is buffered for replay to late/reconnecting subscribers // and fanned out to current subscribers. type runStream struct { - mu sync.Mutex - entries []session.Message - subs map[chan session.Message]struct{} - done bool - summary *PromptRunSummary - errMsg string - endedAt time.Time + mu sync.Mutex + entries []session.Message + subs map[chan session.Message]struct{} + eventSubs map[chan runStreamEvent]struct{} + run PromptRunFrame + chatState *ChatStateFrame + cancel context.CancelFunc + stopRequested bool + done bool + summary *PromptRunSummary + errMsg string + endedAt time.Time } func newRunStream() *runStream { - return &runStream{subs: map[chan session.Message]struct{}{}} + return &runStream{ + subs: map[chan session.Message]struct{}{}, + eventSubs: map[chan runStreamEvent]struct{}{}, + } +} + +type runStreamEvent struct { + name string + data any } // publish appends a frame and fans it out. A subscriber whose buffer is full is @@ -49,6 +62,68 @@ func (s *runStream) publish(e session.Message) { close(ch) } } + s.publishEventLocked(runStreamEvent{name: "entry", data: e}) +} + +func (s *runStream) setRun(frame PromptRunFrame) { + s.mu.Lock() + s.run = frame + s.publishEventLocked(runStreamEvent{name: "run", data: frame}) + s.mu.Unlock() +} + +func (s *runStream) setChatState(frame ChatStateFrame) { + s.mu.Lock() + copy := frame + copy.Queued = append([]ChatQueuedMessage(nil), frame.Queued...) + copy.DiscardedMessageIDs = append([]string(nil), frame.DiscardedMessageIDs...) + s.chatState = © + s.run.SessionID = frame.SessionID + s.run.Capabilities = frame.Capabilities + s.publishEventLocked(runStreamEvent{name: "state", data: copy}) + s.mu.Unlock() +} + +func (s *runStream) setCancel(cancel context.CancelFunc) { + s.mu.Lock() + s.cancel = cancel + requested := s.stopRequested + s.mu.Unlock() + if requested { + cancel() + } +} + +func (s *runStream) requestStop() bool { + s.mu.Lock() + if s.done { + s.mu.Unlock() + return false + } + s.stopRequested = true + cancel := s.cancel + s.mu.Unlock() + if cancel != nil { + cancel() + } + return true +} + +func (s *runStream) wasStopped() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.stopRequested +} + +func (s *runStream) publishEventLocked(event runStreamEvent) { + for ch := range s.eventSubs { + select { + case ch <- event: + default: + delete(s.eventSubs, ch) + close(ch) + } + } } // complete marks the run finished successfully and closes all subscribers. @@ -67,10 +142,22 @@ func (s *runStream) finish(sum *PromptRunSummary, errMsg string) { s.summary = sum s.errMsg = errMsg s.endedAt = time.Now() + if errMsg != "" { + s.run.Status = "error" + } else { + s.run.Status = "done" + } + if sum != nil && sum.SessionID != "" { + s.run.SessionID = sum.SessionID + } for ch := range s.subs { delete(s.subs, ch) close(ch) } + for ch := range s.eventSubs { + delete(s.eventSubs, ch) + close(ch) + } } // subscribe atomically snapshots the replay buffer and registers a live channel @@ -88,14 +175,27 @@ func (s *runStream) subscribe() (replay []session.Message, ch chan session.Messa return replay, ch, false, nil, "" } -func (s *runStream) unsubscribe(ch chan session.Message) { +func (s *runStream) subscribeEvents() (PromptRunFrame, []session.Message, *ChatStateFrame, chan runStreamEvent, bool, *PromptRunSummary, string) { + s.mu.Lock() + defer s.mu.Unlock() + replay := append([]session.Message(nil), s.entries...) + state := cloneChatState(s.chatState) + if s.done { + return s.run, replay, state, nil, true, s.summary, s.errMsg + } + ch := make(chan runStreamEvent, runSubBuffer) + s.eventSubs[ch] = struct{}{} + return s.run, replay, state, ch, false, nil, "" +} + +func (s *runStream) unsubscribeEvents(ch chan runStreamEvent) { if ch == nil { return } s.mu.Lock() defer s.mu.Unlock() - if _, ok := s.subs[ch]; ok { - delete(s.subs, ch) + if _, ok := s.eventSubs[ch]; ok { + delete(s.eventSubs, ch) close(ch) } } @@ -170,7 +270,9 @@ func (b *runBroker) prune(maxAge time.Duration) { } type promptRunSnapshotBody struct { + Run PromptRunFrame `json:"run"` Entries []session.Message `json:"entries"` + State *ChatStateFrame `json:"state,omitempty"` Done bool `json:"done"` Summary *PromptRunSummary `json:"summary,omitempty"` Error string `json:"error,omitempty"` @@ -195,11 +297,15 @@ func handlePromptRunStream(b *runBroker) http.HandlerFunc { } setSSEHeaders(w) - replay, ch, done, summary, errMsg := stream.subscribe() - defer stream.unsubscribe(ch) + run, replay, state, ch, done, summary, errMsg := stream.subscribeEvents() + defer stream.unsubscribeEvents(ch) + writeSSE(w, "run", run) for _, e := range replay { writeSSE(w, "entry", e) } + if state != nil { + writeSSE(w, "state", state) + } if done { writeTerminal(w, summary, errMsg) flusher.Flush() @@ -213,7 +319,7 @@ func handlePromptRunStream(b *runBroker) http.HandlerFunc { select { case <-r.Context().Done(): return - case e, open := <-ch: + case event, open := <-ch: if !open { if d, sum, em := stream.state(); d { writeTerminal(w, sum, em) @@ -221,7 +327,7 @@ func handlePromptRunStream(b *runBroker) http.HandlerFunc { } return } - writeSSE(w, "entry", e) + writeSSE(w, event.name, event.data) flusher.Flush() case <-heartbeat.C: fmt.Fprint(w, ": ping\n\n") @@ -240,10 +346,21 @@ func handlePromptRunSnapshot(b *runBroker) http.HandlerFunc { http.Error(w, "unknown run", http.StatusNotFound) return } - entries, done, summary, errMsg := stream.snapshot() + run, entries, state, ch, done, summary, errMsg := stream.subscribeEvents() + stream.unsubscribeEvents(ch) w.Header().Set("Content-Type", "application/json") - _ = json.NewEncoder(w).Encode(promptRunSnapshotBody{Entries: entries, Done: done, Summary: summary, Error: errMsg}) + _ = json.NewEncoder(w).Encode(promptRunSnapshotBody{Run: run, Entries: entries, State: state, Done: done, Summary: summary, Error: errMsg}) + } +} + +func cloneChatState(state *ChatStateFrame) *ChatStateFrame { + if state == nil { + return nil } + copy := *state + copy.Queued = append([]ChatQueuedMessage(nil), state.Queued...) + copy.DiscardedMessageIDs = append([]string(nil), state.DiscardedMessageIDs...) + return © } func setSSEHeaders(w http.ResponseWriter) { diff --git a/pkg/cli/prompt_run_test.go b/pkg/cli/prompt_run_test.go new file mode 100644 index 0000000..5fe00ca --- /dev/null +++ b/pkg/cli/prompt_run_test.go @@ -0,0 +1,305 @@ +package cli + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + clickyapi "github.com/flanksource/clicky/api" +) + +func TestExecuteSyncRunMultiModelsParallel(t *testing.T) { + withTestCaptainDB(t) + + old := executePromptRequestFunc + t.Cleanup(func() { executePromptRequestFunc = old }) + + started := make(chan struct{}) + var calls atomic.Int32 + executePromptRequestFunc = func(ctx context.Context, req ai.Request, cfg ai.Config, _ time.Duration, noStream bool) (any, error) { + if noStream { + t.Errorf("multi-model run should preserve streaming unless --no-stream is set") + } + if calls.Add(1) == 2 { + close(started) + } + select { + case <-started: + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(time.Second): + return nil, errors.New("executor was not called concurrently") + } + return AIPromptResult{ + Text: cfg.Model.Name, + StructuredOutput: map[string]any{"model": cfg.Model.Name}, + Model: cfg.Model.Name, + Backend: string(cfg.Model.Backend), + Dir: req.Cwd(), + SessionID: "session-" + cfg.Model.Name, + HistoryFile: filepath.Join(req.Cwd(), ".history", cfg.Model.Name+".jsonl"), + InputTokens: 1, + Output: 2, + CostUSD: 0.01, + Duration: "1ms", + }, nil + } + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + rendered := testRenderedPrompt(api.Model{Name: "claude-sonnet-5", Backend: api.BackendAnthropic}) + rendered.Input.SetCwd(t.TempDir()) + got, err := executeSyncRun(ctx, rendered, AIPromptOptions{MultiModels: []string{"cli:sonnet-5,cmux:opus"}}) + if err != nil { + t.Fatalf("executeSyncRun: %v", err) + } + if got.Status != "completed" || got.Total != 2 || got.Succeeded != 2 || got.Failed != 0 { + t.Fatalf("batch status = %s total=%d succeeded=%d failed=%d", got.Status, got.Total, got.Succeeded, got.Failed) + } + want := []struct { + backend api.Backend + model string + }{ + {api.BackendClaudeCLI, "claude-sonnet-5"}, + {api.BackendClaudeCmux, "claude-opus-5"}, + } + for i, w := range want { + if got.Runs[i].Backend != string(w.backend) || got.Runs[i].Model != w.model { + t.Fatalf("run[%d] = %s/%s, want %s/%s", i, got.Runs[i].Backend, got.Runs[i].Model, w.backend, w.model) + } + } + if got.InputTokens != 2 || got.OutputTokens != 4 { + t.Fatalf("aggregate tokens = %d/%d, want 2/4", got.InputTokens, got.OutputTokens) + } + if got.Runs[0].Dir == "" || !strings.Contains(got.Runs[0].HistoryFile, ".history") { + t.Fatalf("run metadata missing dir/history: %+v", got.Runs[0]) + } + if got.Runs[0].CostUSD != 0.01 || got.Runs[0].InputTokens != 1 || got.Runs[0].OutputTokens != 2 { + t.Fatalf("run usage/cost missing: %+v", got.Runs[0]) + } + if got.Runs[0].StructuredOutput["model"] != got.Runs[0].Model { + t.Fatalf("run structured output missing: %+v", got.Runs[0]) + } +} + +func TestExecuteSyncRunMultiModelsHonorsNoStream(t *testing.T) { + old := executePromptRequestFunc + t.Cleanup(func() { executePromptRequestFunc = old }) + + executePromptRequestFunc = func(_ context.Context, _ ai.Request, cfg ai.Config, _ time.Duration, noStream bool) (any, error) { + if !noStream { + t.Errorf("multi-model run should pass explicit --no-stream through") + } + return AIPromptResult{ + Text: "ok", + Model: cfg.Model.Name, + Backend: string(cfg.Model.Backend), + Duration: "1ms", + }, nil + } + + rendered := testRenderedPrompt(api.Model{Name: "claude-sonnet-5", Backend: api.BackendAnthropic}) + got, err := executeSyncRun(context.Background(), rendered, AIPromptOptions{MultiModels: []string{"cli:sonnet-5"}, NoStream: true}) + if err != nil { + t.Fatalf("executeSyncRun: %v", err) + } + if got.Status != "completed" || got.Succeeded != 1 { + t.Fatalf("batch status = %+v", got) + } +} + +func TestExecuteSyncRunMultiModelsPartialFailure(t *testing.T) { + withTestCaptainDB(t) + + old := executePromptRequestFunc + t.Cleanup(func() { executePromptRequestFunc = old }) + + executePromptRequestFunc = func(_ context.Context, _ ai.Request, cfg ai.Config, _ time.Duration, _ bool) (any, error) { + if cfg.Model.Backend == api.BackendCodexCmux { + return nil, errors.New("cmux unavailable") + } + return AIPromptResult{ + Text: "ok", + Model: cfg.Model.Name, + Backend: string(cfg.Model.Backend), + InputTokens: 3, + Output: 4, + Duration: "2ms", + }, nil + } + + rendered := testRenderedPrompt(api.Model{Name: "gpt-5.5", Backend: api.BackendOpenAI}) + got, err := executeSyncRun(context.Background(), rendered, AIPromptOptions{MultiModels: []string{"api:gpt-5.5,cmux:gpt-5.5"}}) + if err != nil { + t.Fatalf("executeSyncRun: %v", err) + } + if got.Status != "partial" || got.Succeeded != 1 || got.Failed != 1 { + t.Fatalf("batch status = %s succeeded=%d failed=%d", got.Status, got.Succeeded, got.Failed) + } + if got.Runs[1].Status != "failed" || !strings.Contains(got.Runs[1].Error, "cmux unavailable") { + t.Fatalf("failed run = %+v", got.Runs[1]) + } +} + +func TestExecuteSyncRunMultiModelsRejectsResume(t *testing.T) { + rendered := testRenderedPrompt(api.Model{Name: "gpt-5.5", Backend: api.BackendOpenAI}) + rendered.Input.SessionID = "session-1" + _, err := executeSyncRun(context.Background(), rendered, AIPromptOptions{MultiModels: []string{"api:gpt-5.5,cmux:gpt-5.5"}}) + if err == nil || !strings.Contains(err.Error(), "--resume") { + t.Fatalf("error = %v, want --resume rejection", err) + } +} + +func TestVariantModelUsesSelectorEffort(t *testing.T) { + selector := api.Model{Name: "gpt-5.6-terra", Backend: api.BackendCodexCmux, Effort: api.EffortUltra} + got := variantModel(selector, nil) + if got.Name != selector.Name || got.Backend != selector.Backend || got.Effort != api.EffortUltra { + t.Fatalf("variant = %+v, want selector model/backend/effort", got) + } +} + +func TestPromptRunResultPrettyRendersRunsAsColumns(t *testing.T) { + result := PromptRunResult{ + Status: "partial", + Total: 2, + Succeeded: 1, + Failed: 1, + Runs: []PromptRunItem{ + { + Selector: "api:sonnet-5", + Status: "completed", + Backend: "anthropic", + Model: "claude-sonnet-5", + Dir: "/repo", + SessionID: "0123456789abcdef", + HistoryFile: "/repo/.claude/session.jsonl", + InputTokens: 8, + OutputTokens: 14, + CostUSD: 0.0002, + Duration: "1s", + Text: "api ok", + }, + { + Selector: "cmux:opus", + Status: "failed", + Backend: "claude-cmux", + Model: "claude-opus-4-8", + Duration: "2s", + Error: "cmux unavailable", + }, + }, + } + + pretty := result.Pretty() + table := firstPrettyTable(t, pretty) + wantFields := []string{"metric", "run1", "run2"} + if strings.Join(table.FieldNames, ",") != strings.Join(wantFields, ",") { + t.Fatalf("field names = %v, want %v", table.FieldNames, wantFields) + } + if len(table.Headers) != 3 { + t.Fatalf("headers = %d, want 3", len(table.Headers)) + } + wantHeaders := []string{"Metric", "api:sonnet-5", "cmux:opus"} + for i, want := range wantHeaders { + if got := table.Headers[i].String(); got != want { + t.Fatalf("header[%d] = %q, want %q", i, got, want) + } + } + + tests := []struct { + metric string + run1 string + run2 string + }{ + {"Status", "completed", "failed"}, + {"Backend", "anthropic", "claude-cmux"}, + {"Model", "claude-sonnet-5", "claude-opus-4-8"}, + {"Error", "", "cmux unavailable"}, + {"Duration", "1s", "2s"}, + {"Tokens", "8/14", ""}, + {"Cost", "$0.0002", ""}, + {"Session", "0123456789ab", ""}, + {"History", "/repo/.claude/session.jsonl", ""}, + {"Dir", "/repo", ""}, + } + for _, tt := range tests { + row := tableRowByMetric(t, table, tt.metric) + if row["run1"].String() != tt.run1 || row["run2"].String() != tt.run2 { + t.Fatalf("%s row = %q / %q, want %q / %q", tt.metric, row["run1"].String(), row["run2"].String(), tt.run1, tt.run2) + } + } + if output := pretty.String(); !strings.Contains(output, "Response — api:sonnet-5\napi ok") { + t.Fatalf("response block missing after metadata table: %q", output) + } +} + +func TestHistoryFileForRun(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + cwd := filepath.Join(home, "repo") + if err := os.MkdirAll(cwd, 0o755); err != nil { + t.Fatalf("mkdir cwd: %v", err) + } + + got := historyFileForRun(api.BackendClaudeCLI, "claude-session", cwd) + if want := filepath.Join(home, ".claude", "projects"); !strings.HasPrefix(got, want) || !strings.HasSuffix(got, "claude-session.jsonl") { + t.Fatalf("claude history path = %q, want under %q", got, want) + } + + codexDir := filepath.Join(home, ".codex", "sessions", "2026", "07", "08") + if err := os.MkdirAll(codexDir, 0o755); err != nil { + t.Fatalf("mkdir codex dir: %v", err) + } + codexFile := filepath.Join(codexDir, "codex-session.jsonl") + if err := os.WriteFile(codexFile, []byte(`{"type":"session_meta","payload":{"id":"codex-session","cwd":"`+cwd+`"}}`+"\n"), 0o644); err != nil { + t.Fatalf("write codex session: %v", err) + } + if got := historyFileForRun(api.BackendCodexAgent, "codex-session", cwd); got != codexFile { + t.Fatalf("codex history path = %q, want %q", got, codexFile) + } +} + +func firstPrettyTable(t *testing.T, text clickyapi.Text) clickyapi.TextTable { + t.Helper() + for _, child := range text.Children { + if table, ok := child.(clickyapi.TextTable); ok { + return table + } + } + t.Fatalf("no table child in %#v", text.Children) + return clickyapi.TextTable{} +} + +func tableRowByMetric(t *testing.T, table clickyapi.TextTable, metric string) clickyapi.TableRow { + t.Helper() + for _, row := range table.Rows { + if row["metric"].String() == metric { + return row + } + } + t.Fatalf("no %q metric row in %#v", metric, table.Rows) + return nil +} + +func testRenderedPrompt(model api.Model) PromptRenderResult { + req := ai.Request{ + Model: model, + Prompt: api.Prompt{User: "hello"}, + Budget: api.Budget{Timeout: "1s"}, + } + cfg := ai.Config{Model: model} + return PromptRenderResult{ + Name: "test", + Model: model.Name, + Backend: string(model.Backend), + Input: req, + Config: cfg, + } +} diff --git a/pkg/cli/prompt_runtimes_ginkgo_test.go b/pkg/cli/prompt_runtimes_ginkgo_test.go new file mode 100644 index 0000000..7c4bab1 --- /dev/null +++ b/pkg/cli/prompt_runtimes_ginkgo_test.go @@ -0,0 +1,168 @@ +package cli + +import ( + "context" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("prompt-declared runtimes", func() { + var path string + + BeforeEach(func() { + configPath := filepath.Join(GinkgoT().TempDir(), ".captain.yaml") + captainconfig.SetPathForTesting(configPath) + DeferCleanup(func() { captainconfig.SetPathForTesting("") }) + path = filepath.Join(GinkgoT().TempDir(), "compare.prompt") + Expect(os.WriteFile(path, []byte(`--- +name: Compare UI audits +model: gemini-3.5-flash +runtimes: + - api:gemini-3.5-flash:high + - model: claude-sonnet-5 + backend: anthropic + effort: medium +--- +{{role "user"}} +Review the screenshot. +`), 0o600)).To(Succeed()) + }) + + It("renders root runtimes as resolved parallel defaults", func() { + rendered, err := renderPromptCLI(context.Background(), path, AIPromptOptions{}, "", "") + + Expect(err).NotTo(HaveOccurred()) + Expect(rendered.Runtimes).To(HaveLen(2)) + Expect(rendered.Runtimes[0]).To(SatisfyAll( + HaveField("Name", "gemini-3.5-flash"), + HaveField("Backend", api.BackendGemini), + HaveField("Effort", api.EffortHigh), + )) + Expect(rendered.Runtimes[1]).To(SatisfyAll( + HaveField("Name", "claude-sonnet-5"), + HaveField("Backend", api.BackendAnthropic), + HaveField("Effort", api.EffortMedium), + )) + }) + + DescribeTable("resolves a discovered prompt by bare filename", + func(id string) { + ctx := ContextWithPromptDirs(context.Background(), []string{filepath.Dir(path)}) + expectedPath, err := filepath.EvalSymlinks(path) + Expect(err).NotTo(HaveOccurred()) + + record, err := resolvePromptRecord(ctx, id) + + Expect(err).NotTo(HaveOccurred()) + Expect(displayPromptPath(record)).To(Equal(expectedPath)) + Expect(record.Rel).To(Equal("compare.prompt")) + }, + Entry("without the extension", "compare"), + Entry("with the extension", "compare.prompt"), + ) + + It("rejects an ambiguous bare filename", func() { + otherDir := GinkgoT().TempDir() + otherPath := filepath.Join(otherDir, filepath.Base(path)) + Expect(os.WriteFile(otherPath, []byte("Review this screenshot."), 0o600)).To(Succeed()) + ctx := ContextWithPromptDirs(context.Background(), []string{filepath.Dir(path), otherDir}) + + _, err := resolvePromptRecord(ctx, "compare") + + Expect(err).To(MatchError(ContainSubstring("prompt name \"compare\" is ambiguous"))) + }) + + It("does not require a singular base model", func() { + content, err := os.ReadFile(path) + Expect(err).NotTo(HaveOccurred()) + Expect(os.WriteFile(path, []byte(strings.Replace(string(content), "model: gemini-3.5-flash\n", "", 1)), 0o600)).To(Succeed()) + + rendered, err := renderPromptCLI(context.Background(), path, AIPromptOptions{}, "", "") + + Expect(err).NotTo(HaveOccurred()) + Expect(rendered.ValidationError).To(BeEmpty()) + Expect(rendered.Runtimes).To(HaveLen(2)) + }) + + It("lets explicit CLI runtimes replace prompt defaults", func() { + rendered, err := renderPromptCLI(context.Background(), path, AIPromptOptions{ + MultiModels: []string{"api:gpt-5.5:high,agent:sol:medium"}, + }, "", "") + + Expect(err).NotTo(HaveOccurred()) + Expect(rendered.Runtimes).To(HaveLen(2)) + Expect(rendered.Runtimes[0].Name).To(Equal("gpt-5.5")) + Expect(rendered.Runtimes[1].Backend).To(Equal(api.BackendCodexAgent)) + }) + + It("lets explicit HTTP runtimes replace prompt defaults", func() { + record, err := filePromptRecord(path) + Expect(err).NotTo(HaveOccurred()) + rendered, err := renderPrompt(context.Background(), record.ID, PromptRenderRequest{Runtimes: []api.Model{ + {Name: "gpt-5.5", Backend: api.BackendOpenAI, Effort: api.EffortHigh}, + {Name: "gpt-5.6-sol", Backend: api.BackendCodexAgent, Effort: api.EffortMedium}, + }}) + + Expect(err).NotTo(HaveOccurred()) + Expect(rendered.Runtimes).To(HaveLen(2)) + Expect(rendered.Runtimes[0].Backend).To(Equal(api.BackendOpenAI)) + Expect(rendered.Runtimes[1].Backend).To(Equal(api.BackendCodexAgent)) + }) + + DescribeTable("executes name-resolved runtimes without CLI selectors", + func(id string) { + withGinkgoCaptainDB() + + originalExecute := executePromptRequestFunc + DeferCleanup(func() { executePromptRequestFunc = originalExecute }) + var mu sync.Mutex + executed := []api.Model{} + executePromptRequestFunc = func(_ context.Context, _ ai.Request, cfg ai.Config, _ time.Duration, _ bool) (any, error) { + mu.Lock() + defer mu.Unlock() + executed = append(executed, cfg.Model) + return AIPromptResult{Text: "ok", Model: cfg.Model.Name, Backend: string(cfg.Model.Backend)}, nil + } + ctx := ContextWithPromptDirs(context.Background(), []string{filepath.Dir(path)}) + rendered, err := renderPromptCLI(ctx, id, AIPromptOptions{}, "", "") + Expect(err).NotTo(HaveOccurred()) + + result, err := executeSyncRun(ctx, rendered, AIPromptOptions{}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.Runs).To(HaveLen(2)) + Expect(executed).To(ConsistOf( + SatisfyAll(HaveField("Name", "gemini-3.5-flash"), HaveField("Backend", api.BackendGemini)), + SatisfyAll(HaveField("Name", "claude-sonnet-5"), HaveField("Backend", api.BackendAnthropic)), + )) + }, + Entry("without the extension", "compare"), + Entry("with the extension", "compare.prompt"), + ) + + It("rejects unknown fields in prompt runtimes", func() { + invalidPath := filepath.Join(GinkgoT().TempDir(), "invalid.prompt") + Expect(os.WriteFile(invalidPath, []byte(`--- +runtimes: + - model: gemini-3.5-flash + backend: gemini + typo: true + - api:sonnet-5:medium +--- +Review the screenshot. +`), 0o600)).To(Succeed()) + + _, err := renderPromptCLI(context.Background(), invalidPath, AIPromptOptions{}, "", "") + + Expect(err).To(MatchError(ContainSubstring("field typo not found"))) + }) +}) diff --git a/pkg/cli/prompt_schema.go b/pkg/cli/prompt_schema.go index ee226ed..ac6f5a2 100644 --- a/pkg/cli/prompt_schema.go +++ b/pkg/cli/prompt_schema.go @@ -66,46 +66,21 @@ func WritePromptSchema(w io.Writer) error { // PromptSchemaDocument assembles the prompt/spec editor schema document from the // live runtime: the cmux "extra args" reflected from Go structs (cached, since // they never change within a process) and the available backends/models probed -// from the environment via whoami (cached with a short TTL so a long-running -// serve process reflects key/model changes without re-probing per request). +// from the environment (ai.CachedAdapters owns the short-TTL cache so a +// long-running serve process reflects key/model changes without re-probing per +// request). func PromptSchemaDocument() (map[string]any, error) { - adapters, err := cachedSchemaAdapters(time.Now()) + adapters, err := schemaAdapters() if err != nil { return nil, err } return buildPromptSchemaDocument(adapters) } -// probeSchemaAdapters is the live probe used to source backends and models. It is -// a package var so tests can substitute a deterministic, network-free stub. -var probeSchemaAdapters = func() ([]AdapterStatus, error) { - return ProbeAdapters(WhoamiOptions{Models: true}, osAuthProbe()) -} - -const schemaAdapterCacheTTL = 60 * time.Second - -var ( - schemaAdapterMu sync.Mutex - schemaAdapterCache []AdapterStatus - schemaAdapterAt time.Time -) - -// cachedSchemaAdapters returns the probed adapters, reusing a cached probe within -// the TTL. A probe error is never cached: the next call retries so a transient -// failure does not permanently empty the schema. -func cachedSchemaAdapters(now time.Time) ([]AdapterStatus, error) { - schemaAdapterMu.Lock() - defer schemaAdapterMu.Unlock() - if schemaAdapterCache != nil && now.Sub(schemaAdapterAt) < schemaAdapterCacheTTL { - return schemaAdapterCache, nil - } - adapters, err := probeSchemaAdapters() - if err != nil { - return nil, err - } - schemaAdapterCache = adapters - schemaAdapterAt = now - return adapters, nil +// schemaAdapters sources the probed adapters through pkg/ai's cache. It is a +// package var so tests can substitute a deterministic, network-free stub. +var schemaAdapters = func() ([]AdapterStatus, error) { + return ai.CachedAdapters(time.Now()) } // reflectedSchemaBytes holds the JSON of the reflection-derived schemas. They are @@ -153,11 +128,6 @@ func promptSchemaExampleSpec() map[string]any { "user": "Update the prompt runtime editor and keep the payload compact.", "system": "Answer with direct implementation guidance.", "appendSystem": "Prefer existing clicky-ui primitives.", - "source": "storybook/demo.prompt", - "metadata": map[string]any{ - "surface": "prompt-workbench", - "owner": "captain", - }, }, "budget": map[string]any{ "cost": 0.5, @@ -192,9 +162,9 @@ func promptSchemaExampleSpec() map[string]any { "cwd": ".", "baseDir": ".shell", "dotenv": []string{".env", ".env.local"}, - "env": []string{ - "CAPTAIN_RUNTIME=storybook", - "GAVEL_PROFILE=local", + "envVars": []map[string]any{ + {"name": "CAPTAIN_RUNTIME", "value": "storybook"}, + {"name": "GAVEL_PROFILE", "value": "local"}, }, }, "cliArgs": map[string]any{ diff --git a/pkg/cli/prompt_schema_build.go b/pkg/cli/prompt_schema_build.go index 84222fe..96d5ad1 100644 --- a/pkg/cli/prompt_schema_build.go +++ b/pkg/cli/prompt_schema_build.go @@ -5,7 +5,6 @@ import ( "fmt" "strings" - "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" clickyrpc "github.com/flanksource/clicky/rpc" ) @@ -13,8 +12,6 @@ import ( // buildPromptSchemaDocument is the pure assembler (no I/O), taking a probed // adapter set so tests can drive it with a deterministic, network-free stub. func buildPromptSchemaDocument(adapters []AdapterStatus) (map[string]any, error) { - adapters = withCatalogModelFallback(adapters) - reflected, err := reflectedSchemas() if err != nil { return nil, err @@ -100,6 +97,9 @@ func buildBackendsCatalog(adapters []AdapterStatus, argsByBackend map[api.Backen if len(a.Models) > 0 { entry["models"] = a.Models } + if len(a.ModelDetails) > 0 { + entry["modelDetails"] = a.ModelDetails + } if a.ModelError != "" { entry["modelError"] = a.ModelError } @@ -151,64 +151,176 @@ func injectSpecConditionals(specMap map[string]any, adapters []AdapterStatus, ar if len(a.Models) > 0 { thenProps["model"] = map[string]any{"enum": toAnySlice(a.Models)} } + thenSchema := map[string]any{"properties": thenProps} + var effortRules []any + for _, model := range a.ModelDetails { + efforts := []any{""} + for _, effort := range model.SupportedEfforts { + efforts = append(efforts, string(effort)) + } + effortRules = append(effortRules, map[string]any{ + "if": map[string]any{ + "required": []any{"model"}, + "properties": map[string]any{"model": map[string]any{"const": model.ID}}, + }, + "then": map[string]any{ + "properties": map[string]any{"effort": map[string]any{"enum": efforts}}, + }, + }) + } + if len(effortRules) > 0 { + thenSchema["allOf"] = effortRules + } allOf = append(allOf, map[string]any{ "if": map[string]any{ "required": []any{"backend"}, "properties": map[string]any{"backend": map[string]any{"const": a.Backend}}, }, - "then": map[string]any{"properties": thenProps}, + "then": thenSchema, }) } specMap["allOf"] = allOf return nil } -// withCatalogModelFallback fills a backend's model list from the static catalog -// when it is unauthenticated and the live probe returned nothing, so the editor -// still offers model choices instead of an empty enum. The catalog answered, so -// the "set your key" hint is cleared. CLI/agent backends already list catalog -// models regardless of auth, so this only affects unauthenticated API backends -// (and no-ops where the catalog has no entry, e.g. the cmux backends). -func withCatalogModelFallback(adapters []AdapterStatus) []AdapterStatus { - out := make([]AdapterStatus, len(adapters)) - for i, a := range adapters { - if !a.Authenticated && len(a.Models) == 0 { - if ids := catalogModelIDs(api.Backend(a.Backend)); len(ids) > 0 { - a.Models = ids - a.ModelError = "" +// flatModels is a convenience union of every available model across adapters, +// shaped like clicky-ui's ChatModel catalog while retaining the legacy backend +// and ready fields for older consumers. +func flatModels(adapters []AdapterStatus) []map[string]any { + type entry struct { + data map[string]any + backends []string + } + out := []entry{} + positions := map[string]int{} + for _, a := range adapters { + provider := modelProviderForBackend(api.Backend(a.Backend)) + for _, model := range flatModelDetails(a) { + id := strings.TrimSpace(model.id) + if id == "" { + continue + } + key := provider + "\x00" + id + if idx, ok := positions[key]; ok { + if !containsString(out[idx].backends, a.Backend) { + out[idx].backends = append(out[idx].backends, a.Backend) + out[idx].data["backends"] = out[idx].backends + } + if a.Ready() { + out[idx].data["configured"] = true + out[idx].data["ready"] = true + } + continue + } + label := strings.TrimSpace(model.label) + if label == "" { + label = id + } + backends := []string{a.Backend} + positions[key] = len(out) + out = append(out, entry{ + backends: backends, + data: map[string]any{ + "id": id, + "label": label, + "provider": provider, + "reasoning": modelSupportsReasoning(id), + "configured": a.Ready(), + "backends": backends, + "backend": a.Backend, + "ready": a.Ready(), + }, + }) + if len(model.supportedEfforts) > 0 { + values := make([]string, 0, len(model.supportedEfforts)) + for _, effort := range model.supportedEfforts { + values = append(values, string(effort)) + } + out[len(out)-1].data["supportedEfforts"] = values + } + if model.defaultEffort != api.EffortNone { + out[len(out)-1].data["defaultEffort"] = string(model.defaultEffort) } } - out[i] = a } - return out + flat := make([]map[string]any, 0, len(out)) + for _, item := range out { + flat = append(flat, item.data) + } + return flat } -// catalogModelIDs returns the catalog's model slugs for a backend, bare (no -// provider prefix) to match how spec.model names models. -func catalogModelIDs(b api.Backend) []string { - var out []string - for _, m := range ai.Catalog() { - if m.Backend == b { - out = append(out, m.BareID()) +type flatModelDetail struct { + id string + label string + supportedEfforts []api.Effort + defaultEffort api.Effort +} + +func flatModelDetails(adapter AdapterStatus) []flatModelDetail { + if len(adapter.ModelDetails) > 0 { + out := make([]flatModelDetail, 0, len(adapter.ModelDetails)) + for _, model := range adapter.ModelDetails { + out = append(out, flatModelDetail{ + id: model.ID, + label: model.Name, + supportedEfforts: model.SupportedEfforts, + defaultEffort: model.DefaultEffort, + }) } + return out + } + out := make([]flatModelDetail, 0, len(adapter.Models)) + for _, id := range adapter.Models { + out = append(out, flatModelDetail{id: id, label: id}) } return out } -// flatModels is a convenience union of every available model across adapters. -func flatModels(adapters []AdapterStatus) []map[string]any { - out := []map[string]any{} - for _, a := range adapters { - for _, id := range a.Models { - out = append(out, map[string]any{ - "id": id, - "backend": a.Backend, - "ready": a.Ready(), - }) +func modelProviderForBackend(backend api.Backend) string { + switch backend { + case api.BackendClaudeAgent, api.BackendClaudeCLI, api.BackendClaudeCmux: + return "claude-agent" + case api.BackendCodexAgent, api.BackendCodexCLI, api.BackendCodexCmux: + return "codex-cli" + case api.BackendGemini, api.BackendGeminiCLI: + return "googleai" + default: + return string(backend) + } +} + +func modelSupportsReasoning(id string) bool { + token := strings.ToLower(strings.TrimSpace(id)) + if slash := strings.LastIndex(token, "/"); slash >= 0 { + token = token[slash+1:] + } + token = strings.TrimPrefix(token, "models/") + switch { + case strings.HasPrefix(token, "claude-"), + strings.HasPrefix(token, "opus"), + strings.HasPrefix(token, "sonnet"), + strings.HasPrefix(token, "haiku"), + strings.HasPrefix(token, "gpt-"), + strings.HasPrefix(token, "o1"), + strings.HasPrefix(token, "o3"), + strings.HasPrefix(token, "o4"), + strings.HasPrefix(token, "gemini-"), + strings.HasPrefix(token, "deepseek"): + return true + default: + return false + } +} + +func containsString(values []string, needle string) bool { + for _, value := range values { + if value == needle { + return true } } - return out + return false } // argDefName is the $defs key for a backend's args schema, e.g. diff --git a/pkg/cli/prompt_schema_http_test.go b/pkg/cli/prompt_schema_http_test.go index cc5dddc..ed30c15 100644 --- a/pkg/cli/prompt_schema_http_test.go +++ b/pkg/cli/prompt_schema_http_test.go @@ -5,19 +5,13 @@ import ( "net/http" "net/http/httptest" "testing" - "time" ) func TestHandlePromptSchemaServesDocument(t *testing.T) { - prevProbe := probeSchemaAdapters - prevCache, prevAt := schemaAdapterCache, schemaAdapterAt - t.Cleanup(func() { - probeSchemaAdapters = prevProbe - schemaAdapterCache, schemaAdapterAt = prevCache, prevAt - }) + prev := schemaAdapters + t.Cleanup(func() { schemaAdapters = prev }) stub := stubbedSchemaAdapters(t) - probeSchemaAdapters = func() ([]AdapterStatus, error) { return stub, nil } - schemaAdapterCache, schemaAdapterAt = nil, time.Time{} + schemaAdapters = func() ([]AdapterStatus, error) { return stub, nil } req := httptest.NewRequest(http.MethodGet, "/api/captain/ai/prompt/schema", nil) rec := httptest.NewRecorder() diff --git a/pkg/cli/prompt_schema_test.go b/pkg/cli/prompt_schema_test.go index cec64cb..1b33df1 100644 --- a/pkg/cli/prompt_schema_test.go +++ b/pkg/cli/prompt_schema_test.go @@ -4,30 +4,31 @@ import ( "bytes" "encoding/json" "io" + "reflect" "strings" "testing" - "time" + "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" "github.com/spf13/cobra" ) -// fakeSchemaProbe reports no API keys and no CLI login files but all CLI binaries -// present. This keeps the probe hermetic (fetchAPIModels makes no network call -// without an API key) while CLI/agent backends still list their static catalog -// models, giving a deterministic adapter set. -func fakeSchemaProbe() authProbe { - return authProbe{ - getenv: func(string) string { return "" }, - lookPath: func(bin string) (string, error) { return "/usr/local/bin/" + bin, nil }, - fileExists: func(string) bool { return false }, - home: "/home/test", +// fakeSchemaProbe reports no API keys and no CLI login files but all CLI +// binaries present. This keeps the probe hermetic: fetchAPIModels makes no +// network call without an API key, API backends stay key-gated, and CLI-style +// backends project exact IDs from Captain's internal registry. +func fakeSchemaProbe() ai.AuthProbe { + return ai.AuthProbe{ + Getenv: func(string) string { return "" }, + LookPath: func(bin string) (string, error) { return "/usr/local/bin/" + bin, nil }, + FileExists: func(string) bool { return false }, + Home: "/home/test", } } func stubbedSchemaAdapters(t *testing.T) []AdapterStatus { t.Helper() - adapters, err := ProbeAdapters(WhoamiOptions{Models: true}, fakeSchemaProbe()) + adapters, err := ai.ProbeAdapters(ai.WhoamiOptions{Models: true}, fakeSchemaProbe()) if err != nil { t.Fatalf("ProbeAdapters: %v", err) } @@ -63,22 +64,41 @@ func TestPromptSchemaDocumentBackendsAndConditionals(t *testing.T) { } } - // codex-cli's models come from the static catalog (one entry, AgentModel - // "gpt-5-codex") - an independent baseline unaffected by API keys. - codexModels := byName[string(api.BackendCodexCLI)]["models"].([]string) - if !contains(codexModels, "gpt-5-codex") { - t.Errorf("codex-cli models = %v, want to contain gpt-5-codex", codexModels) + codexCLIModels, hasModels := byName[string(api.BackendCodexCLI)]["models"].([]string) + if !hasModels || len(codexCLIModels) == 0 { + t.Fatalf("codex-cli should expose exact registry models without API provider data: %+v", byName[string(api.BackendCodexCLI)]) } + if got := codexCLIModels[0]; got != "gpt-5.6-sol" { + t.Errorf("codex-cli first model = %q, want gpt-5.6-sol", got) + } + flat := doc["models"].([]map[string]any) + codexModel := schemaModelForBackend(t, flat, "gpt-5.6-sol", string(api.BackendCodexCLI)) + if got := codexModel["provider"]; got != "codex-cli" { + t.Errorf("flat model provider = %v, want codex-cli", got) + } + if got, ok := codexModel["label"].(string); !ok || got == "" { + t.Errorf("flat model label = %#v, want non-empty string", codexModel["label"]) + } + if _, ok := codexModel["reasoning"].(bool); !ok { + t.Errorf("flat model reasoning = %#v, want bool", codexModel["reasoning"]) + } + if efforts, ok := codexModel["supportedEfforts"].([]string); !ok || !containsString(efforts, "max") || !containsString(efforts, "ultra") { + t.Errorf("flat model supportedEfforts = %#v, want patched Codex ultra", codexModel["supportedEfforts"]) + } + if _, ok := codexModel["defaultEffort"]; ok { + t.Errorf("flat model should not have a locally patched default effort: %#v", codexModel["defaultEffort"]) + } + if got, ok := codexModel["configured"].(bool); !ok || got { + t.Errorf("flat model configured = %#v, want false for fake unauthenticated CLI", codexModel["configured"]) + } + assertSchemaModelBackends(t, codexModel, string(api.BackendCodexCLI), string(api.BackendCodexAgent), string(api.BackendCodexCmux)) - // anthropic is unauthenticated in the fake probe, so its models fall back to - // the catalog and the "set your key" hint is cleared. anthropic := byName[string(api.BackendAnthropic)] - anthropicModels, _ := anthropic["models"].([]string) - if !contains(anthropicModels, "claude-sonnet-5") { - t.Errorf("unauthenticated anthropic models = %v, want catalog fallback incl claude-sonnet-5", anthropicModels) + if _, hasModels := anthropic["models"]; hasModels { + t.Errorf("anthropic should not synthesize models without live provider data: %+v", anthropic) } - if _, hasErr := anthropic["modelError"]; hasErr { - t.Errorf("anthropic should carry no modelError once the catalog answers") + if errText, _ := anthropic["modelError"].(string); !strings.Contains(errText, "ANTHROPIC_API_KEY") { + t.Errorf("anthropic modelError = %q, want missing-key hint", errText) } spec := doc["spec"].(map[string]any) @@ -134,6 +154,64 @@ func TestPromptSchemaDocumentBackendsAndConditionals(t *testing.T) { } } +func TestInjectSpecConditionalsUsesOnlyModelBackedEfforts(t *testing.T) { + spec := map[string]any{} + adapters := []AdapterStatus{{ + Backend: string(api.BackendCodexAgent), + ModelDetails: []ai.ModelDef{ + {ID: "gpt-5.6-sol", SupportedEfforts: []api.Effort{api.EffortLow, api.EffortMax}}, + {ID: "no-effort-model"}, + }, + }} + if err := injectSpecConditionals(spec, adapters, nil); err != nil { + t.Fatalf("injectSpecConditionals: %v", err) + } + rules := spec["allOf"].([]any)[0].(map[string]any)["then"].(map[string]any)["allOf"].([]any) + got := map[string][]any{} + for _, raw := range rules { + rule := raw.(map[string]any) + model := rule["if"].(map[string]any)["properties"].(map[string]any)["model"].(map[string]any)["const"].(string) + got[model] = rule["then"].(map[string]any)["properties"].(map[string]any)["effort"].(map[string]any)["enum"].([]any) + } + if want := []any{"", "low", "max"}; !reflect.DeepEqual(got["gpt-5.6-sol"], want) { + t.Errorf("Sol effort enum = %v, want %v", got["gpt-5.6-sol"], want) + } + if want := []any{""}; !reflect.DeepEqual(got["no-effort-model"], want) { + t.Errorf("no-effort model enum = %v, want %v", got["no-effort-model"], want) + } +} + +func schemaModelForBackend(t *testing.T, models []map[string]any, id, backend string) map[string]any { + t.Helper() + for _, model := range models { + if model["id"] != id { + continue + } + backends, ok := model["backends"].([]string) + if !ok { + t.Fatalf("model %s backends = %T, want []string", id, model["backends"]) + } + if containsString(backends, backend) { + return model + } + } + t.Fatalf("models[] missing id %q with backend %q: %+v", id, backend, models) + return nil +} + +func assertSchemaModelBackends(t *testing.T, model map[string]any, want ...string) { + t.Helper() + backends, ok := model["backends"].([]string) + if !ok { + t.Fatalf("model backends = %T, want []string", model["backends"]) + } + for _, backend := range want { + if !containsString(backends, backend) { + t.Errorf("model backends = %v, missing %s", backends, backend) + } + } +} + func TestPromptSchemaExampleIsPortable(t *testing.T) { ex := promptSchemaExampleSpec() @@ -152,6 +230,19 @@ func TestPromptSchemaExampleIsPortable(t *testing.T) { if _, err := api.InferBackend(model); err != nil { t.Errorf("example model %q does not resolve to a backend: %v", model, err) } + prompt := ex["prompt"].(map[string]any) + for _, excluded := range []string{"source", "metadata"} { + if _, ok := prompt[excluded]; ok { + t.Errorf("example prompt includes editor-excluded %q", excluded) + } + } + setup := ex["setup"].(map[string]any) + if _, ok := setup["env"]; ok { + t.Errorf("example setup uses stale env key: %+v", setup) + } + if _, ok := setup["envVars"]; !ok { + t.Errorf("example setup missing envVars: %+v", setup) + } // The example selects claude-cmux, so its cliArgs must validate as // ClaudeCmuxOptions — matching the schema's per-backend conditional. @@ -167,50 +258,11 @@ func TestPromptSchemaExampleIsPortable(t *testing.T) { } } -func TestCachedSchemaAdaptersReusesWithinTTL(t *testing.T) { - prevProbe := probeSchemaAdapters - prevCache, prevAt := schemaAdapterCache, schemaAdapterAt - t.Cleanup(func() { - probeSchemaAdapters = prevProbe - schemaAdapterCache, schemaAdapterAt = prevCache, prevAt - }) - - stub := stubbedSchemaAdapters(t) - calls := 0 - probeSchemaAdapters = func() ([]AdapterStatus, error) { - calls++ - return stub, nil - } - schemaAdapterCache, schemaAdapterAt = nil, time.Time{} - - base := time.Unix(1_000_000, 0) - if _, err := cachedSchemaAdapters(base); err != nil { - t.Fatalf("cachedSchemaAdapters: %v", err) - } - if _, err := cachedSchemaAdapters(base.Add(10 * time.Second)); err != nil { - t.Fatalf("cachedSchemaAdapters: %v", err) - } - if calls != 1 { - t.Errorf("probe called %d times within TTL, want 1", calls) - } - if _, err := cachedSchemaAdapters(base.Add(2 * schemaAdapterCacheTTL)); err != nil { - t.Fatalf("cachedSchemaAdapters: %v", err) - } - if calls != 2 { - t.Errorf("probe called %d times after TTL expiry, want 2", calls) - } -} - func TestWritePromptSchemaEmitsValidJSON(t *testing.T) { - prevProbe := probeSchemaAdapters - prevCache, prevAt := schemaAdapterCache, schemaAdapterAt - t.Cleanup(func() { - probeSchemaAdapters = prevProbe - schemaAdapterCache, schemaAdapterAt = prevCache, prevAt - }) + prev := schemaAdapters + t.Cleanup(func() { schemaAdapters = prev }) stub := stubbedSchemaAdapters(t) - probeSchemaAdapters = func() ([]AdapterStatus, error) { return stub, nil } - schemaAdapterCache, schemaAdapterAt = nil, time.Time{} + schemaAdapters = func() ([]AdapterStatus, error) { return stub, nil } var buf bytes.Buffer if err := WritePromptSchema(&buf); err != nil { diff --git a/pkg/cli/prompt_source.go b/pkg/cli/prompt_source.go index ffd4dbc..5fdcae4 100644 --- a/pkg/cli/prompt_source.go +++ b/pkg/cli/prompt_source.go @@ -32,20 +32,21 @@ func readStdinIfCLI(ctx context.Context) string { } // loadPromptContent resolves the prompt source for the unified prompt commands. -// Precedence: the positional (a .prompt filepath or a registry id) > --prompt/-p -// text > piped stdin. usedStdin reports whether stdin became the prompt body (so -// the caller does not also expose it as the {{input}} variable). +// Precedence: the positional (a discovered name, .prompt filepath, or registry +// id) > --prompt/-p text > piped stdin. usedStdin reports whether stdin became +// the prompt body (so the caller does not also expose it as the {{input}} +// variable). func loadPromptContent(ctx context.Context, id string, opts AIPromptOptions, stdin string) (content, source string, usedStdin bool, record promptRecord, err error) { switch { case strings.TrimSpace(id) != "": - record, err := resolvePromptRecord(ctx, id) // .prompt filepath or registry id + record, err := resolvePromptRecord(ctx, id) if err != nil { return "", "", false, promptRecord{}, err } if record.Source.Kind == "file" { log.Debugf("prompt source: file %s (positional %q)", record.Path, id) } else { - log.Debugf("prompt source: registry id %q → %s/%s", id, record.Source.Kind, record.Rel) + log.Debugf("prompt source: resolved %q → %s/%s", id, record.Source.Kind, record.Rel) } c, err := readPromptContent(record) if err != nil { @@ -58,8 +59,10 @@ func loadPromptContent(ctx context.Context, id string, opts AIPromptOptions, std case strings.TrimSpace(stdin) != "": log.Debugf("prompt source: stdin (%d chars)", len(stdin)) return stdin, "", true, promptRecord{Rel: "stdin.prompt"}, nil + case len(opts.Attach) > 0: + return ephemeralPromptContent(), "", false, promptRecord{Rel: "attachment.prompt"}, nil default: - return "", "", false, promptRecord{}, fmt.Errorf("prompt required: pass a .prompt file/id, --prompt/-p text, or pipe via stdin") + return "", "", false, promptRecord{}, fmt.Errorf("prompt or attachment required: pass a prompt name, .prompt file, id, --prompt/-p text, --attach/-A, or pipe via stdin") } } @@ -107,20 +110,6 @@ func renderLoadedContent(content, source string, vars map[string]any, opts AIPro return req, cfg, nil } -// renderPromptSource is the single render pipeline shared by run and the -// deprecated ai-prompt alias: load content → render → overlay → normalize. -func renderPromptSource(ctx context.Context, id string, opts AIPromptOptions, varsJSON, stdin string) (ai.Request, ai.Config, error) { - content, source, usedStdin, _, err := loadPromptContent(ctx, id, opts, stdin) - if err != nil { - return ai.Request{}, ai.Config{}, err - } - vars, err := promptVars(opts, varsJSON, stdin, usedStdin) - if err != nil { - return ai.Request{}, ai.Config{}, err - } - return renderLoadedContent(content, source, vars, opts) -} - // actionFlagsToOptions reconstructs the typed AIPromptOptions from the entity // action's stringly-typed flag map (clicky CSV-encodes []string and "true"/"false" // for bool), so the render/run core can reuse overlayCLI. @@ -161,6 +150,10 @@ func actionFlagsToOptions(f map[string]string) (AIPromptOptions, error) { o.System = f["system"] o.AppendSystem = f["append-system"] o.Var = flagSlice(f["var"]) + if attach := strings.TrimSpace(f["attach"]); attach != "" { + o.Attach = []string{attach} + } + o.MultiModels = flagSlice(f["multi-models"]) o.Timeout = f["timeout"] o.NoStream = flagBool(f["no-stream"]) return o, nil diff --git a/pkg/cli/prompt_source_test.go b/pkg/cli/prompt_source_test.go index 2e6ff09..461927d 100644 --- a/pkg/cli/prompt_source_test.go +++ b/pkg/cli/prompt_source_test.go @@ -17,6 +17,7 @@ func TestActionFlagsToOptions_DecodesSlicesBoolsInts(t *testing.T) { "skill-dir": "/a,/b", "max-tokens": "1234", "var": "a=1,b=2", + "multi-models": "cli:sonnet-5,cmux:opus", "prompt": "hi", "permission-mode": "plan", } @@ -42,6 +43,9 @@ func TestActionFlagsToOptions_DecodesSlicesBoolsInts(t *testing.T) { if len(o.Var) != 2 || o.Var[0] != "a=1" { t.Errorf("var = %v", o.Var) } + if len(o.MultiModels) != 2 || o.MultiModels[0] != "cli:sonnet-5" || o.MultiModels[1] != "cmux:opus" { + t.Errorf("multi-models = %v", o.MultiModels) + } if _, err := actionFlagsToOptions(map[string]string{"max-tokens": "nope"}); err == nil { t.Error("expected error on non-numeric max-tokens") } diff --git a/pkg/cli/prompt_sources.go b/pkg/cli/prompt_sources.go new file mode 100644 index 0000000..b81e88a --- /dev/null +++ b/pkg/cli/prompt_sources.go @@ -0,0 +1,254 @@ +package cli + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "fmt" + "io/fs" + "os" + "path/filepath" + "strings" + + promptlib "github.com/flanksource/captain/pkg/ai/prompt" + "github.com/flanksource/captain/pkg/captainconfig" + "github.com/flanksource/clicky" + clickyapi "github.com/flanksource/clicky/api" +) + +func buildPromptSources(ctx context.Context) ([]promptSource, error) { + sources := []promptSource{{ + Kind: "embedded", + ID: "embedded", + Label: "Embedded examples", + WalkRoot: "testdata", + FS: promptlib.Examples, + Writable: false, + }} + + seen := map[string]bool{} + addLocal := func(raw, base string) error { + dir, err := resolvePromptDir(raw, base) + if err != nil { + return err + } + if seen[dir] { + return nil + } + seen[dir] = true + sources = append(sources, promptSource{ + Kind: "local", + ID: hashPromptDir(dir), + Label: dir, + Root: dir, + Writable: true, + }) + return nil + } + + cfg, exists, err := captainconfig.Load() + if err != nil { + return nil, err + } + if exists { + configPath, err := captainconfig.Path() + if err != nil { + return nil, err + } + base := filepath.Dir(configPath) + for _, dir := range cfg.Prompts.Dirs { + if err := addLocal(dir, base); err != nil { + return nil, err + } + } + } + cwd, err := os.Getwd() + if err != nil { + return nil, err + } + for _, dir := range promptDirsFromContext(ctx) { + if err := addLocal(dir, cwd); err != nil { + return nil, err + } + } + if _, ok := firstWritableSource(sources); !ok { + dir := filepath.Join(cwd, ".captain", "prompts") + sources = append(sources, promptSource{ + Kind: "local", + ID: hashPromptDir(dir), + Label: dir, + Root: dir, + Writable: true, + Implicit: true, + }) + } + return sources, nil +} + +func promptDirsFromContext(ctx context.Context) []string { + if ctx == nil { + return nil + } + if dirs, ok := ctx.Value(promptDirsContextKey{}).([]string); ok { + return dirs + } + return nil +} + +func resolvePromptDir(raw, base string) (string, error) { + dir := strings.TrimSpace(raw) + if dir == "" { + return "", fmt.Errorf("prompt dir cannot be empty") + } + if strings.HasPrefix(dir, "~") { + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + switch { + case dir == "~": + dir = home + case strings.HasPrefix(dir, "~/"): + dir = filepath.Join(home, dir[2:]) + default: + return "", fmt.Errorf("unsupported home-relative prompt dir %q", raw) + } + } + if !filepath.IsAbs(dir) { + dir = filepath.Join(base, dir) + } + abs, err := filepath.Abs(dir) + if err != nil { + return "", err + } + info, err := os.Stat(abs) + if err != nil { + return "", fmt.Errorf("prompt dir %s: %w", abs, err) + } + if !info.IsDir() { + return "", fmt.Errorf("prompt dir %s is not a directory", abs) + } + if resolved, err := filepath.EvalSymlinks(abs); err == nil { + abs = resolved + } + return filepath.Clean(abs), nil +} + +func writableSourceByID(sources []promptSource, id string) (promptSource, bool) { + for _, source := range sources { + if source.ID == id && source.Writable { + return source, true + } + } + return promptSource{}, false +} + +func firstWritableSource(sources []promptSource) (promptSource, bool) { + for _, source := range sources { + if source.Writable { + return source, true + } + } + return promptSource{}, false +} + +func safeLocalPromptPath(source promptSource, rel string) (string, error) { + cleanRel := strings.TrimPrefix(filepath.Clean(filepath.FromSlash(rel)), string(filepath.Separator)) + if cleanRel == "." || cleanRel == "" || filepath.IsAbs(rel) || strings.HasPrefix(cleanRel, ".."+string(filepath.Separator)) || cleanRel == ".." { + return "", fmt.Errorf("invalid prompt path %q", rel) + } + if filepath.Ext(cleanRel) != ".prompt" { + return "", fmt.Errorf("prompt path must end with .prompt") + } + full := filepath.Join(source.Root, cleanRel) + abs, err := filepath.Abs(full) + if err != nil { + return "", err + } + relToRoot, err := filepath.Rel(source.Root, abs) + if err != nil { + return "", err + } + if strings.HasPrefix(relToRoot, ".."+string(filepath.Separator)) || relToRoot == ".." || filepath.IsAbs(relToRoot) { + return "", fmt.Errorf("prompt path escapes source root") + } + if info, err := os.Lstat(abs); err == nil && info.Mode()&fs.ModeSymlink != 0 { + return "", fmt.Errorf("prompt symlinks are not supported") + } + return abs, nil +} + +func normalizeWriteRelPath(relPath, name string) (string, error) { + rel := strings.TrimSpace(relPath) + if rel == "" { + rel = slugPromptName(name) + } + if rel == "" { + return "", fmt.Errorf("prompt name or path required") + } + if !strings.HasSuffix(rel, ".prompt") { + rel += ".prompt" + } + rel = filepath.ToSlash(filepath.Clean(filepath.FromSlash(rel))) + if strings.HasPrefix(rel, "../") || rel == ".." || strings.HasPrefix(rel, "/") { + return "", fmt.Errorf("invalid prompt path %q", relPath) + } + return rel, nil +} + +func slugPromptName(name string) string { + name = strings.ToLower(strings.TrimSpace(name)) + var b strings.Builder + lastDash := false + for _, r := range name { + ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') + if ok { + b.WriteRune(r) + lastDash = false + continue + } + if !lastDash { + b.WriteRune('-') + lastDash = true + } + } + return strings.Trim(b.String(), "-") +} + +func encodePromptID(kind, sourceID, rel string) string { + return base64.RawURLEncoding.EncodeToString([]byte(kind + "\x00" + sourceID + "\x00" + filepath.ToSlash(rel))) +} + +func decodePromptID(id string) (promptRef, error) { + data, err := base64.RawURLEncoding.DecodeString(id) + if err != nil { + return promptRef{}, fmt.Errorf("invalid prompt id: %w", err) + } + parts := strings.SplitN(string(data), "\x00", 3) + if len(parts) != 3 || parts[0] == "" || parts[1] == "" || parts[2] == "" { + return promptRef{}, fmt.Errorf("invalid prompt id") + } + return promptRef{Kind: parts[0], SourceID: parts[1], RelPath: filepath.ToSlash(parts[2])}, nil +} + +func hashPromptDir(dir string) string { + sum := sha256.Sum256([]byte(dir)) + return hex.EncodeToString(sum[:])[:12] +} + +func ValidatePromptDirs(dirs []string) error { + cwd, err := os.Getwd() + if err != nil { + return err + } + for _, dir := range dirs { + if _, err := resolvePromptDir(dir, cwd); err != nil { + return err + } + } + return nil +} + +var _ clicky.EntityItem = PromptSummary{} +var _ clickyapi.TableProvider = PromptSummary{} diff --git a/pkg/cli/prompt_spec.go b/pkg/cli/prompt_spec.go new file mode 100644 index 0000000..0d94414 --- /dev/null +++ b/pkg/cli/prompt_spec.go @@ -0,0 +1,241 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "io" + "sort" + "strconv" + "strings" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/collections" + clickyrpc "github.com/flanksource/clicky/rpc" +) + +func applyPromptDefaults(req *ai.Request, cfg *ai.Config) error { + savedCfg := loadSavedConfig() + saved := savedCfg.AI + promptModel := req.Model + if promptModel.Name == "" { + promptModel.Name = cfg.Model.Name + } + if promptModel.ID == "" { + promptModel.ID = cfg.Model.ID + } + if promptModel.Backend == "" { + promptModel.Backend = cfg.Model.Backend + } + identity := selectModelIdentity( + api.Model{Name: promptModel.Name, ID: promptModel.ID, Backend: promptModel.Backend}, + ) + req.Name, req.ID, req.Backend = identity.Name, identity.ID, identity.Backend + if cfg.Model.Effort != api.EffortNone { + // An effort-qualified model selector (for example agent:sol:high) + // is model-local and intentionally overrides the request-wide flag/default. + req.Effort = cfg.Model.Effort + } else if req.Effort == "" { + req.Effort = cfg.Model.Effort + } + resolved, err := applyProviderDefaults(req.Model, saved) + if err != nil { + return err + } + req.Model = resolved + req.NoCache = req.NoCache || saved.NoCache + if req.Budget.MaxTokens == 0 { + req.Budget.MaxTokens = firstPositive(cfg.Budget.MaxTokens, saved.MaxTokens, 4096) + } + if req.Budget.Cost == 0 { + req.Budget.Cost = firstPositiveFloat(cfg.Budget.Cost, saved.BudgetUSD) + } + + cfg.Model = req.Model + cfg.Budget = req.Budget + cfg.NoCache = req.NoCache + if isZeroSchemaRepair(cfg.SchemaRepair) { + cfg.SchemaRepair = schemaRepairConfig(savedCfg.Prompts.SchemaRepair) + } + return nil +} + +func mergeStringMaps(base, overlay map[string]string) map[string]string { + if len(overlay) == 0 { + return base + } + out := make(map[string]string, collections.SafeAdd(len(base), len(overlay))) + for k, v := range base { + out[k] = v + } + for k, v := range overlay { + out[k] = v + } + return out +} + +func mergeToolModes(base, overlay map[string]api.ToolMode) map[string]api.ToolMode { + if len(overlay) == 0 { + return base + } + out := make(map[string]api.ToolMode, collections.SafeAdd(len(base), len(overlay))) + for k, v := range base { + out[k] = v + } + for k, v := range overlay { + out[k] = v + } + return out +} + +func mergePresets(base, overlay []api.Preset) []api.Preset { + if len(overlay) == 0 { + return base + } + seen := make(map[api.Preset]bool, collections.SafeAdd(len(base), len(overlay))) + out := make([]api.Preset, 0, collections.SafeAdd(len(base), len(overlay))) + for _, preset := range base { + if seen[preset] { + continue + } + seen[preset] = true + out = append(out, preset) + } + for _, preset := range overlay { + if seen[preset] { + continue + } + seen[preset] = true + out = append(out, preset) + } + return out +} + +func sortedStringKeys[V any](m map[string]V) []string { + keys := make([]string, 0, len(m)) + for key := range m { + if key != "" { + keys = append(keys, key) + } + } + sort.Strings(keys) + return keys +} + +func enabledResourcePolicies(in api.ResourcePolicies) api.ResourcePolicies { + out := api.ResourcePolicies{} + for _, key := range sortedStringKeys(in) { + if in[key] == api.ResourceEnabled { + out[key] = api.ResourceEnabled + } + } + return out +} + +func dedupeStrings(in []string) []string { + seen := map[string]bool{} + var out []string + for _, item := range in { + if item == "" || seen[item] { + continue + } + seen[item] = true + out = append(out, item) + } + return out +} + +func readRenderRequest(ctx context.Context, flags map[string]string) (PromptRenderRequest, error) { + var req PromptRenderRequest + if err := decodePromptBody(ctx, map[string]any{}, &req); err != nil { + return PromptRenderRequest{}, err + } + if req.Variables == nil { + req.Variables = map[string]any{} + } + if err := mergePromptActionFlags(&req, flags); err != nil { + return PromptRenderRequest{}, err + } + return req, nil +} + +func mergePromptActionFlags(req *PromptRenderRequest, flags map[string]string) error { + if len(flags) == 0 { + return nil + } + if raw := strings.TrimSpace(flags["vars"]); raw != "" { + var vars map[string]any + if err := json.Unmarshal([]byte(raw), &vars); err != nil { + return fmt.Errorf("parse --vars JSON: %w", err) + } + req.Variables = vars + } + if v := strings.TrimSpace(flags["model"]); v != "" { + ensureRenderSpec(req).Name = v + } + if v := strings.TrimSpace(flags["fallback"]); v != "" { + ensureRenderSpec(req).Fallbacks = fallbackModelsFromFlags([]string{v}) + } + if v := strings.TrimSpace(flags["backend"]); v != "" { + ensureRenderSpec(req).Backend = api.Backend(v) + } + if v := strings.TrimSpace(flags["timeout"]); v != "" { + ensureRenderSpec(req).Budget.Timeout = v + } + if v := strings.TrimSpace(flags["max-tokens"]); v != "" { + n, err := strconv.Atoi(v) + if err != nil { + return fmt.Errorf("invalid --max-tokens %q: %w", v, err) + } + ensureRenderSpec(req).Budget.MaxTokens = n + } + return nil +} + +func ensureRenderSpec(req *PromptRenderRequest) *api.Spec { + if req.Spec == nil { + req.Spec = &api.Spec{} + } + return req.Spec +} + +func decodePromptBody(ctx context.Context, flat map[string]any, dst any) error { + if r, ok := clickyrpc.RequestFromContext(ctx); ok && r.Body != nil { + body, err := io.ReadAll(r.Body) + if err != nil { + return fmt.Errorf("read request body: %w", err) + } + r.Body = io.NopCloser(strings.NewReader(string(body))) + if len(strings.TrimSpace(string(body))) > 0 { + decoder := json.NewDecoder(strings.NewReader(string(body))) + decoder.DisallowUnknownFields() + if err := decoder.Decode(dst); err != nil { + return fmt.Errorf("decode request body: %w", err) + } + return nil + } + } + if len(flat) == 0 { + return nil + } + data, err := json.Marshal(flat) + if err != nil { + return err + } + decoder := json.NewDecoder(strings.NewReader(string(data))) + decoder.DisallowUnknownFields() + if err := decoder.Decode(dst); err != nil { + return fmt.Errorf("decode command body: %w", err) + } + return nil +} + +func runtimeTimeout(raw string) time.Duration { + timeout, _ := time.ParseDuration(raw) + if timeout <= 0 { + return 120 * time.Second + } + return timeout +} diff --git a/pkg/cli/provider_defaults.go b/pkg/cli/provider_defaults.go new file mode 100644 index 0000000..3099234 --- /dev/null +++ b/pkg/cli/provider_defaults.go @@ -0,0 +1,126 @@ +package cli + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" +) + +type ProviderDefaultView struct { + Agent string `json:"agent"` + Model string `json:"model"` + Effort string `json:"effort"` + Configured bool `json:"configured"` +} + +func effectiveProviderDefaults(saved captainconfig.AIDefaults, provider api.Backend) (ProviderDefaultView, error) { + if provider.Provider() != provider { + return ProviderDefaultView{}, fmt.Errorf("invalid provider %q", provider) + } + configured, exists := saved.Providers[string(provider)] + legacy := saved.Provider(string(provider)) + if configured.Agent == "" { + configured.Agent = legacy.Agent + } + if configured.Model == "" { + configured.Model = legacy.Model + } + if configured.ReasoningEffort == "" { + configured.ReasoningEffort = legacy.ReasoningEffort + } + agent := api.Backend(strings.TrimSpace(configured.Agent)) + if agent == "" { + agent = provider + } + if agent.Provider() != provider { + return ProviderDefaultView{}, fmt.Errorf("agent %q does not belong to provider %q", agent, provider) + } + model := strings.TrimSpace(configured.Model) + if model == "" { + model = defaultModelFor(agent) + } + effort := api.Effort(strings.TrimSpace(configured.ReasoningEffort)) + if err := ai.ValidateModelEffort(agent, model, effort); err != nil { + return ProviderDefaultView{}, err + } + return ProviderDefaultView{ + Agent: string(agent), Model: model, Effort: string(effort), Configured: exists, + }, nil +} + +func applyProviderDefaults(model api.Model, saved captainconfig.AIDefaults) (api.Model, error) { + var err error + if strings.TrimSpace(model.Name) != "" { + model, err = model.Expand() + if err != nil { + return api.Model{}, err + } + } + model, err = applyCandidateDefaults(model, saved, true) + if err != nil { + return api.Model{}, err + } + for i := range model.Fallbacks { + fallback := model.Fallbacks[i] + fallback, err = fallback.Expand() + if err != nil { + return api.Model{}, fmt.Errorf("fallback[%d]: %w", i, err) + } + fallback, err = applyCandidateDefaults(fallback, saved, false) + if err != nil { + return api.Model{}, fmt.Errorf("fallback[%d]: %w", i, err) + } + model.Fallbacks[i] = fallback + } + return model, nil +} + +func applyCandidateDefaults(model api.Model, saved captainconfig.AIDefaults, allowActive bool) (api.Model, error) { + provider := model.Backend.Provider() + if provider == "" && strings.TrimSpace(model.Name) != "" { + backend, err := api.InferBackend(model.Name) + if err != nil { + return api.Model{}, err + } + provider = backend.Provider() + } + if provider == "" && allowActive { + provider = api.Backend(saved.ActiveProvider()) + } + if provider == "" { + return api.Model{}, fmt.Errorf("provider cannot be resolved for model %q", model.Name) + } + defaults, err := effectiveProviderDefaults(saved, provider) + if err != nil { + return api.Model{}, err + } + if model.Backend == "" { + model.Backend = api.Backend(defaults.Agent) + } + if strings.TrimSpace(model.Name) == "" { + model.Name = defaults.Model + } + if model.Effort == api.EffortNone { + model.Effort = api.Effort(defaults.Effort) + } + if err := ai.ValidateModelEffort(model.Backend, model.Name, model.Effort); err != nil { + return api.Model{}, err + } + return model, nil +} + +func allProviderDefaults(saved captainconfig.AIDefaults) (map[string]ProviderDefaultView, error) { + providers := []api.Backend{api.AnthropicProvider, api.OpenAIProvider, api.GeminiProvider, api.DeepSeekProvider} + out := make(map[string]ProviderDefaultView, len(providers)) + for _, provider := range providers { + defaults, err := effectiveProviderDefaults(saved, provider) + if err != nil { + return nil, err + } + out[string(provider)] = defaults + } + return out, nil +} diff --git a/pkg/cli/provider_defaults_test.go b/pkg/cli/provider_defaults_test.go new file mode 100644 index 0000000..f4d2fb4 --- /dev/null +++ b/pkg/cli/provider_defaults_test.go @@ -0,0 +1,51 @@ +package cli + +import ( + "testing" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" +) + +func TestApplyProviderDefaultsFillsMissingFieldsAndFallbacks(t *testing.T) { + saved := captainconfig.AIDefaults{ + DefaultProvider: "anthropic", + Providers: map[string]captainconfig.ProviderDefaults{ + "anthropic": {Agent: "claude-agent", Model: "claude-sonnet-5", ReasoningEffort: "high"}, + "openai": {Agent: "codex-agent", Model: "gpt-5.6-sol", ReasoningEffort: "medium"}, + }, + } + got, err := applyProviderDefaults(api.Model{Fallbacks: api.ModelList{{Name: "gpt-5.6"}}}, saved) + if err != nil { + t.Fatalf("applyProviderDefaults: %v", err) + } + if got.Backend != api.BackendClaudeAgent || got.Name != "claude-sonnet-5" || got.Effort != api.EffortHigh { + t.Fatalf("primary = %+v", got) + } + if fallback := got.Fallbacks[0]; fallback.Backend != api.BackendCodexAgent || fallback.Effort != api.EffortMedium { + t.Fatalf("fallback = %+v", fallback) + } +} + +func TestApplyProviderDefaultsPreservesExplicitFields(t *testing.T) { + saved := captainconfig.AIDefaults{Providers: map[string]captainconfig.ProviderDefaults{ + "openai": {Agent: "codex-agent", Model: "gpt-5.6-sol", ReasoningEffort: "medium"}, + }} + want := api.Model{Name: "gpt-explicit", Backend: api.BackendOpenAI, Effort: api.EffortHigh} + got, err := applyProviderDefaults(want, saved) + if err != nil { + t.Fatalf("applyProviderDefaults: %v", err) + } + if got.Name != want.Name || got.Backend != want.Backend || got.Effort != want.Effort { + t.Fatalf("explicit model changed: got %+v want %+v", got, want) + } +} + +func TestEffectiveProviderDefaultsRejectsCrossFamilyAgent(t *testing.T) { + saved := captainconfig.AIDefaults{Providers: map[string]captainconfig.ProviderDefaults{ + "anthropic": {Agent: "codex-agent", Model: "claude-sonnet-5"}, + }} + if _, err := effectiveProviderDefaults(saved, api.BackendAnthropic); err == nil { + t.Fatal("expected cross-family agent error") + } +} diff --git a/pkg/cli/secret_catalog.go b/pkg/cli/secret_catalog.go index a3cc00e..d3038d0 100644 --- a/pkg/cli/secret_catalog.go +++ b/pkg/cli/secret_catalog.go @@ -9,6 +9,7 @@ import ( "strings" "unicode/utf8" + "github.com/flanksource/captain/pkg/ai" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes" @@ -169,7 +170,7 @@ func previewsForConfigMap(item corev1.ConfigMap) []secretKeyPreview { out := make([]secretKeyPreview, 0, len(keys)) for _, key := range keys { if value, ok := item.Data[key]; ok { - out = append(out, secretKeyPreview{Key: key, Value: maskKey(value)}) + out = append(out, secretKeyPreview{Key: key, Value: ai.MaskKey(value)}) continue } out = append(out, secretKeyPreview{Key: key, Value: byteCountPreview(item.BinaryData[key])}) @@ -212,7 +213,7 @@ func maskPreviewBytes(value []byte) string { if !utf8.Valid(value) { return byteCountPreview(value) } - return maskKey(string(value)) + return ai.MaskKey(string(value)) } func byteCountPreview(value []byte) string { diff --git a/pkg/cli/secret_catalog_test.go b/pkg/cli/secret_catalog_test.go index 0ea50fc..260e2a6 100644 --- a/pkg/cli/secret_catalog_test.go +++ b/pkg/cli/secret_catalog_test.go @@ -3,6 +3,7 @@ package cli import ( "testing" + "github.com/flanksource/captain/pkg/ai" corev1 "k8s.io/api/core/v1" ) @@ -18,7 +19,7 @@ func TestPreviewsForByteDataMasksTextAndBinary(t *testing.T) { if previews[0].Key != "binary" || previews[0].Value != "3 bytes" { t.Fatalf("binary preview = %+v", previews[0]) } - if previews[1].Key != "token" || previews[1].Value != maskKey("sk-ant-api03-ABCDEFGH") { + if previews[1].Key != "token" || previews[1].Value != ai.MaskKey("sk-ant-api03-ABCDEFGH") { t.Fatalf("token preview = %+v", previews[1]) } } @@ -38,7 +39,7 @@ func TestConfigMapKeysIncludeBinaryData(t *testing.T) { if previews[0].Key != "cert" || previews[0].Value != "2 bytes" { t.Fatalf("cert preview = %+v", previews[0]) } - if previews[1].Key != "workspace" || previews[1].Value != maskKey("/repo/workspace") { + if previews[1].Key != "workspace" || previews[1].Value != ai.MaskKey("/repo/workspace") { t.Fatalf("workspace preview = %+v", previews[1]) } } diff --git a/pkg/cli/serve.go b/pkg/cli/serve.go index b981371..7c458d8 100644 --- a/pkg/cli/serve.go +++ b/pkg/cli/serve.go @@ -8,7 +8,6 @@ import ( "io" "io/fs" "net/http" - "net/url" "os" "os/exec" "os/signal" @@ -19,16 +18,20 @@ import ( "syscall" "time" - "github.com/flanksource/clicky/aichat" + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/monitor" "github.com/flanksource/clicky/rpc" rpchttp "github.com/flanksource/clicky/rpc/http" "github.com/flanksource/clicky/task" "github.com/spf13/cobra" ) -// all: keeps the committed dist/.gitkeep placeholder in the embed so the binary -// compiles without a built webapp (dist is gitignored and built on demand or in -// the release workflow). When index.html is absent, serve.go reports it at runtime. +// The built webapp under webapp/dist is generated locally and ignored. The +// tracked .gitkeep lets the embed pattern compile before `task www:build` +// generates the Vite assets. all: ensures the placeholder is embedded too. +// When index.html is absent, serve.go reports it at runtime. // //go:embed all:webapp/dist var captainWebappFS embed.FS @@ -41,6 +44,7 @@ type ServeOptions struct { Open bool ThreadsFile string PromptDirs []string + MCPServers []aichat.MCPServer } func NewServeCommand(version string) *cobra.Command { @@ -64,15 +68,17 @@ session. With --dev, Captain also starts the Vite dev server from pkg/cli/webapp and proxies /api back to this Go process.`, RunE: func(cmd *cobra.Command, args []string) error { - if err := opts.validate(); err != nil { + runOpts := opts + runOpts.Port = effectiveServePort(runOpts.Dev, cmd.Flags().Changed("port"), runOpts.Port) + if err := runOpts.validate(); err != nil { return err } - return RunServe(cmd.Context(), cmd.Root(), opts, version, cmd.OutOrStdout(), cmd.ErrOrStderr()) + return RunServe(cmd.Context(), cmd.Root(), runOpts, version, cmd.OutOrStdout(), cmd.ErrOrStderr()) }, } cmd.Flags().StringVar(&opts.Host, "host", opts.Host, "Host to bind the API server to") - cmd.Flags().IntVarP(&opts.Port, "port", "p", opts.Port, "Port to bind the API server to") + cmd.Flags().IntVarP(&opts.Port, "port", "p", opts.Port, "Port to bind the API server to (random when --dev is set)") cmd.Flags().BoolVar(&opts.Dev, "dev", false, "Launch the Vite dev server with /api proxied to Captain") cmd.Flags().IntVar(&opts.UIPort, "ui-port", opts.UIPort, "Port for the Vite dev server when --dev is set") cmd.Flags().BoolVar(&opts.Open, "open", false, "Open the web UI in the default browser") @@ -82,25 +88,6 @@ proxies /api back to this Go process.`, return cmd } -func (o ServeOptions) validate() error { - if strings.TrimSpace(o.Host) == "" { - return fmt.Errorf("host cannot be empty") - } - if o.Port < 1 || o.Port > 65535 { - return fmt.Errorf("invalid --port %d", o.Port) - } - if o.Dev && (o.UIPort < 1 || o.UIPort > 65535) { - return fmt.Errorf("invalid --ui-port %d", o.UIPort) - } - if strings.TrimSpace(o.ThreadsFile) == "" { - return fmt.Errorf("threads file cannot be empty") - } - if err := ValidatePromptDirs(o.PromptDirs); err != nil { - return err - } - return nil -} - func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, version string, stdout, stderr io.Writer) error { if ctx == nil { ctx = context.Background() @@ -117,6 +104,15 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve return err } threadStore := newFileThreadStore(opts.ThreadsFile) + attachmentStore, err := newAttachmentStore(cwd) + if err != nil { + return err + } + listener, addr, servePort, err := listenCaptainServer(opts.Host, opts.Port) + if err != nil { + return err + } + defer listener.Close() openAPIConfig := &rpc.OpenAPIConfig{ Title: "Captain", Description: "Captain command and agent launcher API.", @@ -124,7 +120,7 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve } serveConfig := &rpc.ServeConfig{ Host: opts.Host, - Port: opts.Port, + Port: servePort, Title: openAPIConfig.Title, Description: openAPIConfig.Description, Version: openAPIConfig.Version, @@ -136,18 +132,22 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve } rpcServer := rpc.NewSwaggerServer(serveConfig, rootCmd, openAPIConfig) - chat := aichat.NewServer(aichat.Options{ - RootCmd: rootCmd, - System: "You are Captain's coding-agent launcher assistant. Use Captain and Clicky tools when useful, " + - "prefer read-only inspection unless the user explicitly asks for edits, and keep follow-up guidance concise.", - Threads: threadStore, - ToolFilter: captainChatToolEnabled, - ToolApprovalPolicy: captainChatRequiresApproval, - Agent: aichat.AgentOptions{ - Cwd: cwd, - }, - }) - defer chat.Close() + openAPISpec, err := rpc.NewOpenAPIGenerator(openAPIConfig).GenerateFromCobraWithConfig(rootCmd, rpcServer.ConverterConfig()) + if err != nil { + return err + } + addCaptainPromptRunPaths(openAPISpec) + addCaptainProviderTokenPaths(openAPISpec) + addCaptainProviderDefaultsPaths(openAPISpec) + chat, mcpTools, err := newCaptainChatService(ctx, rootCmd, opts, cwd, threadStore, attachmentStore) + if err != nil { + return err + } + defer func() { + if err := mcpTools.Close(); err != nil { + log.Errorf("close Captain MCP tools: %v", err) + } + }() // Prompt runs are tracked as clicky task groups. Disable terminal rendering: // serve runs on a TTY, so otherwise the task manager draws progress bars over @@ -156,21 +156,35 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve task.SetNoRender(true) mux := http.NewServeMux() - rpcServer.RegisterRoutes(mux) + mux.Handle("GET /api/openapi.json", handleCaptainOpenAPI(openAPISpec, false)) + mux.Handle("GET /api/openapi.yaml", handleCaptainOpenAPI(openAPISpec, true)) + mux.HandleFunc("GET /api/entities", rpcServer.HandleEntities) + mux.HandleFunc("GET /health", rpcServer.HandleHealth) + rpcServer.RegisterExecutionRoutes(mux) mux.HandleFunc("POST /api/captain/chat/threads/from-agent", handleThreadFromAgent(threadStore)) + mux.HandleFunc("GET /api/captain/projects", handleProjects()) mux.HandleFunc("GET /api/captain/sessions/live", handleSessionsLive()) + mux.HandleFunc("GET /api/captain/sessions/throughput", handleSessionsThroughput()) mux.HandleFunc("GET /api/captain/sessions/{id}", handleSessionGet()) + mux.HandleFunc("POST /api/captain/hooks/{provider}", handleMonitorHookEvent()) mux.HandleFunc("GET /api/captain/ai/permissions/catalog", handlePermissionCatalog(cwd)) mux.HandleFunc("GET /api/captain/ai/prompt/schema", handlePromptSchema()) + registerProviderTokenHandlers(mux) + registerProviderDefaultsHandlers(mux) mux.HandleFunc("GET /api/captain/secrets/resources", handleSecretResources()) mux.HandleFunc("GET /api/captain/secrets/preview", handleSecretPreview()) - // Task tracking: /api/captain/tasks, /tasks/stream, /tasks/{id}; plus the - // run-list SSE the clicky-ui useTaskRuns hook subscribes to. + mux.HandleFunc("POST /api/attachments", handleAttachmentUpload(attachmentStore)) + mux.HandleFunc("GET /api/attachments/{id}", handleAttachmentGet(attachmentStore)) + // Task tracking: /api/captain/tasks, /tasks/stream, /tasks/{id}, and the + // /tasks/runs/stream run-list SSE the clicky-ui useTaskRuns hook subscribes to. task.RegisterHandlers(mux, "/api/captain") - mux.Handle("GET /api/captain/tasks/runs/stream", task.RunsSSEHandler(nil)) // Live session history for a prompt run. mux.Handle("GET /api/captain/prompt/runs/{runId}/stream", handlePromptRunStream(promptRuns)) mux.Handle("GET /api/captain/prompt/runs/{runId}", handlePromptRunSnapshot(promptRuns)) + mux.Handle("POST /api/captain/prompt/runs/{runId}/message", handlePromptRunMessage(promptChats)) + mux.Handle("POST /api/captain/prompt/runs/{runId}/interrupt", handlePromptRunInterrupt(promptChats)) + mux.Handle("POST /api/captain/prompt/runs/{runId}/stop", handlePromptRunStop(promptRuns, promptChats)) + mux.Handle("POST /api/captain/sessions/{id}/message", handleSessionMessage(promptChats)) mux.Handle("/api/chat", chat.Handler()) mux.Handle("/api/chat/", chat.Handler()) @@ -180,7 +194,12 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve } mux.Handle("/", uiHandler) - addr := fmt.Sprintf("%s:%d", opts.Host, opts.Port) + // Export the serve URL so every captain-launched agent (and the hook + // receiver subprocesses its sessions spawn) delivers hook events to this + // instance even off the default port. + if err := os.Setenv(api.ServeURLEnv, "http://"+addr); err != nil { + return err + } httpSrv := &http.Server{ Addr: addr, Handler: rpchttp.TimingMiddleware(PromptDirsMiddleware(mux, opts.PromptDirs)), @@ -192,7 +211,34 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve ctx, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM) defer stop() + db, err := captainServeDB(ctx) + if err != nil { + return err + } + mon, err := monitor.New(monitor.Config{DB: db, HostID: captainHostID()}) + if err != nil { + return err + } + setServeMonitor(mon) + go func() { + if err := mon.Run(ctx); err != nil { + log.Errorf("session monitor stopped: %v", err) + } + }() + select { + case <-mon.Ready(): + case <-ctx.Done(): + return ctx.Err() + } + liveSessions, err := db.CountLiveRootSessions(ctx) + if err != nil { + return err + } + dsn, source := captainDatabaseIdentity() + log.Infof("Database Info: source=%q dsn=%q live_sessions=%d", source, database.MaskDSN(dsn), liveSessions) + go prunePromptRuns(ctx, promptRuns) + defer promptChats.stopAll() errCh := make(chan error, 1) go func() { @@ -200,14 +246,14 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve fmt.Fprintf(stdout, " UI: http://%s/\n", addr) fmt.Fprintf(stdout, " OpenAPI JSON: http://%s/api/openapi.json\n", addr) fmt.Fprintf(stdout, " AI Chat: http://%s/api/chat\n", addr) - if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed { + if err := httpSrv.Serve(listener); err != nil && err != http.ErrServerClosed { errCh <- err } }() var vite *exec.Cmd if opts.Dev { - vite, err = startCaptainViteDevServer(ctx, opts.Host, opts.Port, opts.UIPort, stdout, stderr) + vite, err = startCaptainViteDevServer(ctx, opts.Host, servePort, opts.UIPort, stdout, stderr) if err != nil { _ = httpSrv.Close() return err @@ -243,157 +289,6 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve } } -func handleSessionsLive() http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - query := r.URL.Query() - opts := SessionLiveOptions{ - Source: strings.TrimSpace(query.Get("source")), - Query: strings.TrimSpace(query.Get("q")), - Limit: 100, - } - if opts.Source == "" { - opts.Source = "all" - } - if raw := strings.TrimSpace(query.Get("all")); raw != "" { - opts.All, _ = strconv.ParseBool(raw) - } - if raw := strings.TrimSpace(query.Get("limit")); raw != "" { - limit, err := strconv.Atoi(raw) - if err != nil { - http.Error(w, fmt.Sprintf("invalid limit %q", raw), http.StatusBadRequest) - return - } - opts.Limit = limit - } - - result, err := RunSessionLive(r.Context(), opts) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - writeServeJSON(w, http.StatusOK, result) - } -} - -// handleSessionGet serves the unified session model for one session id at -// GET /api/captain/sessions/{id}, so the web UI can render the same model the -// CLI `sessions get` returns. -func handleSessionGet() http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - id := strings.TrimSpace(r.PathValue("id")) - if id == "" { - http.Error(w, "session id is required", http.StatusBadRequest) - return - } - source := strings.TrimSpace(r.URL.Query().Get("source")) - if source == "" { - source = "all" - } - s, err := RunSessionGet(r.Context(), SessionGetOptions{ID: id, Source: source}) - if err != nil { - http.Error(w, err.Error(), http.StatusNotFound) - return - } - writeServeJSON(w, http.StatusOK, s) - } -} - -func handleThreadFromAgent(store aichat.ThreadStore) http.HandlerFunc { - type request struct { - Title string `json:"title"` - ProviderSessionID string `json:"providerSessionId"` - Model string `json:"model"` - } - type response struct { - *aichat.Thread - LaunchURL string `json:"launchUrl"` - } - - return func(w http.ResponseWriter, r *http.Request) { - var body request - if err := json.NewDecoder(r.Body).Decode(&body); err != nil { - http.Error(w, fmt.Sprintf("invalid request body: %v", err), http.StatusBadRequest) - return - } - body.ProviderSessionID = strings.TrimSpace(body.ProviderSessionID) - if body.ProviderSessionID == "" { - http.Error(w, "providerSessionId is required", http.StatusBadRequest) - return - } - if strings.TrimSpace(body.Title) == "" { - body.Title = "Captain agent " + body.ProviderSessionID - } - thread, err := store.Create(r.Context(), body.Title) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - if err := store.SetProviderSession(r.Context(), thread.ID, body.ProviderSessionID); err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - thread, err = store.Get(r.Context(), thread.ID) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - launchURL := "/chat/" + url.PathEscape(thread.ID) - if model := strings.TrimSpace(body.Model); model != "" { - launchURL += "?model=" + url.QueryEscape(model) - } - writeServeJSON(w, http.StatusCreated, response{ - Thread: thread, - LaunchURL: launchURL, - }) - } -} - -func captainChatToolEnabled(tool aichat.ToolInfo) bool { - raw := strings.ToLower(strings.TrimSpace(tool.OperationName)) - if raw == "" { - raw = strings.ToLower(strings.TrimSpace(tool.Name)) - } - raw = strings.ReplaceAll(raw, "/", " ") - fields := strings.FieldsFunc(raw, func(r rune) bool { - return r == ' ' || r == '_' || r == '-' - }) - if len(fields) == 0 { - return true - } - switch fields[0] { - case "serve", "mcp", "hook", "container", "sandbox", "port", "configure", "ai": - return false - default: - return true - } -} - -func captainChatRequiresApproval(tool aichat.ToolInfo, _ any) bool { - switch strings.ToUpper(tool.Method) { - case http.MethodGet, http.MethodHead, http.MethodOptions: - return false - } - if isReadOnlyCaptainTool(tool.ClickyVerb) || - isReadOnlyCaptainTool(tool.Name) || - isReadOnlyCaptainTool(tool.OperationName) || - isReadOnlyCaptainTool(tool.Path) { - return false - } - return true -} - -func isReadOnlyCaptainTool(value string) bool { - for _, part := range strings.FieldsFunc(strings.ToLower(value), func(r rune) bool { - return r == '_' || r == '-' || r == '/' || r == '.' || r == ' ' - }) { - switch part { - case "list", "get", "read", "show", "lookup", "search", "history", "info", "cost", "changes", "models", "status", "check": - return true - } - } - return false -} - func startCaptainViteDevServer(ctx context.Context, apiHost string, apiPort, uiPort int, stdout, stderr io.Writer) (*exec.Cmd, error) { webappDir, err := captainWebappDevDir() if err != nil { diff --git a/pkg/cli/serve_chat.go b/pkg/cli/serve_chat.go new file mode 100644 index 0000000..ca95395 --- /dev/null +++ b/pkg/cli/serve_chat.go @@ -0,0 +1,150 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + + "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/api/registry" + "github.com/flanksource/captain/pkg/attachments" + clickyaichat "github.com/flanksource/clicky/aichat" + "github.com/flanksource/commons-db/shell" + "github.com/spf13/cobra" +) + +func newCaptainChatService( + ctx context.Context, + rootCmd *cobra.Command, + opts ServeOptions, + cwd string, + threadStore aichat.ThreadStore, + attachmentStore *attachments.Store, +) (*aichat.Service, *aichat.MCPToolProvider, error) { + chatTools, err := clickyaichat.NewCobraToolProvider(clickyaichat.CobraToolProviderOptions{ + Root: rootCmd, Filter: captainChatToolEnabled, Permission: captainChatToolPermission, + }) + if err != nil { + return nil, nil, err + } + mcpTools := aichat.NewMCPToolProvider(aichat.MCPToolProviderOptions{Servers: opts.MCPServers}) + if _, err := mcpTools.ToolSet(ctx); err != nil { + return nil, nil, err + } + chat := aichat.NewService(aichat.ServiceOptions{ + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{ + System: "You are Captain's coding-agent launcher assistant. Use Captain and Clicky tools when useful, " + + "prefer read-only inspection unless the user explicitly asks for edits, and keep follow-up guidance concise.", + Spec: api.Spec{ + Model: api.Model{Name: "sol", Mode: registry.ModeAgent}, + Setup: &shell.Setup{Cwd: cwd}, + }, + }, nil + }), + Tools: chatTools, MCP: mcpTools, Threads: threadStore, + Attachments: chatAttachmentResolver{store: attachmentStore}, + }) + return chat, mcpTools, nil +} + +func handleThreadFromAgent(store aichat.ThreadStore) http.HandlerFunc { + type request struct { + Title string `json:"title"` + ProviderSessionID string `json:"providerSessionId"` + Model string `json:"model"` + } + type response struct { + *aichat.Thread + LaunchURL string `json:"launchUrl"` + } + + return func(w http.ResponseWriter, r *http.Request) { + var body request + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, fmt.Sprintf("invalid request body: %v", err), http.StatusBadRequest) + return + } + body.ProviderSessionID = strings.TrimSpace(body.ProviderSessionID) + if body.ProviderSessionID == "" { + http.Error(w, "providerSessionId is required", http.StatusBadRequest) + return + } + if strings.TrimSpace(body.Title) == "" { + body.Title = "Captain agent " + body.ProviderSessionID + } + thread, err := store.Create(r.Context(), body.Title) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := store.SetProviderSession(r.Context(), thread.ID, body.ProviderSessionID); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + thread, err = store.Get(r.Context(), thread.ID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + launchURL := "/chat/" + url.PathEscape(thread.ID) + if model := strings.TrimSpace(body.Model); model != "" { + launchURL += "?model=" + url.QueryEscape(model) + } + writeServeJSON(w, http.StatusCreated, response{Thread: thread, LaunchURL: launchURL}) + } +} + +func captainChatToolEnabled(tool tools.ToolInfo) bool { + raw := strings.ToLower(strings.TrimSpace(tool.Annotation("clicky/operation"))) + if raw == "" { + raw = strings.ToLower(strings.TrimSpace(tool.Name)) + } + raw = strings.ReplaceAll(raw, "/", " ") + fields := strings.FieldsFunc(raw, func(r rune) bool { + return r == ' ' || r == '_' || r == '-' + }) + if len(fields) == 0 { + return true + } + switch fields[0] { + case "serve", "mcp", "hook", "container", "sandbox", "port", "configure", "ai": + return false + default: + return true + } +} + +func captainChatToolPermission(tool tools.ToolInfo) api.ToolMode { + switch strings.ToUpper(tool.Annotation("clicky/method")) { + case http.MethodGet, http.MethodHead, http.MethodOptions: + return api.ToolModeOn + } + if tool.DefaultPermission == api.ToolModeOn { + return api.ToolModeOn + } + if isReadOnlyCaptainTool(tool.Annotation("clicky/verb")) || + isReadOnlyCaptainTool(tool.Name) || + isReadOnlyCaptainTool(tool.Annotation("clicky/operation")) || + isReadOnlyCaptainTool(tool.Annotation("clicky/path")) { + return api.ToolModeOn + } + return api.ToolModeAsk +} + +func isReadOnlyCaptainTool(value string) bool { + for _, part := range strings.FieldsFunc(strings.ToLower(value), func(r rune) bool { + return r == '_' || r == '-' || r == '/' || r == '.' || r == ' ' + }) { + switch part { + case "list", "get", "read", "show", "lookup", "search", "history", "info", "cost", "changes", "models", "status", "check": + return true + } + } + return false +} diff --git a/pkg/cli/serve_hooks.go b/pkg/cli/serve_hooks.go new file mode 100644 index 0000000..3fe7d98 --- /dev/null +++ b/pkg/cli/serve_hooks.go @@ -0,0 +1,42 @@ +package cli + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "github.com/flanksource/captain/pkg/monitor" +) + +// handleMonitorHookEvent accepts normalized hook events from `captain hook +// monitor notify` and enqueues them on the serve monitor. Delivery is +// fire-and-forget for the sender: 202 means enqueued, not ingested. +func handleMonitorHookEvent() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + provider := strings.TrimSpace(r.PathValue("provider")) + if provider != "claude" && provider != "codex" { + http.Error(w, fmt.Sprintf("unknown hook provider %q", provider), http.StatusBadRequest) + return + } + var ev monitor.HookEvent + if err := json.NewDecoder(r.Body).Decode(&ev); err != nil { + http.Error(w, "invalid hook event payload: "+err.Error(), http.StatusBadRequest) + return + } + if strings.TrimSpace(ev.Event) == "" { + http.Error(w, "hook event name is required", http.StatusBadRequest) + return + } + mon := serveMonitor() + if mon == nil { + http.Error(w, "session monitor is not running", http.StatusServiceUnavailable) + return + } + ev.Provider = provider + ev.ReceivedAt = time.Now().UTC() + mon.NotifyHookEvent(ev) + w.WriteHeader(http.StatusAccepted) + } +} diff --git a/pkg/cli/serve_hooks_test.go b/pkg/cli/serve_hooks_test.go new file mode 100644 index 0000000..8d078b7 --- /dev/null +++ b/pkg/cli/serve_hooks_test.go @@ -0,0 +1,48 @@ +package cli + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/monitor" + "github.com/stretchr/testify/assert" +) + +func postHookEventRequest(t *testing.T, provider, body string) *httptest.ResponseRecorder { + t.Helper() + mux := http.NewServeMux() + mux.HandleFunc("POST /api/captain/hooks/{provider}", handleMonitorHookEvent()) + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/captain/hooks/"+provider, strings.NewReader(body)) + mux.ServeHTTP(rec, req) + return rec +} + +func TestHandleMonitorHookEvent(t *testing.T) { + validBody := `{"event":"Stop","sessionId":"abc123","transcriptPath":"/p.jsonl"}` + + t.Run("unknown provider is rejected", func(t *testing.T) { + assert.Equal(t, http.StatusBadRequest, postHookEventRequest(t, "gemini", validBody).Code) + }) + + t.Run("malformed body is rejected", func(t *testing.T) { + assert.Equal(t, http.StatusBadRequest, postHookEventRequest(t, "claude", "not json").Code) + }) + + t.Run("missing event name is rejected", func(t *testing.T) { + assert.Equal(t, http.StatusBadRequest, postHookEventRequest(t, "claude", `{"sessionId":"abc123"}`).Code) + }) + + t.Run("no running monitor is unavailable", func(t *testing.T) { + setServeMonitor(nil) + assert.Equal(t, http.StatusServiceUnavailable, postHookEventRequest(t, "claude", validBody).Code) + }) + + t.Run("valid event is accepted", func(t *testing.T) { + setServeMonitor(&monitor.Monitor{}) + t.Cleanup(func() { setServeMonitor(nil) }) + assert.Equal(t, http.StatusAccepted, postHookEventRequest(t, "claude", validBody).Code) + }) +} diff --git a/pkg/cli/serve_openapi.go b/pkg/cli/serve_openapi.go new file mode 100644 index 0000000..9ee030c --- /dev/null +++ b/pkg/cli/serve_openapi.go @@ -0,0 +1,240 @@ +package cli + +import ( + "net/http" + + "github.com/flanksource/clicky/rpc" +) + +func addCaptainPromptRunPaths(spec *rpc.OpenAPISpec) { + if spec.Paths == nil { + spec.Paths = map[string]rpc.OpenAPIPath{} + } + pathParameter := rpc.OpenAPIParameter{ + Name: "runId", In: "path", Required: true, + Schema: &rpc.OpenAPISchema{Type: "string"}, + } + spec.Paths["/api/captain/prompt/runs/{runId}"] = rpc.OpenAPIPath{ + "get": promptRunOperation("getPromptRun", "Get a prompt run snapshot", pathParameter, nil, jsonResponse(promptRunSnapshotSchema())), + } + spec.Paths["/api/captain/prompt/runs/{runId}/stream"] = rpc.OpenAPIPath{ + "get": promptRunOperation("streamPromptRun", "Stream prompt run events", pathParameter, nil, rpc.OpenAPIResponse{ + Description: "SSE events: run, entry, state, done, and error", + Content: map[string]rpc.OpenAPIMediaType{ + "text/event-stream": {Schema: &rpc.OpenAPISchema{Type: "string"}}, + }, + }), + } + messageBody := &rpc.OpenAPIRequestBody{ + Required: true, + Content: map[string]rpc.OpenAPIMediaType{ + "application/json": {Schema: chatMessageRequestSchema()}, + }, + } + spec.Paths["/api/captain/prompt/runs/{runId}/message"] = rpc.OpenAPIPath{ + "post": promptRunOperation("messagePromptRun", "Send or queue a prompt run message", pathParameter, messageBody, jsonResponse(chatMessageResponseSchema())), + } + spec.Paths["/api/captain/prompt/runs/{runId}/interrupt"] = rpc.OpenAPIPath{ + "post": promptRunOperation("interruptPromptRun", "Pause the active prompt turn", pathParameter, nil, jsonResponse(statusSchema())), + } + spec.Paths["/api/captain/prompt/runs/{runId}/stop"] = rpc.OpenAPIPath{ + "post": promptRunOperation("stopPromptRun", "Stop the entire prompt run", pathParameter, nil, jsonResponse(statusSchema())), + } + sessionParameter := pathParameter + sessionParameter.Name = "id" + spec.Paths["/api/captain/sessions/{id}/message"] = rpc.OpenAPIPath{ + "post": promptRunOperation("messageSession", "Continue a saved or active session", sessionParameter, messageBody, jsonResponse(chatMessageResponseSchema())), + } +} + +func addCaptainProviderTokenPaths(spec *rpc.OpenAPISpec) { + if spec.Paths == nil { + spec.Paths = map[string]rpc.OpenAPIPath{} + } + provider := rpc.OpenAPIParameter{ + Name: "provider", In: "path", Required: true, + Schema: &rpc.OpenAPISchema{Type: "string", Enum: []any{"anthropic", "openai", "gemini", "deepseek"}}, + } + tokenSchema := &rpc.OpenAPISchema{Type: "string", Format: "password", Extensions: map[string]any{"writeOnly": true}} + request := func(required bool) *rpc.OpenAPIRequestBody { + schema := &rpc.OpenAPISchema{Type: "object", Properties: map[string]*rpc.OpenAPISchema{"token": tokenSchema}} + if required { + schema.Required = []string{"token"} + } + return &rpc.OpenAPIRequestBody{Required: true, Content: map[string]rpc.OpenAPIMediaType{ + "application/json": {Schema: schema}, + }} + } + operation := func(id, summary string, body *rpc.OpenAPIRequestBody) rpc.OpenAPIOperation { + return rpc.OpenAPIOperation{ + Tags: []string{"Provider credentials"}, Summary: summary, OperationID: id, + Parameters: []rpc.OpenAPIParameter{provider}, RequestBody: body, + Responses: map[string]rpc.OpenAPIResponse{ + "200": jsonResponse(providerTokenResponseSchema()), + "400": {Description: "Invalid request"}, "403": {Description: "Local same-origin access required"}, + "422": {Description: "Credential rejected"}, "500": {Description: "Vault persistence failed"}, + "502": {Description: "Provider request failed"}, + }, + } + } + spec.Paths["/api/captain/ai/providers/{provider}/token"] = rpc.OpenAPIPath{ + "put": operation("saveProviderToken", "Validate and save a provider token", request(true)), + } + spec.Paths["/api/captain/ai/providers/{provider}/token/test"] = rpc.OpenAPIPath{ + "post": operation("testProviderToken", "Test a candidate or configured provider token", request(false)), + } +} + +func providerTokenResponseSchema() *rpc.OpenAPISchema { + return &rpc.OpenAPISchema{Type: "object", Required: []string{"provider", "valid", "saved", "source", "maskedToken", "modelCount"}, Properties: map[string]*rpc.OpenAPISchema{ + "provider": {Type: "string"}, "valid": {Type: "boolean"}, "saved": {Type: "boolean"}, + "source": {Type: "string"}, "maskedToken": {Type: "string"}, "modelCount": {Type: "integer"}, + }} +} + +func addCaptainProviderDefaultsPaths(spec *rpc.OpenAPISpec) { + if spec.Paths == nil { + spec.Paths = map[string]rpc.OpenAPIPath{} + } + provider := rpc.OpenAPIParameter{ + Name: "provider", In: "path", Required: true, + Schema: &rpc.OpenAPISchema{Type: "string", Enum: []any{"anthropic", "openai", "gemini", "deepseek"}}, + } + stringField := func() *rpc.OpenAPISchema { return &rpc.OpenAPISchema{Type: "string"} } + defaultsRequest := &rpc.OpenAPIRequestBody{Required: true, Content: map[string]rpc.OpenAPIMediaType{ + "application/json": {Schema: &rpc.OpenAPISchema{ + Type: "object", Required: []string{"agent", "model", "effort"}, + Properties: map[string]*rpc.OpenAPISchema{"agent": stringField(), "model": stringField(), "effort": stringField()}, + }}, + }} + defaultsResponse := &rpc.OpenAPISchema{ + Type: "object", Required: []string{"provider", "agent", "model", "effort", "active"}, + Properties: map[string]*rpc.OpenAPISchema{ + "provider": stringField(), "agent": stringField(), "model": stringField(), + "effort": stringField(), "active": {Type: "boolean"}, + }, + } + spec.Paths["/api/captain/ai/providers/{provider}/defaults"] = rpc.OpenAPIPath{"put": { + Tags: []string{"Provider configuration"}, Summary: "Save provider runtime defaults", OperationID: "saveProviderDefaults", + Parameters: []rpc.OpenAPIParameter{provider}, RequestBody: defaultsRequest, + Responses: map[string]rpc.OpenAPIResponse{ + "200": jsonResponse(defaultsResponse), "400": {Description: "Invalid request"}, + "403": {Description: "Local same-origin access required"}, "422": {Description: "Defaults rejected"}, + "500": {Description: "Configuration persistence failed"}, + }, + }} + activeRequest := &rpc.OpenAPIRequestBody{Required: true, Content: map[string]rpc.OpenAPIMediaType{ + "application/json": {Schema: &rpc.OpenAPISchema{ + Type: "object", Required: []string{"provider"}, Properties: map[string]*rpc.OpenAPISchema{"provider": stringField()}, + }}, + }} + spec.Paths["/api/captain/ai/default-provider"] = rpc.OpenAPIPath{"put": { + Tags: []string{"Provider configuration"}, Summary: "Select the provider for flagless runs", OperationID: "saveDefaultProvider", + RequestBody: activeRequest, Responses: map[string]rpc.OpenAPIResponse{ + "200": jsonResponse(&rpc.OpenAPISchema{Type: "object", Required: []string{"provider"}, Properties: map[string]*rpc.OpenAPISchema{"provider": stringField()}}), + "400": {Description: "Invalid request"}, "403": {Description: "Local same-origin access required"}, + "500": {Description: "Configuration persistence failed"}, + }, + }} +} + +func promptRunOperation( + id, summary string, + parameter rpc.OpenAPIParameter, + body *rpc.OpenAPIRequestBody, + response rpc.OpenAPIResponse, +) rpc.OpenAPIOperation { + return rpc.OpenAPIOperation{ + Tags: []string{"Prompt runs"}, Summary: summary, OperationID: id, + Parameters: []rpc.OpenAPIParameter{parameter}, RequestBody: body, + Responses: map[string]rpc.OpenAPIResponse{ + "202": response, + "200": response, + "400": {Description: "Invalid request"}, + "404": {Description: "Run or session not found"}, + "409": {Description: "Run state conflict"}, + "422": {Description: "Session cannot be resumed"}, + "500": {Description: "Provider operation failed"}, + }, + } +} + +func jsonResponse(schema *rpc.OpenAPISchema) rpc.OpenAPIResponse { + return rpc.OpenAPIResponse{ + Description: "Successful response", + Content: map[string]rpc.OpenAPIMediaType{ + "application/json": {Schema: schema}, + }, + } +} + +func chatMessageRequestSchema() *rpc.OpenAPISchema { + return &rpc.OpenAPISchema{ + Type: "object", Required: []string{"text"}, + Properties: map[string]*rpc.OpenAPISchema{ + "text": {Type: "string"}, "messageId": {Type: "string"}, + "model": {Type: "string"}, "backend": {Type: "string"}, + }, + } +} + +func chatMessageResponseSchema() *rpc.OpenAPISchema { + return &rpc.OpenAPISchema{ + Type: "object", Required: []string{"runId", "messageId", "status", "capabilities"}, + Properties: map[string]*rpc.OpenAPISchema{ + "runId": {Type: "string"}, "messageId": {Type: "string"}, + "status": {Type: "string", Enum: []any{"steered", "queued", "started"}}, + "capabilities": chatCapabilitiesSchema(), + }, + } +} + +func chatCapabilitiesSchema() *rpc.OpenAPISchema { + boolean := func() *rpc.OpenAPISchema { return &rpc.OpenAPISchema{Type: "boolean"} } + return &rpc.OpenAPISchema{ + Type: "object", + Properties: map[string]*rpc.OpenAPISchema{ + "interrupt": boolean(), "steer": boolean(), "followUp": boolean(), "resume": boolean(), + }, + } +} + +func statusSchema() *rpc.OpenAPISchema { + return &rpc.OpenAPISchema{ + Type: "object", Required: []string{"status"}, + Properties: map[string]*rpc.OpenAPISchema{"status": {Type: "string"}}, + } +} + +func promptRunSnapshotSchema() *rpc.OpenAPISchema { + return &rpc.OpenAPISchema{ + Type: "object", Required: []string{"run", "entries", "done"}, + Properties: map[string]*rpc.OpenAPISchema{ + "run": {Type: "object"}, + "entries": {Type: "array", Items: &rpc.OpenAPISchema{Type: "object"}}, + "state": {Type: "object", Nullable: true}, + "done": {Type: "boolean"}, + "summary": {Type: "object", Nullable: true}, + "error": {Type: "string"}, + }, + } +} + +func handleCaptainOpenAPI(spec *rpc.OpenAPISpec, yaml bool) http.HandlerFunc { + return func(w http.ResponseWriter, _ *http.Request) { + var body []byte + var err error + if yaml { + w.Header().Set("Content-Type", "application/yaml") + body, err = spec.ToYAML() + } else { + w.Header().Set("Content-Type", "application/json") + body, err = spec.ToJSON() + } + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + _, _ = w.Write(body) + } +} diff --git a/pkg/cli/serve_port.go b/pkg/cli/serve_port.go new file mode 100644 index 0000000..f9b0ab2 --- /dev/null +++ b/pkg/cli/serve_port.go @@ -0,0 +1,44 @@ +package cli + +import ( + "fmt" + "net" + "strconv" + "strings" +) + +func effectiveServePort(dev, portFlagSet bool, configuredPort int) int { + if dev && !portFlagSet { + return 0 + } + return configuredPort +} + +func (o ServeOptions) validate() error { + if strings.TrimSpace(o.Host) == "" { + return fmt.Errorf("host cannot be empty") + } + if o.Port < 0 || o.Port > 65535 || (!o.Dev && o.Port == 0) { + return fmt.Errorf("invalid --port %d", o.Port) + } + if o.Dev && (o.UIPort < 1 || o.UIPort > 65535) { + return fmt.Errorf("invalid --ui-port %d", o.UIPort) + } + if strings.TrimSpace(o.ThreadsFile) == "" { + return fmt.Errorf("threads file cannot be empty") + } + return ValidatePromptDirs(o.PromptDirs) +} + +func listenCaptainServer(host string, port int) (net.Listener, string, int, error) { + listener, err := net.Listen("tcp", net.JoinHostPort(host, strconv.Itoa(port))) + if err != nil { + return nil, "", 0, fmt.Errorf("listen on %s:%d: %w", host, port, err) + } + tcpAddr, ok := listener.Addr().(*net.TCPAddr) + if !ok { + _ = listener.Close() + return nil, "", 0, fmt.Errorf("unexpected listener address %T", listener.Addr()) + } + return listener, net.JoinHostPort(host, strconv.Itoa(tcpAddr.Port)), tcpAddr.Port, nil +} diff --git a/pkg/cli/serve_port_ginkgo_test.go b/pkg/cli/serve_port_ginkgo_test.go new file mode 100644 index 0000000..faf446b --- /dev/null +++ b/pkg/cli/serve_port_ginkgo_test.go @@ -0,0 +1,41 @@ +package cli + +import ( + "net" + "strconv" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("serve ports", func() { + DescribeTable("selecting the API port", + func(dev, portFlagSet bool, configuredPort, expectedPort int) { + Expect(effectiveServePort(dev, portFlagSet, configuredPort)).To(Equal(expectedPort)) + }, + Entry("uses an ephemeral port for development", true, false, 9020, 0), + Entry("preserves an explicit development port", true, true, 9021, 9021), + Entry("preserves the default production port", false, false, 9020, 9020), + ) + + It("accepts an ephemeral port only in development", func() { + options := ServeOptions{Host: "localhost", Port: 0, Dev: true, UIPort: 5183, ThreadsFile: "threads.json"} + Expect(options.validate()).To(Succeed()) + + options.Dev = false + Expect(options.validate()).To(MatchError("invalid --port 0")) + }) + + It("keeps the ephemeral API port reserved for the server", func() { + listener, addr, port, err := listenCaptainServer("127.0.0.1", 0) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(listener.Close) + + Expect(port).To(BeNumerically(">", 0)) + Expect(addr).To(Equal(net.JoinHostPort("127.0.0.1", strconv.Itoa(port)))) + + second, err := net.Listen("tcp", addr) + Expect(err).To(HaveOccurred()) + Expect(second).To(BeNil()) + }) +}) diff --git a/pkg/cli/serve_provider_defaults.go b/pkg/cli/serve_provider_defaults.go new file mode 100644 index 0000000..13aca0c --- /dev/null +++ b/pkg/cli/serve_provider_defaults.go @@ -0,0 +1,122 @@ +package cli + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "mime" + "net/http" + "strings" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" +) + +type providerDefaultsRequest struct { + Agent string `json:"agent"` + Model string `json:"model"` + Effort string `json:"effort"` +} + +type providerDefaultsResponse struct { + Provider string `json:"provider"` + Agent string `json:"agent"` + Model string `json:"model"` + Effort string `json:"effort"` + Active bool `json:"active"` +} + +type defaultProviderRequest struct { + Provider string `json:"provider"` +} + +func registerProviderDefaultsHandlers(mux *http.ServeMux) { + mux.HandleFunc("PUT /api/captain/ai/providers/{provider}/defaults", handleProviderDefaults) + mux.HandleFunc("PUT /api/captain/ai/default-provider", handleDefaultProvider) +} + +func handleProviderDefaults(w http.ResponseWriter, r *http.Request) { + if err := validateLocalConfigurationRequest(r); err != nil { + http.Error(w, err.Error(), http.StatusForbidden) + return + } + provider := api.Backend(strings.TrimSpace(r.PathValue("provider"))) + if !configurableAPIBackend(provider) { + http.Error(w, "provider must be one of: anthropic, openai, gemini, deepseek", http.StatusBadRequest) + return + } + var request providerDefaultsRequest + if err := decodeConfigurationRequest(w, r, &request); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + defaults := ProviderDefaultView{ + Agent: strings.TrimSpace(request.Agent), Model: strings.TrimSpace(request.Model), + Effort: strings.TrimSpace(request.Effort), Configured: true, + } + if err := validateProviderDefaults(r.Context(), provider, defaults); err != nil { + http.Error(w, err.Error(), http.StatusUnprocessableEntity) + return + } + active := false + if err := captainconfig.Update(func(cfg *captainconfig.Config) error { + if cfg.AI.Providers == nil { + cfg.AI.Providers = map[string]captainconfig.ProviderDefaults{} + } + cfg.AI.Providers[string(provider)] = captainconfig.ProviderDefaults{ + Agent: defaults.Agent, Model: defaults.Model, ReasoningEffort: defaults.Effort, + } + active = cfg.AI.ActiveProvider() == string(provider) + return nil + }); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + writeConfigurationJSON(w, providerDefaultsResponse{ + Provider: string(provider), Agent: defaults.Agent, Model: defaults.Model, + Effort: defaults.Effort, Active: active, + }) +} + +func handleDefaultProvider(w http.ResponseWriter, r *http.Request) { + if err := validateLocalConfigurationRequest(r); err != nil { + http.Error(w, err.Error(), http.StatusForbidden) + return + } + var request defaultProviderRequest + if err := decodeConfigurationRequest(w, r, &request); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + provider := api.Backend(strings.TrimSpace(request.Provider)) + if !configurableAPIBackend(provider) { + http.Error(w, "provider must be one of: anthropic, openai, gemini, deepseek", http.StatusBadRequest) + return + } + if err := captainconfig.Update(func(cfg *captainconfig.Config) error { + cfg.AI.DefaultProvider = string(provider) + return nil + }); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + writeConfigurationJSON(w, defaultProviderRequest{Provider: string(provider)}) +} + +func decodeConfigurationRequest(w http.ResponseWriter, r *http.Request, target any) error { + mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type")) + if err != nil || mediaType != "application/json" { + return fmt.Errorf("Content-Type must be application/json") + } + r.Body = http.MaxBytesReader(w, r.Body, providerTokenBodyLimit) + decoder := json.NewDecoder(r.Body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(target); err != nil { + return fmt.Errorf("decode request: %w", err) + } + if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { + return fmt.Errorf("request must contain one JSON object") + } + return nil +} diff --git a/pkg/cli/serve_provider_defaults_test.go b/pkg/cli/serve_provider_defaults_test.go new file mode 100644 index 0000000..bd6e9c6 --- /dev/null +++ b/pkg/cli/serve_provider_defaults_test.go @@ -0,0 +1,117 @@ +package cli + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/captainconfig" + "github.com/flanksource/clicky/rpc" +) + +func TestProviderDefaultsPutValidatesAndSaves(t *testing.T) { + setupProviderDefaultsTest(t) + previous := configureDefaultsModels + configureDefaultsModels = func(_ context.Context, agent api.Backend) ([]ai.ModelDef, error) { + if agent != api.BackendCodexAgent { + t.Fatalf("agent = %s", agent) + } + return []ai.ModelDef{{ID: "gpt-5.6-sol"}}, nil + } + t.Cleanup(func() { configureDefaultsModels = previous }) + + response := serveProviderDefaultsRequest(t, "/api/captain/ai/providers/openai/defaults", + `{"agent":"codex-agent","model":"gpt-5.6-sol","effort":"high"}`) + if response.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", response.Code, response.Body.String()) + } + var result providerDefaultsResponse + if err := json.Unmarshal(response.Body.Bytes(), &result); err != nil { + t.Fatalf("decode response: %v", err) + } + if result.Agent != "codex-agent" || result.Model != "gpt-5.6-sol" || result.Effort != "high" { + t.Fatalf("response = %+v", result) + } + config, _, err := captainconfig.Load() + if err != nil || config.AI.Providers["openai"].Model != "gpt-5.6-sol" { + t.Fatalf("config=%+v err=%v", config, err) + } +} + +func TestProviderDefaultsPutFailurePreservesExistingDefaults(t *testing.T) { + setupProviderDefaultsTest(t) + if err := captainconfig.Save(captainconfig.Config{AI: captainconfig.AIDefaults{ + Providers: map[string]captainconfig.ProviderDefaults{ + "openai": {Agent: "openai", Model: "gpt-existing", ReasoningEffort: "medium"}, + }, + }}); err != nil { + t.Fatalf("seed config: %v", err) + } + previous := configureDefaultsModels + configureDefaultsModels = func(context.Context, api.Backend) ([]ai.ModelDef, error) { + return []ai.ModelDef{{ID: "gpt-valid"}}, nil + } + t.Cleanup(func() { configureDefaultsModels = previous }) + + response := serveProviderDefaultsRequest(t, "/api/captain/ai/providers/openai/defaults", + `{"agent":"openai","model":"gpt-invalid","effort":"high"}`) + if response.Code != http.StatusUnprocessableEntity { + t.Fatalf("status=%d body=%s", response.Code, response.Body.String()) + } + config, _, err := captainconfig.Load() + if err != nil || config.AI.Providers["openai"].Model != "gpt-existing" { + t.Fatalf("config=%+v err=%v", config, err) + } +} + +func TestDefaultProviderPutChangesOnlyActiveProvider(t *testing.T) { + setupProviderDefaultsTest(t) + if err := captainconfig.Save(captainconfig.Config{Prompts: captainconfig.PromptDefaults{Dirs: []string{"/repo/prompts"}}}); err != nil { + t.Fatalf("seed config: %v", err) + } + response := serveProviderDefaultsRequest(t, "/api/captain/ai/default-provider", `{"provider":"gemini"}`) + if response.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", response.Code, response.Body.String()) + } + config, _, err := captainconfig.Load() + if err != nil || config.AI.DefaultProvider != "gemini" || len(config.Prompts.Dirs) != 1 { + t.Fatalf("config=%+v err=%v", config, err) + } +} + +func TestProviderDefaultsOpenAPIPaths(t *testing.T) { + spec := &rpc.OpenAPISpec{} + addCaptainProviderDefaultsPaths(spec) + if _, ok := spec.Paths["/api/captain/ai/providers/{provider}/defaults"]; !ok { + t.Fatal("missing provider defaults OpenAPI path") + } + if _, ok := spec.Paths["/api/captain/ai/default-provider"]; !ok { + t.Fatal("missing default provider OpenAPI path") + } +} + +func setupProviderDefaultsTest(t *testing.T) { + t.Helper() + captainconfig.SetPathForTesting(filepath.Join(t.TempDir(), ".captain.yaml")) + t.Cleanup(func() { captainconfig.SetPathForTesting("") }) +} + +func serveProviderDefaultsRequest(t *testing.T, path, body string) *httptest.ResponseRecorder { + t.Helper() + mux := http.NewServeMux() + registerProviderDefaultsHandlers(mux) + request := httptest.NewRequest(http.MethodPut, path, strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Origin", "http://localhost:9020") + request.Host = "localhost:9020" + request.RemoteAddr = "127.0.0.1:1234" + response := httptest.NewRecorder() + mux.ServeHTTP(response, request) + return response +} diff --git a/pkg/cli/serve_provider_tokens.go b/pkg/cli/serve_provider_tokens.go new file mode 100644 index 0000000..d197f33 --- /dev/null +++ b/pkg/cli/serve_provider_tokens.go @@ -0,0 +1,155 @@ +package cli + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "mime" + "net" + "net/http" + "net/url" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/credentials" +) + +const providerTokenBodyLimit = 8 << 10 + +type providerTokenRequest struct { + Token string `json:"token"` +} + +type providerTokenResponse struct { + Provider string `json:"provider"` + Valid bool `json:"valid"` + Saved bool `json:"saved"` + Source string `json:"source"` + MaskedToken string `json:"maskedToken"` + ModelCount int `json:"modelCount"` +} + +func registerProviderTokenHandlers(mux *http.ServeMux) { + mux.Handle("PUT /api/captain/ai/providers/{provider}/token", handleProviderToken(false)) + mux.Handle("POST /api/captain/ai/providers/{provider}/token/test", handleProviderToken(true)) +} + +func handleProviderToken(testOnly bool) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := validateLocalConfigurationRequest(r); err != nil { + http.Error(w, err.Error(), http.StatusForbidden) + return + } + backend := ai.Backend(strings.TrimSpace(r.PathValue("provider"))) + if !configurableAPIBackend(backend) { + http.Error(w, "provider must be one of: anthropic, openai, gemini, deepseek", http.StatusBadRequest) + return + } + request, err := decodeProviderTokenRequest(w, r) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + token, source := strings.TrimSpace(request.Token), "candidate" + if token == "" { + if !testOnly { + http.Error(w, "token cannot be empty", http.StatusBadRequest) + return + } + resolved, err := ai.ResolveAPIKey(backend) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + token, source = resolved.Token, resolved.Source + if token == "" { + http.Error(w, fmt.Sprintf("no credential configured for %s", backend), http.StatusBadRequest) + return + } + } + + models, err := configureTokenModels(r.Context(), backend, token) + if err != nil { + http.Error(w, fmt.Sprintf("validate %s credential: %v", backend, err), providerValidationStatus(err)) + return + } + if !testOnly { + vault, err := credentials.DefaultVault() + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := vault.Set(string(backend), token); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + source = credentials.SourceVault + } + writeConfigurationJSON(w, providerTokenResponse{ + Provider: string(backend), Valid: true, Saved: !testOnly, + Source: source, MaskedToken: ai.MaskKey(token), ModelCount: len(models), + }) + }) +} + +func validateLocalConfigurationRequest(r *http.Request) error { + host, _, err := net.SplitHostPort(r.RemoteAddr) + if err != nil { + host = r.RemoteAddr + } + ip := net.ParseIP(host) + if ip == nil || !ip.IsLoopback() { + return fmt.Errorf("configuration changes are restricted to loopback clients") + } + requestHost := r.Host + if parsed, _, err := net.SplitHostPort(requestHost); err == nil { + requestHost = parsed + } + requestIP := net.ParseIP(requestHost) + if !strings.EqualFold(requestHost, "localhost") && (requestIP == nil || !requestIP.IsLoopback()) { + return fmt.Errorf("configuration changes require a loopback request host") + } + origin := strings.TrimSpace(r.Header.Get("Origin")) + if origin == "" { + return nil + } + parsed, err := url.Parse(origin) + if err != nil || !strings.EqualFold(parsed.Host, r.Host) { + return fmt.Errorf("configuration changes require a same-origin request") + } + return nil +} + +func decodeProviderTokenRequest(w http.ResponseWriter, r *http.Request) (providerTokenRequest, error) { + mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type")) + if err != nil || mediaType != "application/json" { + return providerTokenRequest{}, fmt.Errorf("Content-Type must be application/json") + } + r.Body = http.MaxBytesReader(w, r.Body, providerTokenBodyLimit) + decoder := json.NewDecoder(r.Body) + decoder.DisallowUnknownFields() + var request providerTokenRequest + if err := decoder.Decode(&request); err != nil { + return providerTokenRequest{}, fmt.Errorf("decode request: %w", err) + } + if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { + return providerTokenRequest{}, fmt.Errorf("request must contain one JSON object") + } + return request, nil +} + +func providerValidationStatus(err error) int { + var httpErr ai.ModelHTTPError + if errors.As(err, &httpErr) && (httpErr.StatusCode == http.StatusUnauthorized || httpErr.StatusCode == http.StatusForbidden) { + return http.StatusUnprocessableEntity + } + return http.StatusBadGateway +} + +func writeConfigurationJSON(w http.ResponseWriter, response any) { + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(response); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } +} diff --git a/pkg/cli/serve_provider_tokens_test.go b/pkg/cli/serve_provider_tokens_test.go new file mode 100644 index 0000000..21d690c --- /dev/null +++ b/pkg/cli/serve_provider_tokens_test.go @@ -0,0 +1,166 @@ +package cli + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/credentials" + "github.com/flanksource/clicky/rpc" +) + +func TestProviderTokenPutValidatesAndSavesWithoutEchoingSecret(t *testing.T) { + setupProviderTokenTest(t) + const secret = "submitted-provider-secret" + previous := configureTokenModels + configureTokenModels = func(_ context.Context, backend ai.Backend, token string) ([]ai.ModelDef, error) { + if backend != ai.BackendGemini || token != secret { + t.Fatalf("validation got backend=%s token=%q", backend, token) + } + return []ai.ModelDef{{ID: "gemini-example"}}, nil + } + t.Cleanup(func() { configureTokenModels = previous }) + + response := serveProviderTokenRequest(t, http.MethodPut, "/api/captain/ai/providers/gemini/token", `{"token":"`+secret+`"}`) + if response.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", response.Code, response.Body.String()) + } + if strings.Contains(response.Body.String(), secret) { + t.Fatalf("response exposed token: %s", response.Body.String()) + } + var result providerTokenResponse + if err := json.Unmarshal(response.Body.Bytes(), &result); err != nil { + t.Fatalf("decode response: %v", err) + } + if !result.Valid || !result.Saved || result.ModelCount != 1 || result.Source != credentials.SourceVault { + t.Fatalf("response = %+v", result) + } + vault, _ := credentials.DefaultVault() + values, err := vault.Load() + if err != nil || values["gemini"] != secret { + t.Fatalf("vault=%v err=%v", values, err) + } +} + +func TestProviderTokenPutRejectsCredentialWithoutReplacingOldToken(t *testing.T) { + setupProviderTokenTest(t) + vault, _ := credentials.DefaultVault() + if err := vault.Set("openai", "existing-provider-secret"); err != nil { + t.Fatalf("seed vault: %v", err) + } + previous := configureTokenModels + configureTokenModels = func(context.Context, ai.Backend, string) ([]ai.ModelDef, error) { + return nil, ai.ModelHTTPError{Backend: ai.BackendOpenAI, StatusCode: http.StatusUnauthorized} + } + t.Cleanup(func() { configureTokenModels = previous }) + + response := serveProviderTokenRequest(t, http.MethodPut, "/api/captain/ai/providers/openai/token", `{"token":"invalid-provider-secret"}`) + if response.Code != http.StatusUnprocessableEntity { + t.Fatalf("status=%d body=%s", response.Code, response.Body.String()) + } + values, err := vault.Load() + if err != nil || values["openai"] != "existing-provider-secret" { + t.Fatalf("vault=%v err=%v", values, err) + } +} + +func TestProviderTokenTestCurrentDoesNotWrite(t *testing.T) { + setupProviderTokenTest(t) + t.Setenv("ANTHROPIC_API_KEY", "environment-provider-secret") + previous := configureTokenModels + configureTokenModels = func(_ context.Context, _ ai.Backend, token string) ([]ai.ModelDef, error) { + if token != "environment-provider-secret" { + t.Fatalf("token=%q", token) + } + return []ai.ModelDef{{ID: "claude-example"}}, nil + } + t.Cleanup(func() { configureTokenModels = previous }) + + response := serveProviderTokenRequest(t, http.MethodPost, "/api/captain/ai/providers/anthropic/token/test", `{}`) + if response.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", response.Code, response.Body.String()) + } + var result providerTokenResponse + if err := json.Unmarshal(response.Body.Bytes(), &result); err != nil { + t.Fatalf("decode response: %v", err) + } + if result.Saved || result.Source != credentials.SourceEnvironment { + t.Fatalf("response = %+v", result) + } +} + +func TestProviderTokenEndpointsRejectRemoteAndCrossOriginRequests(t *testing.T) { + setupProviderTokenTest(t) + mux := http.NewServeMux() + registerProviderTokenHandlers(mux) + + remote := httptest.NewRequest(http.MethodPost, "/api/captain/ai/providers/openai/token/test", strings.NewReader(`{}`)) + remote.Header.Set("Content-Type", "application/json") + remote.RemoteAddr = "203.0.113.10:1234" + remoteResponse := httptest.NewRecorder() + mux.ServeHTTP(remoteResponse, remote) + if remoteResponse.Code != http.StatusForbidden { + t.Fatalf("remote status=%d", remoteResponse.Code) + } + + crossOrigin := httptest.NewRequest(http.MethodPost, "/api/captain/ai/providers/openai/token/test", strings.NewReader(`{}`)) + crossOrigin.Header.Set("Content-Type", "application/json") + crossOrigin.Header.Set("Origin", "http://example.invalid") + crossOrigin.Host = "localhost:9020" + crossOrigin.RemoteAddr = "127.0.0.1:1234" + crossOriginResponse := httptest.NewRecorder() + mux.ServeHTTP(crossOriginResponse, crossOrigin) + if crossOriginResponse.Code != http.StatusForbidden { + t.Fatalf("cross-origin status=%d", crossOriginResponse.Code) + } + + rebound := httptest.NewRequest(http.MethodPost, "/api/captain/ai/providers/openai/token/test", strings.NewReader(`{}`)) + rebound.Header.Set("Content-Type", "application/json") + rebound.Header.Set("Origin", "http://captain.example") + rebound.Host = "captain.example" + rebound.RemoteAddr = "127.0.0.1:1234" + reboundResponse := httptest.NewRecorder() + mux.ServeHTTP(reboundResponse, rebound) + if reboundResponse.Code != http.StatusForbidden { + t.Fatalf("non-loopback host status=%d", reboundResponse.Code) + } +} + +func TestProviderTokenOpenAPIPaths(t *testing.T) { + spec := &rpc.OpenAPISpec{} + addCaptainProviderTokenPaths(spec) + if _, ok := spec.Paths["/api/captain/ai/providers/{provider}/token"]; !ok { + t.Fatal("missing token save OpenAPI path") + } + if _, ok := spec.Paths["/api/captain/ai/providers/{provider}/token/test"]; !ok { + t.Fatal("missing token test OpenAPI path") + } +} + +func setupProviderTokenTest(t *testing.T) { + t.Helper() + credentials.SetPathForTesting(filepath.Join(t.TempDir(), "vault")) + t.Cleanup(func() { credentials.SetPathForTesting("") }) + for _, name := range []string{"ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY", "GOOGLE_API_KEY", "DEEPSEEK_API_KEY"} { + t.Setenv(name, "") + } +} + +func serveProviderTokenRequest(t *testing.T, method, path, body string) *httptest.ResponseRecorder { + t.Helper() + mux := http.NewServeMux() + registerProviderTokenHandlers(mux) + req := httptest.NewRequest(method, path, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Origin", "http://localhost:9020") + req.Host = "localhost:9020" + req.RemoteAddr = "127.0.0.1:1234" + response := httptest.NewRecorder() + mux.ServeHTTP(response, req) + return response +} diff --git a/pkg/cli/serve_sessions.go b/pkg/cli/serve_sessions.go new file mode 100644 index 0000000..45e7d1d --- /dev/null +++ b/pkg/cli/serve_sessions.go @@ -0,0 +1,121 @@ +package cli + +import ( + "context" + "fmt" + "net/http" + "strconv" + "strings" +) + +func handleSessionsLive() http.HandlerFunc { + return handleSessionsLiveWithRunner(RunSessionLive) +} + +func handleSessionsLiveWithRunner(run func(context.Context, SessionLiveOptions) (SessionLiveResult, error)) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + query := r.URL.Query() + opts := SessionLiveOptions{ + Source: strings.TrimSpace(query.Get("source")), Project: strings.TrimSpace(query.Get("project")), + Query: strings.TrimSpace(query.Get("q")), Limit: 100, Cursor: strings.TrimSpace(query.Get("cursor")), + } + if opts.Source == "" { + opts.Source = "all" + } + if raw := strings.TrimSpace(query.Get("all")); raw != "" { + opts.All, _ = strconv.ParseBool(raw) + } + if strings.EqualFold(opts.Project, "all") { + opts.Project, opts.All = "", true + } else if opts.Project != "" { + opts.All = false + } + if raw := strings.TrimSpace(query.Get("limit")); raw != "" { + limit, err := strconv.Atoi(raw) + if err != nil { + http.Error(w, fmt.Sprintf("invalid limit %q", raw), http.StatusBadRequest) + return + } + opts.Limit = limit + } + result, err := run(r.Context(), opts) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + writeServeJSON(w, http.StatusOK, result) + } +} + +func handleSessionsThroughput() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + query := r.URL.Query() + opts := SessionThroughputOptions{ + Source: strings.TrimSpace(query.Get("source")), Project: strings.TrimSpace(query.Get("project")), + Query: strings.TrimSpace(query.Get("q")), Limit: defaultSessionThroughputLimit, + } + if opts.Source == "" { + opts.Source = "all" + } + if raw := strings.TrimSpace(query.Get("all")); raw != "" { + opts.All, _ = strconv.ParseBool(raw) + } + if strings.EqualFold(opts.Project, "all") { + opts.Project, opts.All = "", true + } else if opts.Project != "" { + opts.All = false + } + if raw := strings.TrimSpace(query.Get("limit")); raw != "" { + limit, err := strconv.Atoi(raw) + if err != nil { + http.Error(w, fmt.Sprintf("invalid limit %q", raw), http.StatusBadRequest) + return + } + opts.Limit = limit + } + result, err := RunSessionThroughput(r.Context(), opts) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + writeServeJSON(w, http.StatusOK, result) + } +} + +func handleProjects() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + result, err := RunProjectOptions(r.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + writeServeJSON(w, http.StatusOK, result) + } +} + +func handleSessionGet() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + id := strings.TrimSpace(r.PathValue("id")) + if id == "" { + http.Error(w, "session id is required", http.StatusBadRequest) + return + } + query := r.URL.Query() + s, err := RunSessionGet(r.Context(), SessionGetOptions{ + ID: id, Offset: queryInt(query.Get("offset")), Limit: queryInt(query.Get("limit")), Tail: queryInt(query.Get("tail")), + }) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + writeServeJSON(w, http.StatusOK, s) + } +} + +func queryInt(value string) int { + n, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil || n < 0 { + return 0 + } + return n +} diff --git a/pkg/cli/serve_test.go b/pkg/cli/serve_test.go index 139c165..1a3258c 100644 --- a/pkg/cli/serve_test.go +++ b/pkg/cli/serve_test.go @@ -10,7 +10,7 @@ import ( "testing" "github.com/flanksource/captain/pkg/claude" - "github.com/flanksource/captain/pkg/session" + "github.com/flanksource/captain/pkg/database" ) func TestHandleThreadFromAgentCreatesThread(t *testing.T) { @@ -62,9 +62,10 @@ func TestHandleThreadFromAgentRequiresProviderSession(t *testing.T) { } } -func TestHandleSessionGetReturnsUnifiedModel(t *testing.T) { +func TestHandleSessionGetReturnsAllMatches(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + db := withTestCaptainDB(t) project := filepath.Join(home, "work", "project") if err := os.MkdirAll(project, 0o755); err != nil { t.Fatal(err) @@ -72,16 +73,26 @@ func TestHandleSessionGetReturnsUnifiedModel(t *testing.T) { t.Chdir(project) markProjectRoot(t, project) - writeJSONL(t, filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-web.jsonl"), + sessionFile := filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-web.jsonl") + writeJSONL(t, sessionFile, map[string]any{ "type": "assistant", "sessionId": "sess-web", "uuid": "a1", "timestamp": "2026-07-06T10:00:00Z", "cwd": project, "message": map[string]any{ - "role": "assistant", "model": "claude-opus-4", + "role": "assistant", "model": "claude-opus-4-5", "content": []any{map[string]any{"type": "text", "text": "hello from the model"}}, }, }, ) + for _, input := range []database.CreateSessionInput{ + {ProviderSessionID: "sess-web", Source: "claude", HostID: "test-host", Path: sessionFile, Project: "example", CWD: project}, + {ProviderSessionID: "sess-web", Source: "gavel", Provider: "cmux", HostID: "local"}, + {ProviderSessionID: "sess-web", Source: "claude", Provider: "cmux", HostID: "local"}, + } { + if _, err := db.CreateOrGetSession(t.Context(), input); err != nil { + t.Fatalf("create duplicate session: %v", err) + } + } req := httptest.NewRequest(http.MethodGet, "/api/captain/sessions/sess-web", nil) req.SetPathValue("id", "sess-web") @@ -91,14 +102,56 @@ func TestHandleSessionGetReturnsUnifiedModel(t *testing.T) { if rec.Code != http.StatusOK { t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String()) } - var got session.Session + var got SessionGetResult if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { t.Fatalf("unmarshal: %v (body=%s)", err, rec.Body.String()) } - if got.ID != "sess-web" || got.Source != "claude" { - t.Fatalf("session = %+v", got) + if got.Total != 3 || len(got.Sessions) != 3 { + t.Fatalf("result = %+v", got) + } + detail := got.Sessions[0].Detail + if detail == nil || detail.ID != "sess-web" || detail.Source != "claude" { + t.Fatalf("session = %+v", detail) + } + if len(detail.Messages) == 0 || detail.Messages[0].Parts[0].Text != "hello from the model" { + t.Fatalf("messages = %+v", detail.Messages) + } +} + +func TestHandleProjectsReturnsOptions(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + withTestCaptainDB(t) + project := filepath.Join(home, "work", "project") + if err := os.MkdirAll(project, 0o755); err != nil { + t.Fatal(err) + } + markProjectRoot(t, project) + writeJSONL(t, filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-web.jsonl"), + map[string]any{ + "type": "assistant", + "sessionId": "sess-web", + "timestamp": "2026-07-06T10:00:00Z", + "cwd": project, + "message": map[string]any{ + "role": "assistant", + "content": []any{map[string]any{"type": "text", "text": "hello"}}, + }, + }, + ) + + req := httptest.NewRequest(http.MethodGet, "/api/captain/projects", nil) + rec := httptest.NewRecorder() + handleProjects()(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String()) + } + var got ProjectOptionsResult + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("unmarshal: %v (body=%s)", err, rec.Body.String()) } - if len(got.Messages) == 0 || got.Messages[0].Parts[0].Text != "hello from the model" { - t.Fatalf("messages = %+v", got.Messages) + if got.Total != 1 || got.Projects[0].Value != project { + t.Fatalf("projects = %+v, want %q", got, project) } } diff --git a/pkg/cli/session_cache.go b/pkg/cli/session_cache.go deleted file mode 100644 index f448663..0000000 --- a/pkg/cli/session_cache.go +++ /dev/null @@ -1,106 +0,0 @@ -package cli - -import ( - "context" - "fmt" - "os" - "runtime" - "sync" - - "github.com/flanksource/captain/pkg/session" - rpchttp "github.com/flanksource/clicky/rpc/http" -) - -// sessionFileRef identifies a session file and its source for summarization. -type sessionFileRef struct { - source string - path string -} - -// summarizeSessionFileCached returns the persisted summary for a session file, -// reusing the stored row when the file's mtime+size are unchanged and otherwise -// rebuilding the rich rows (root + sub-agents) and upserting them. When the -// store is unavailable it degrades to building the summary uncached. -func summarizeSessionFileCached(ref sessionFileRef) (SessionRecord, error) { - info, err := os.Stat(ref.path) - if err != nil { - return SessionRecord{}, err - } - if info.IsDir() { - return SessionRecord{}, fmt.Errorf("%s is a directory", ref.path) - } - modUnix, size := info.ModTime().UnixNano(), info.Size() - - st := sessionStore() - if st != nil { - if row, ok := st.lookupFresh(ref.path, modUnix, size); ok { - return row.toRecord(), nil - } - } - - rows, err := session.RowsFromFile(ref.path, ref.source) - if err != nil { - return SessionRecord{}, err - } - if st != nil { - st.upsertRows(rows) - } - for _, r := range rows { - if r.Path == ref.path { - return recordFromRow(r), nil - } - } - if len(rows) > 0 { - return recordFromRow(rows[0]), nil - } - return SessionRecord{}, fmt.Errorf("no summary for %s", ref.path) -} - -// summarizeSessionRefs sampled-summarizes files concurrently (bounded by -// GOMAXPROCS) through the summary cache, preserving input order. Files that -// fail to summarize or yield no session id are dropped. It records the "parse" -// phase for the wall time of the whole batch. -func summarizeSessionRefs(ctx context.Context, refs []sessionFileRef) []sessionCandidate { - defer rpchttp.Track(ctx, "parse")() - - records := make([]SessionRecord, len(refs)) - found := make([]bool, len(refs)) - - workers := runtime.GOMAXPROCS(0) - if workers > len(refs) { - workers = len(refs) - } - if workers < 1 { - workers = 1 - } - - jobs := make(chan int) - var wg sync.WaitGroup - for w := 0; w < workers; w++ { - wg.Add(1) - go func() { - defer wg.Done() - for i := range jobs { - record, err := summarizeSessionFileCached(refs[i]) - if err != nil || record.ID == "" { - continue - } - records[i] = record - found[i] = true - } - }() - } - for i := range refs { - jobs <- i - } - close(jobs) - wg.Wait() - - candidates := make([]sessionCandidate, 0, len(refs)) - for i := range refs { - if found[i] { - candidates = append(candidates, sessionCandidate{record: records[i], path: refs[i].path}) - } - } - return candidates -} diff --git a/pkg/cli/session_cache_test.go b/pkg/cli/session_cache_test.go deleted file mode 100644 index 420abcd..0000000 --- a/pkg/cli/session_cache_test.go +++ /dev/null @@ -1,83 +0,0 @@ -package cli - -import ( - "context" - "fmt" - "path/filepath" - "testing" -) - -// claudeSummaryRow is a minimal on-disk claude session line carrying the given -// session id, enough for the fast summarizer to extract an ID. -func claudeSummaryRow(id string) map[string]any { - return map[string]any{ - "type": "user", - "sessionId": id, - "timestamp": "2026-06-01T10:00:00Z", - "message": map[string]any{ - "role": "user", - "content": []any{map[string]any{"type": "text", "text": "hi"}}, - }, - } -} - -// TestSummarizeSessionFileCachedDegradesUncached checks the store-disabled path -// (TestMain sets CAPTAIN_SESSION_DB_URL=off): summarization still works, uses the -// in-file session id, and reflects the current content on every call (no stale -// cache). Cache reuse/invalidation against a real DB is covered by the gated -// integration test. -func TestSummarizeSessionFileCachedDegradesUncached(t *testing.T) { - path := filepath.Join(t.TempDir(), "aa.jsonl") - ref := sessionFileRef{source: "claude", path: path} - - writeJSONL(t, path, claudeSummaryRow("aa")) - first, err := summarizeSessionFileCached(ref) - if err != nil { - t.Fatalf("first summarize: %v", err) - } - if first.ID != "aa" { - t.Fatalf("first.ID = %q, want aa (from the in-file sessionId)", first.ID) - } - - // With no store, a rewritten file is re-read every time (uncached). - writeJSONL(t, path, claudeSummaryRow("bb")) - again, err := summarizeSessionFileCached(ref) - if err != nil { - t.Fatalf("second summarize: %v", err) - } - if again.ID != "bb" { - t.Fatalf("second.ID = %q, want bb (uncached re-read)", again.ID) - } -} - -func TestSummarizeSessionRefsParallelFindsEverySession(t *testing.T) { - dir := t.TempDir() - const count = 60 - refs := make([]sessionFileRef, 0, count) - want := make(map[string]bool, count) - for i := 0; i < count; i++ { - id := fmt.Sprintf("sess-%03d", i) - path := filepath.Join(dir, id+".jsonl") - writeJSONL(t, path, claudeSummaryRow(id)) - refs = append(refs, sessionFileRef{source: "claude", path: path}) - want[id] = true - } - - got := summarizeSessionRefs(context.Background(), refs) - if len(got) != count { - t.Fatalf("got %d candidates, want %d", len(got), count) - } - seen := make(map[string]bool, count) - for i, candidate := range got { - seen[candidate.record.ID] = true - // Order must match the input refs so downstream sorting is deterministic. - if candidate.path != refs[i].path { - t.Fatalf("candidate[%d].path = %q, want %q (order not preserved)", i, candidate.path, refs[i].path) - } - } - for id := range want { - if !seen[id] { - t.Fatalf("session %q missing from parallel summary", id) - } - } -} diff --git a/pkg/cli/session_env.go b/pkg/cli/session_env.go new file mode 100644 index 0000000..497c0bb --- /dev/null +++ b/pkg/cli/session_env.go @@ -0,0 +1,102 @@ +package cli + +import ( + "bytes" + "os" + "os/exec" + "regexp" + "runtime" + "strconv" + "strings" +) + +// cmuxEnvPattern matches a CMUX_* environment assignment in the flattened +// `ps -Eww` output on macOS. CMUX values are single whitespace-free tokens +// (ids, ports, paths), so this reliably skips decoy references inside the +// argv (e.g. `${CMUX_CLAUDE_HOOK_CMUX_BIN:-cmux}`, which uses ':' not '='). +var cmuxEnvPattern = regexp.MustCompile(`\bCMUX_[A-Z0-9_]+=\S+`) + +// processSurface returns the cmux surface hosting the process, or nil when the +// process has no CMUX_* environment (not launched under cmux) or its environment +// cannot be read. Enrichment is best-effort and never fails the caller. +func processSurface(pid int) *CmuxSurface { + if pid <= 0 { + return nil + } + env, err := processEnviron(pid) + if err != nil { + return nil + } + return parseCmuxSurface(env) +} + +// processEnviron reads another process's environment. On Linux it reads the +// NUL-separated /proc//environ (full environment). On macOS/other it shells +// out to `ps -Eww` and extracts only the CMUX_* assignments (the full argv+env +// blob cannot be split reliably, but the CMUX subset is unambiguous). +func processEnviron(pid int) (map[string]string, error) { + if runtime.GOOS == "linux" { + data, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/environ") + if err != nil { + return nil, err + } + return parseProcEnviron(data), nil + } + out, err := exec.Command("ps", "-Eww", "-o", "command=", "-p", strconv.Itoa(pid)).Output() + if err != nil { + return nil, err + } + return extractCmuxEnv(string(out)), nil +} + +// parseProcEnviron parses NUL-separated KEY=VALUE pairs from /proc//environ. +func parseProcEnviron(data []byte) map[string]string { + env := make(map[string]string) + for _, pair := range bytes.Split(data, []byte{0}) { + if len(pair) == 0 { + continue + } + key, value, ok := strings.Cut(string(pair), "=") + if !ok || key == "" { + continue + } + env[key] = value + } + return env +} + +// extractCmuxEnv pulls CMUX_* assignments out of a flattened `ps -Eww` line. +func extractCmuxEnv(psOutput string) map[string]string { + env := make(map[string]string) + for _, match := range cmuxEnvPattern.FindAllString(psOutput, -1) { + key, value, ok := strings.Cut(match, "=") + if !ok { + continue + } + env[key] = value + } + return env +} + +// parseCmuxSurface builds a CmuxSurface from an environment map, returning nil +// when no CMUX surface markers are present. +func parseCmuxSurface(env map[string]string) *CmuxSurface { + surface := CmuxSurface{ + SurfaceID: env["CMUX_SURFACE_ID"], + WorkspaceID: env["CMUX_WORKSPACE_ID"], + TabID: env["CMUX_TAB_ID"], + PanelID: env["CMUX_PANEL_ID"], + AgentKind: env["CMUX_AGENT_LAUNCH_KIND"], + SocketPath: env["CMUX_SOCKET_PATH"], + } + if port, err := strconv.Atoi(env["CMUX_PORT"]); err == nil { + surface.Port = port + } + if pid, err := strconv.Atoi(env["CMUX_CLAUDE_PID"]); err == nil { + surface.ClaudePID = pid + } + if surface == (CmuxSurface{}) { + return nil + } + return &surface +} diff --git a/pkg/cli/session_env_test.go b/pkg/cli/session_env_test.go new file mode 100644 index 0000000..a38da8c --- /dev/null +++ b/pkg/cli/session_env_test.go @@ -0,0 +1,72 @@ +package cli + +import "testing" + +func TestParseProcEnviron(t *testing.T) { + data := []byte("PATH=/usr/bin\x00CMUX_SURFACE_ID=SURFACE-1\x00CMUX_PORT=9150\x00EMPTY\x00") + env := parseProcEnviron(data) + if env["PATH"] != "/usr/bin" { + t.Fatalf("PATH = %q", env["PATH"]) + } + if env["CMUX_SURFACE_ID"] != "SURFACE-1" { + t.Fatalf("CMUX_SURFACE_ID = %q", env["CMUX_SURFACE_ID"]) + } + if _, ok := env["EMPTY"]; ok { + t.Fatal("entry without '=' should be skipped") + } +} + +func TestExtractCmuxEnvIgnoresArgvDecoys(t *testing.T) { + // Flattened `ps -Eww` output: the claude argv embeds ${CMUX_..:-cmux} + // references (colon, not '='), which must NOT be captured as env vars. + psOutput := `claude --settings {"hooks":[{"command":"${CMUX_CLAUDE_HOOK_CMUX_BIN:-cmux} hooks"}]} ` + + `CMUX_SURFACE_ID=1D727ED7-0DCB CMUX_WORKSPACE_ID=4D2B7ACC CMUX_TAB_ID=4D2B7ACC ` + + `CMUX_PANEL_ID=1D727ED7 CMUX_PORT=9150 CMUX_AGENT_LAUNCH_KIND=claude ` + + `CMUX_SOCKET_PATH=/Users/moshe/.local/state/cmux/cmux.sock ` + + `CMUX_AGENT_LAUNCH_ARGV_B64=L1VzZXJzL21vc2hlAA==` + env := extractCmuxEnv(psOutput) + + if env["CMUX_SURFACE_ID"] != "1D727ED7-0DCB" { + t.Fatalf("CMUX_SURFACE_ID = %q", env["CMUX_SURFACE_ID"]) + } + if _, ok := env["CMUX_CLAUDE_HOOK_CMUX_BIN"]; ok { + t.Fatal("decoy ${CMUX_...:-cmux} reference should not be captured") + } + if env["CMUX_AGENT_LAUNCH_ARGV_B64"] != "L1VzZXJzL21vc2hlAA==" { + t.Fatalf("base64 value with '==' padding not captured: %q", env["CMUX_AGENT_LAUNCH_ARGV_B64"]) + } +} + +func TestParseCmuxSurface(t *testing.T) { + env := map[string]string{ + "CMUX_SURFACE_ID": "SURFACE-1", + "CMUX_WORKSPACE_ID": "WS-1", + "CMUX_TAB_ID": "TAB-1", + "CMUX_PANEL_ID": "PANEL-1", + "CMUX_PORT": "9150", + "CMUX_CLAUDE_PID": "33088", + "CMUX_AGENT_LAUNCH_KIND": "claude", + "CMUX_SOCKET_PATH": "/tmp/cmux.sock", + } + surface := parseCmuxSurface(env) + if surface == nil { + t.Fatal("expected surface") + } + if surface.SurfaceID != "SURFACE-1" || surface.WorkspaceID != "WS-1" || + surface.TabID != "TAB-1" || surface.PanelID != "PANEL-1" || + surface.AgentKind != "claude" || surface.SocketPath != "/tmp/cmux.sock" { + t.Fatalf("surface = %+v", surface) + } + if surface.Port != 9150 { + t.Fatalf("port = %d", surface.Port) + } + if surface.ClaudePID != 33088 { + t.Fatalf("claude pid = %d", surface.ClaudePID) + } +} + +func TestParseCmuxSurfaceNilWhenAbsent(t *testing.T) { + if surface := parseCmuxSurface(map[string]string{"PATH": "/usr/bin"}); surface != nil { + t.Fatalf("expected nil surface, got %+v", surface) + } +} diff --git a/pkg/cli/session_get.go b/pkg/cli/session_get.go new file mode 100644 index 0000000..53c6b36 --- /dev/null +++ b/pkg/cli/session_get.go @@ -0,0 +1,440 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/flanksource/clicky" + clickyapi "github.com/flanksource/clicky/api" + rpchttp "github.com/flanksource/clicky/rpc/http" +) + +type SessionGetResult struct { + RootSessionID string `json:"rootSessionId,omitempty"` + Sessions []SessionGetItem `json:"sessions"` + Total int `json:"total"` +} + +type SessionGetItem struct { + CaptainID string `json:"captainId"` + ParentSessionID string `json:"parentSessionId,omitempty"` + RootSessionID string `json:"rootSessionId,omitempty"` + ProviderSessionID string `json:"providerSessionId,omitempty"` + Host string `json:"host,omitempty"` + Aggregate bool `json:"aggregate,omitempty"` + DetailAvailable bool `json:"detailAvailable"` + Summary SessionRecord `json:"summary"` + Detail *session.Session `json:"detail,omitempty"` + ActiveRunID string `json:"activeRunId,omitempty"` + Chat *ChatCapabilities `json:"chat,omitempty"` + ChatState *ChatStateFrame `json:"chatState,omitempty"` +} + +type sessionGetStore interface { + sessionOverviewStore + ListPromptRuns(context.Context, database.PromptRunFilter) ([]database.PromptRun, error) +} + +// RunSessionGet returns every Captain session matching an exact Captain UUID +// or provider-session-id prefix. Transcript-less matches remain visible via +// their overview metadata; recorded transcripts and native prompt runs are +// projected into the same detail model and paged. +func RunSessionGet(ctx context.Context, opts SessionGetOptions) (SessionGetResult, error) { + if strings.TrimSpace(opts.ID) == "" { + return SessionGetResult{}, fmt.Errorf("id is required") + } + stopDatabase := rpchttp.Track(ctx, "database") + db, err := captainDB(ctx) + stopDatabase() + if err != nil { + return SessionGetResult{}, err + } + return runSessionGet(ctx, db, opts) +} + +func runSessionGet(ctx context.Context, db sessionGetStore, opts SessionGetOptions) (SessionGetResult, error) { + id := strings.TrimSpace(opts.ID) + if id == "" { + return SessionGetResult{}, fmt.Errorf("id is required") + } + stopLookup := rpchttp.Track(ctx, "lookup") + overviews, err := resolveOverviewsByIdentity(ctx, db, id) + stopLookup() + if err != nil { + return SessionGetResult{}, err + } + + items := make([]SessionGetItem, 0, len(overviews)) + for i := range overviews { + stopHydrate := rpchttp.Track(ctx, "hydrate") + item, itemErr := buildSessionGetItem(ctx, db, overviews[i], opts) + stopHydrate() + if itemErr != nil { + return SessionGetResult{}, itemErr + } + items = append(items, item) + } + rootID := "" + if len(items) > 1 && items[0].ParentSessionID == "" { + rootID = items[0].CaptainID + } + return SessionGetResult{RootSessionID: rootID, Sessions: items, Total: len(items)}, nil +} + +func buildSessionGetItem(ctx context.Context, db sessionGetStore, overview database.SessionOverview, opts SessionGetOptions) (SessionGetItem, error) { + item := SessionGetItem{ + CaptainID: overview.ID.String(), ProviderSessionID: stringOr(overview.ProviderSessionID, ""), + Host: overview.HostID, Aggregate: stringOr(overview.AgentType, "") == "batch", + Summary: recordFromOverview(overview), + } + if overview.ParentSessionID != nil { + item.ParentSessionID = overview.ParentSessionID.String() + } + if overview.RootSessionID != nil { + item.RootSessionID = overview.RootSessionID.String() + } + capabilities := sessionChatCapabilities(item.Summary) + item.Chat = &capabilities + active, ok := promptChats.getRun(item.CaptainID) + if !ok && item.ProviderSessionID != "" { + active, ok = promptChats.getSession(item.ProviderSessionID) + } + if ok { + var activeCapabilities ChatCapabilities + item.ActiveRunID, activeCapabilities, item.ChatState = active.projection() + item.Chat = &activeCapabilities + } + detail, err := loadSessionDetail(ctx, db, overview) + if err != nil { + return SessionGetItem{}, err + } + if detail == nil { + return item, nil + } + enrichSessionDetail(detail, item.Summary) + item.DetailAvailable = true + item.Summary.DetailAvailable = true + item.Summary.Messages = max(item.Summary.Messages, len(detail.Messages)) + item.Summary.Provider = firstNonEmpty(item.Summary.Provider, detail.Provider) + item.Summary.Backend = firstNonEmpty(item.Summary.Backend, detail.Backend) + item.Summary.Model = firstNonEmpty(item.Summary.Model, detail.Model) + item.Summary.ReasoningEffort = firstNonEmpty(item.Summary.ReasoningEffort, detail.ReasoningEffort) + pageSessionTranscript(detail, opts) + item.Detail = detail + return item, nil +} + +func loadSessionDetail(ctx context.Context, db sessionGetStore, overview database.SessionOverview) (*session.Session, error) { + path := stringOr(overview.HistoryFile, stringOr(overview.Path, "")) + var detail *session.Session + if path != "" { + stopParse := rpchttp.Track(ctx, "parse") + parsed, err := buildSessionModel(candidateFromOverview(overview)) + stopParse() + if err != nil { + return nil, fmt.Errorf("parse Captain session %s: %w", overview.ID, err) + } + detail = parsed + } + stopPromptRuns := rpchttp.Track(ctx, "prompt_runs") + runs, err := db.ListPromptRuns(ctx, database.PromptRunFilter{SessionID: &overview.ID}) + stopPromptRuns() + if err != nil { + return nil, fmt.Errorf("list prompt runs for Captain session %s: %w", overview.ID, err) + } + if len(runs) == 0 { + return detail, nil + } + if detail == nil { + return sessionFromPromptRun(overview, runs[0]) + } + if err := attachPromptRunData(detail, runs[0]); err != nil { + return nil, fmt.Errorf("attach prompt run %s to Captain session %s: %w", runs[0].ID, overview.ID, err) + } + return detail, nil +} + +func sessionFromPromptRun(overview database.SessionOverview, run database.PromptRun) (*session.Session, error) { + resolved := run.Runtime.Resolved + requested := run.Runtime.Requested + detail := &session.Session{ + ID: stringOr(overview.ProviderSessionID, overview.ID.String()), + Source: overview.Source, + Project: stringOr(overview.Project, ""), + CWD: stringOr(overview.CWD, ""), + Slug: stringOr(overview.Slug, ""), + Title: stringOr(overview.Title, ""), + InitialPrompt: stringOr(overview.InitialPrompt, run.PromptMarkdown), + Version: stringOr(overview.CLIVersion, ""), + Provider: firstNonEmpty(overview.Provider, resolved.Provider, requested.Provider), + Backend: firstNonEmpty(stringOr(overview.Backend, ""), resolved.Backend, requested.Backend), + Model: firstNonEmpty(stringOr(overview.Model, ""), resolved.Model, requested.Model), + ReasoningEffort: firstNonEmpty(stringOr(overview.Effort, ""), resolved.Effort, requested.Effort), + StartedAt: firstTime(overview.StartedAt, run.StartedAt, &run.QueuedAt), + EndedAt: firstTime(overview.EndedAt, run.FinishedAt), + } + if run.PromptMarkdown != "" { + detail.Messages = append(detail.Messages, promptRunMessage(run, "user", run.PromptMarkdown)) + } + resultText, err := promptRunResultText(run) + if err != nil { + return nil, err + } + if resultText != "" { + detail.Messages = append(detail.Messages, promptRunMessage(run, "assistant", resultText)) + } + if run.Error != "" { + detail.Events = append(detail.Events, session.Event{ + Type: "error", Scope: "prompt_run", Timestamp: run.FinishedAt, UUID: run.ID.String(), + Data: map[string]any{"message": run.Error, "state": run.State}, + }) + } + if err := attachPromptRunData(detail, run); err != nil { + return nil, fmt.Errorf("encode prompt run %s: %w", run.ID, err) + } + return detail, nil +} + +func promptRunMessage(run database.PromptRun, role, text string) session.Message { + return session.Message{ + ID: run.ID.String() + "-" + role, Role: role, + Parts: []session.Part{{Type: session.PartText, Text: text}}, + } +} + +func promptRunResultText(run database.PromptRun) (string, error) { + if run.ResultText != "" { + return run.ResultText, nil + } + if len(run.ResultJSON) == 0 { + return "", nil + } + raw, err := json.Marshal(run.ResultJSON) + if err != nil { + return "", fmt.Errorf("encode prompt run %s result: %w", run.ID, err) + } + return string(raw), nil +} + +func attachPromptRunData(detail *session.Session, run database.PromptRun) error { + if len(run.RenderedSpec) > 0 { + raw, err := json.Marshal(run.RenderedSpec) + if err != nil { + return err + } + detail.Prompt = raw + } + detail.StructuredOutput = promptRunStructuredOutput(run) + return nil +} + +func promptRunStructuredOutput(run database.PromptRun) map[string]any { + if run.ResultJSON != nil { + return run.ResultJSON + } + if !promptRunDeclaresOutputSchema(run.RenderedSpec) || !json.Valid([]byte(run.ResultText)) { + return nil + } + var output map[string]any + if err := json.Unmarshal([]byte(run.ResultText), &output); err != nil { + return nil + } + return output +} + +func promptRunDeclaresOutputSchema(rendered map[string]any) bool { + schema, ok := rendered["outputSchema"].(map[string]any) + return ok && len(schema) > 0 +} + +func firstTime(values ...*time.Time) *time.Time { + for _, value := range values { + if value != nil && !value.IsZero() { + return value + } + } + return nil +} + +func sessionChatCapabilities(summary SessionRecord) ChatCapabilities { + capabilities := chatCapabilitiesForBackend(summary.Backend) + if summary.Source == "claude" || summary.Source == "codex" { + capabilities.Resume = true + } + return capabilities +} + +func enrichSessionDetail(detail *session.Session, summary SessionRecord) { + if detail.Provider == "" { + detail.Provider = summary.Provider + } + if detail.Model == "" { + detail.Model = summary.Model + } + if detail.Backend == "" { + detail.Backend = summary.Backend + } + if detail.ReasoningEffort == "" { + detail.ReasoningEffort = summary.ReasoningEffort + } + if summary.Live != nil { + detail.Live = &session.LiveProcess{ + PID: summary.Live.PID, Status: summary.Live.Status, Active: summary.Live.Active, + CPUPercent: summary.Live.CPUPercent, MemoryPercent: summary.Live.MemoryPercent, + StartedAt: summary.Live.StartedAt, CWD: summary.Live.CWD, Command: summary.Live.Command, + } + } + for i := range detail.Turns { + if detail.Turns[i].Backend == "" { + detail.Turns[i].Backend = summary.Backend + } + if detail.Turns[i].ReasoningEffort == "" { + detail.Turns[i].ReasoningEffort = summary.ReasoningEffort + } + } +} + +// Pretty renders the flat list form. NOTE: terminal and HTML output do not go +// through here — clicky's TryTypedValue matches TreeMixin before Pretty, so +// Tree below wins for any formatter-driven render. Both delegate to +// SessionGetItem.Pretty, which is where per-session layout changes belong. +func (r SessionGetResult) Pretty() clickyapi.Text { + list := clicky.List() + list.Unstyled = true + list.MaxInline = 1 + for i := range r.Sessions { + list.Items = append(list.Items, sessionGetListItem{text: r.Sessions[i].Pretty()}) + } + return clickyapi.Text{}.Add(list) +} + +func (r SessionGetResult) Tree() clickyapi.TreeNode { + byParent := map[string][]SessionGetItem{} + for i := range r.Sessions { + byParent[r.Sessions[i].ParentSessionID] = append(byParent[r.Sessions[i].ParentSessionID], r.Sessions[i]) + } + children := make([]clickyapi.TreeNode, 0, len(byParent[""])) + for _, item := range byParent[""] { + children = append(children, sessionGetTreeNode{item: item, byParent: byParent}) + } + return &clickyapi.ConcreteBranchNode{Children: children} +} + +type sessionGetTreeNode struct { + item SessionGetItem + byParent map[string][]SessionGetItem +} + +func (n sessionGetTreeNode) Pretty() clickyapi.Text { return n.item.Pretty() } + +func (n sessionGetTreeNode) GetChildren() []clickyapi.TreeNode { + items := n.byParent[n.item.CaptainID] + children := make([]clickyapi.TreeNode, len(items)) + for i := range items { + children[i] = sessionGetTreeNode{item: items[i], byParent: n.byParent} + } + return children +} + +func (i SessionGetItem) Pretty() clickyapi.Text { + text := clickyapi.Text{}. + Append("Captain ", "text-muted"). + Append(i.CaptainID, "font-bold text-blue-600") + if strings.TrimSpace(i.Host) != "" { + text = text.Append(" "+i.Host, "text-muted") + } + // The detail body opens with its own header and Summary rows carrying + // source, project, cwd and the provider session id, so repeating them here + // would duplicate four of the first eight lines of output. + if i.Detail != nil { + return text.NewLine().NewLine().Add(i.Detail.Pretty()).Add(i.hiddenRowsNotice()) + } + if i.Summary.Source != "" { + text = text.Append(" ", "").Append(strings.ToUpper(i.Summary.Source), "text-muted") + } + for _, metadata := range []struct{ label, value string }{ + {label: "Provider session", value: i.ProviderSessionID}, + {label: "Project", value: i.Summary.Project}, + {label: "CWD", value: i.Summary.CWD}, + } { + if strings.TrimSpace(metadata.value) != "" { + text = text.NewLine(). + Append(" "+metadata.label+": ", "text-muted"). + Append(metadata.value, "") + } + } + if i.Aggregate { + return text.NewLine().Append(" Aggregate session; child results are shown below", "text-muted") + } + return text.NewLine().Append(" Transcript: unavailable", "text-amber-600") +} + +// hiddenRowsNotice reports messages dropped by the transcript window so a +// bounded view never reads as the whole session. Summary.Messages holds the +// full count; the detail holds only the retained slice. +func (i SessionGetItem) hiddenRowsNotice() clickyapi.Text { + if i.Detail == nil { + return clickyapi.Text{} + } + hidden := i.Summary.Messages - len(i.Detail.Messages) + if hidden <= 0 { + return clickyapi.Text{} + } + return clickyapi.Text{}.NewLine().Append( + fmt.Sprintf(" … %d of %d messages hidden — use --limit 0 for the full transcript", + hidden, i.Summary.Messages), "text-amber-600") +} + +type sessionGetListItem struct { + text clickyapi.Text +} + +func (i sessionGetListItem) String() string { return i.text.String() + "\n" } +func (i sessionGetListItem) ANSI() string { return i.text.ANSI() } +func (i sessionGetListItem) HTML() string { return i.text.HTML() } +func (i sessionGetListItem) Markdown() string { return i.text.Markdown() } + +// pageSessionTranscript windows both transcript collections: the last Tail +// rows, or an Offset/Limit slice from the start. Events are windowed alongside +// messages because provider state rows (titles, skill listings, prompt +// checkpoints) grow with the session and would otherwise ignore the window +// entirely — a --tail 10 could still emit hundreds of event lines. +func pageSessionTranscript(s *session.Session, opts SessionGetOptions) { + full := session.TranscriptWindow{ + Messages: len(s.Messages), Events: len(s.Events), + ToolCalls: session.CountToolParts(s.Messages), + } + s.Messages = pageTranscriptRows(s.Messages, opts) + s.Events = pageTranscriptRows(s.Events, opts) + // Recorded only when rows were actually dropped, so summaries of a complete + // transcript stay free of window annotations. + if len(s.Messages) != full.Messages || len(s.Events) != full.Events { + s.Window = &full + } +} + +func pageTranscriptRows[T any](rows []T, opts SessionGetOptions) []T { + if opts.Tail > 0 { + if len(rows) > opts.Tail { + return rows[len(rows)-opts.Tail:] + } + return rows + } + if opts.Offset <= 0 && opts.Limit <= 0 { + return rows + } + offset := max(opts.Offset, 0) + if offset >= len(rows) { + return nil + } + rows = rows[offset:] + if opts.Limit > 0 && len(rows) > opts.Limit { + rows = rows[:opts.Limit] + } + return rows +} diff --git a/pkg/cli/session_get_compact_test.go b/pkg/cli/session_get_compact_test.go new file mode 100644 index 0000000..75603e4 --- /dev/null +++ b/pkg/cli/session_get_compact_test.go @@ -0,0 +1,145 @@ +package cli + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/session" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// These specs join the package's existing suite (see session_get_multi_test.go); +// Ginkgo permits only one RunSpecs per package. + +// transcriptFixture builds a session detail with the given number of messages +// and events so paging can be observed on both collections. +func transcriptFixture(messages, events int) *session.Session { + s := &session.Session{ID: "ae95d3f5-9303-45ae-8c93-a83e13c79811", Source: "claude"} + for i := 0; i < messages; i++ { + s.Messages = append(s.Messages, session.Message{ + ID: fmt.Sprintf("m%d", i), + Role: "user", + Parts: []session.Part{{Type: session.PartText, Text: fmt.Sprintf("message %d", i)}}, + }) + } + for i := 0; i < events; i++ { + s.Events = append(s.Events, session.Event{ + Type: "task_started", Scope: "session", UUID: fmt.Sprintf("e%d", i), + }) + } + return s +} + +var _ = Describe("session get transcript paging", func() { + It("bounds events alongside messages under --tail", func() { + detail := transcriptFixture(50, 40) + + pageSessionTranscript(detail, SessionGetOptions{Tail: 5}) + + Expect(detail.Messages).To(HaveLen(5)) + Expect(detail.Events).To(HaveLen(5), + "events were previously unbounded, so --tail 5 still dumped every state row") + Expect(detail.Messages[0].ID).To(Equal("m45")) + Expect(detail.Events[0].UUID).To(Equal("e35")) + }) + + It("bounds events alongside messages under --offset/--limit", func() { + detail := transcriptFixture(50, 40) + + pageSessionTranscript(detail, SessionGetOptions{Offset: 10, Limit: 5}) + + Expect(detail.Messages).To(HaveLen(5)) + Expect(detail.Events).To(HaveLen(5)) + Expect(detail.Messages[0].ID).To(Equal("m10")) + Expect(detail.Events[0].UUID).To(Equal("e10")) + }) + + It("returns everything when the limit is explicitly cleared", func() { + detail := transcriptFixture(50, 40) + + pageSessionTranscript(detail, SessionGetOptions{Limit: 0}) + + Expect(detail.Messages).To(HaveLen(50)) + Expect(detail.Events).To(HaveLen(40)) + }) + + It("leaves a collection untouched when it is shorter than the window", func() { + detail := transcriptFixture(3, 2) + + pageSessionTranscript(detail, SessionGetOptions{Tail: 10}) + + Expect(detail.Messages).To(HaveLen(3)) + Expect(detail.Events).To(HaveLen(2)) + }) +}) + +var _ = Describe("session get header", func() { + It("does not repeat detail fields in the Captain header", func() { + item := SessionGetItem{ + CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", + ProviderSessionID: "ae95d3f5-9303-45ae-8c93-a83e13c79811", + Host: "MacBook-Pro.local", + DetailAvailable: true, + Summary: SessionRecord{ + ID: "ae95d3f5-9303-45ae-8c93-a83e13c79811", Source: "claude", + Project: "captain", CWD: "/repo/captain", Messages: 4, + }, + Detail: &session.Session{ + ID: "ae95d3f5-9303-45ae-8c93-a83e13c79811", Source: "claude", + Project: "captain", CWD: "/repo/captain", + }, + } + + rendered := item.Pretty().String() + + Expect(rendered).To(ContainSubstring("d81f885d-3d60-47c0-8122-a8124f2fbdd1")) + Expect(rendered).To(ContainSubstring("MacBook-Pro.local")) + Expect(strings.Count(rendered, "/repo/captain")).To(Equal(1), + "CWD is a Summary row; the header must not repeat it") + Expect(rendered).NotTo(ContainSubstring("Provider session:"), + "the provider session id is already the detail's ID row") + }) + + It("keeps identifying metadata when no transcript is available", func() { + item := SessionGetItem{ + CaptainID: "7ca78c55-e280-50ff-a19a-9f355a6fc55e", + ProviderSessionID: "ad4c854e-cde6-4b99-99f3-667bf74112e3", + Host: "local", + Summary: SessionRecord{Source: "gavel", Project: "xero-cli", CWD: "/repo/xero"}, + } + + rendered := item.Pretty().String() + + Expect(rendered).To(ContainSubstring("Transcript: unavailable")) + Expect(rendered).To(ContainSubstring("ad4c854e-cde6-4b99-99f3-667bf74112e3")) + Expect(rendered).To(ContainSubstring("xero-cli")) + Expect(rendered).To(ContainSubstring("/repo/xero")) + }) + + It("reports how many transcript rows the window hid", func() { + detail := transcriptFixture(200, 0) + pageSessionTranscript(detail, SessionGetOptions{Tail: 5}) + item := SessionGetItem{ + CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", + Summary: SessionRecord{Messages: 200}, + Detail: detail, + } + + rendered := item.Pretty().String() + + Expect(rendered).To(ContainSubstring("195 of 200 messages hidden")) + Expect(rendered).To(ContainSubstring("--limit 0")) + }) + + It("stays silent when the whole transcript is shown", func() { + detail := transcriptFixture(4, 0) + item := SessionGetItem{ + CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", + Summary: SessionRecord{Messages: 4}, + Detail: detail, + } + + Expect(item.Pretty().String()).NotTo(ContainSubstring("hidden")) + }) +}) diff --git a/pkg/cli/session_get_multi_test.go b/pkg/cli/session_get_multi_test.go new file mode 100644 index 0000000..54e58cc --- /dev/null +++ b/pkg/cli/session_get_multi_test.go @@ -0,0 +1,368 @@ +package cli + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/flanksource/clicky" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + . "github.com/onsi/gomega/gstruct" +) + +func TestSessionGetMulti(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Session Get Multi Suite") +} + +var _ = Describe("session get multi-result output", func() { + It("threads an opaque cursor through the live-session HTTP collection", func() { + var received SessionLiveOptions + handler := handleSessionsLiveWithRunner(func(_ context.Context, opts SessionLiveOptions) (SessionLiveResult, error) { + received = opts + return SessionLiveResult{Total: 3, NextCursor: "cursor-next"}, nil + }) + request := httptest.NewRequest(http.MethodGet, "/api/captain/sessions/live?source=codex&all=true&limit=2&cursor=cursor-current", nil) + response := httptest.NewRecorder() + + handler.ServeHTTP(response, request) + + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(received).To(MatchFields(IgnoreExtras, Fields{ + "Source": Equal("codex"), "All": BeTrue(), "Limit": Equal(2), "Cursor": Equal("cursor-current"), + })) + var result SessionLiveResult + Expect(json.Unmarshal(response.Body.Bytes(), &result)).To(Succeed()) + Expect(result.Total).To(Equal(3)) + Expect(result.NextCursor).To(Equal("cursor-next")) + }) + + It("projects overview runtime data into the unified session detail", func() { + detail := &session.Session{Turns: []session.Turn{{ID: "turn-1", Index: 1}}} + summary := SessionRecord{ + Backend: "codex-cmux", + ReasoningEffort: "high", + Live: &SessionLiveWire{ + PID: 4821, Status: "running", Active: true, CWD: "/repo", Command: "codex", + }, + } + + enrichSessionDetail(detail, summary) + + Expect(detail.Backend).To(Equal("codex-cmux")) + Expect(detail.ReasoningEffort).To(Equal("high")) + Expect(detail.Live).To(Equal(&session.LiveProcess{ + PID: 4821, Status: "running", Active: true, CWD: "/repo", Command: "codex", + })) + Expect(detail.Turns).To(Equal([]session.Turn{{ + ID: "turn-1", Index: 1, Backend: "codex-cmux", ReasoningEffort: "high", + }})) + }) + + It("passes UUID-prefix searches through the bounded list filter", func() { + providerID := "ad4c854e-cde6-4b99-99f3-667bf74112e3" + store := &sessionGetOverviewStore{ + list: []database.SessionListSummary{ + {ID: uuid.New(), ProviderSessionID: &providerID, Source: "claude"}, + {ID: uuid.New(), ProviderSessionID: &providerID, Source: "gavel"}, + }, + } + + page, err := dbSessionRecords(context.Background(), store, sessionRecordQuery{ + Source: "all", Query: "ad4c854e", + }) + Expect(err).NotTo(HaveOccurred()) + Expect(page.Records).To(HaveLen(2)) + Expect(store.listFilter.Query).To(Equal("ad4c854e")) + Expect(store.listFilter.RootsOnly).To(BeTrue()) + }) + + It("loads every bounded page when full live output is explicit", func(ctx SpecContext) { + store := &sessionGetOverviewStore{listPages: map[string]database.SessionListPage{ + "": { + Rows: []database.SessionListSummary{{ID: uuid.New()}, {ID: uuid.New()}}, + Total: 3, NextCursor: "cursor-2", + }, + "cursor-2": {Rows: []database.SessionListSummary{{ID: uuid.New()}}, Total: 3}, + }} + + page, err := dbAllSessionRecords(ctx, store, sessionRecordQuery{Limit: 2}) + + Expect(err).NotTo(HaveOccurred()) + Expect(page.Records).To(HaveLen(3)) + Expect(page.Total).To(Equal(3)) + Expect(store.listFilters).To(HaveLen(2)) + Expect(store.listFilters[0].Limit).To(Equal(2)) + Expect(store.listFilters[0].Cursor).To(BeEmpty()) + Expect(store.listFilters[1].Limit).To(Equal(2)) + Expect(store.listFilters[1].Cursor).To(Equal("cursor-2")) + }) + + It("keeps live summaries and database coverage scoped to the returned page", func() { + free := 30 + page := sessionRecordPage{ + Records: []SessionRecord{ + {Tokens: &SessionTokensWire{InputTokens: 10, TotalTokens: 10}, Context: &SessionContextWire{FreePercent: free}, CostUSD: 0.25}, + {Tokens: &SessionTokensWire{OutputTokens: 5, TotalTokens: 5}}, + }, + Total: 9, + } + + result := buildSessionLiveResult(sessionLiveResultOptions{ + Page: page, Source: "all", Scope: "all", ReadAt: time.Date(2026, time.July, 16, 16, 0, 0, 0, time.UTC), + }) + + Expect(result.Total).To(Equal(9)) + Expect(result.Summary.TotalSessions).To(Equal(2)) + Expect(result.Summary.TotalTokens).To(Equal(15)) + Expect(result.Summary.LowestContextFree).To(PointTo(Equal(free))) + Expect(result.Database.Coverage).To(Equal("page")) + }) + + It("reports only records analyzed by a bounded throughput page", func() { + startedAt := time.Date(2026, time.July, 16, 16, 0, 0, 0, time.UTC) + endedAt := startedAt.Add(10 * time.Second) + page := sessionRecordPage{ + Records: []SessionRecord{{ + ID: "analyzed", Source: "codex", Model: "gpt-5", StartedAt: &startedAt, EndedAt: &endedAt, + Tokens: &SessionTokensWire{OutputTokens: 5, TotalTokens: 5}, + }}, + Total: 9, + } + + result := buildSessionThroughputResult(sessionThroughputResultOptions{Page: page, Source: "all", Scope: "all"}) + + Expect(result.Total).To(Equal(1)) + Expect(result.Groups).To(HaveLen(1)) + }) + + It("expands an exact root session ID into its complete thread", func() { + rootID := uuid.MustParse("055781c7-360a-4eb2-80be-452b3937fcfe") + childID := uuid.MustParse("7ca78c55-e280-50ff-a19a-9f355a6fc55e") + store := &sessionGetOverviewStore{ + identity: []database.SessionOverview{{ID: rootID, Source: "captain"}}, + thread: []database.SessionOverview{ + {ID: rootID, Source: "captain"}, + {ID: childID, ParentSessionID: &rootID, RootSessionID: &rootID, Source: "codex"}, + }, + } + + overviews, err := resolveOverviewsByIdentity(context.Background(), store, rootID.String()) + Expect(err).NotTo(HaveOccurred()) + Expect(overviews).To(Equal(store.thread)) + Expect(store.threadRoots).To(Equal([]uuid.UUID{rootID})) + }) + + It("hydrates transcript-less API sessions from persisted prompt runs", func(ctx SpecContext) { + rootID := uuid.MustParse("055781c7-360a-4eb2-80be-452b3937fcfe") + childID := uuid.MustParse("7ca78c55-e280-50ff-a19a-9f355a6fc55e") + runID := uuid.MustParse("293b06b4-f6b7-4f69-a531-7499bd5a473a") + agentType := "batch" + startedAt := time.Date(2026, time.July, 19, 9, 30, 0, 0, time.UTC) + finishedAt := startedAt.Add(8 * time.Second) + store := &sessionGetOverviewStore{ + identity: []database.SessionOverview{{ID: rootID, Source: "captain", AgentType: &agentType}}, + thread: []database.SessionOverview{ + {ID: rootID, Source: "captain", AgentType: &agentType}, + { + ID: childID, ParentSessionID: &rootID, RootSessionID: &rootID, Source: "captain", + Provider: "google", PromptRunCount: 1, + }, + }, + promptRuns: map[uuid.UUID][]database.PromptRun{ + childID: {{ + ID: runID, SessionID: childID, RootSessionID: rootID, + RenderedSpec: map[string]any{ + "name": "structured-ui-review", + "outputSchema": map[string]any{"type": "object"}, + }, + PromptMarkdown: "Review the attached form screenshot.", + ResultText: `{"summary":"Use a single-column form layout."}`, + Runtime: database.PromptRunRuntime{Resolved: database.PromptRunRuntimeSelection{ + Provider: "google", Backend: "gemini-api", Model: "gemini-2.5-pro", Effort: "high", + }}, + State: database.PromptRunStateSucceeded, Phase: database.PromptRunPhaseFinished, + StartedAt: &startedAt, FinishedAt: &finishedAt, + }}, + }, + } + + result, err := runSessionGet(ctx, store, SessionGetOptions{ID: rootID.String()}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.Sessions).To(HaveLen(2)) + Expect(result.Sessions[0].Aggregate).To(BeTrue()) + Expect(result.Pretty().String()).To(ContainSubstring("Aggregate session; child results are shown below")) + Expect(result.Pretty().String()).NotTo(ContainSubstring("Transcript: unavailable")) + child := result.Sessions[1] + Expect(child.DetailAvailable).To(BeTrue()) + Expect(child.Summary.DetailAvailable).To(BeTrue()) + Expect(child.Summary.Messages).To(Equal(2)) + Expect(child.Summary.Model).To(Equal("gemini-2.5-pro")) + Expect(child.Summary.Backend).To(Equal("gemini-api")) + Expect(child.Summary.Provider).To(Equal("google")) + Expect(child.Detail).NotTo(BeNil()) + Expect(child.Detail.Messages).To(Equal([]session.Message{ + { + ID: runID.String() + "-user", Role: "user", + Parts: []session.Part{{Type: session.PartText, Text: "Review the attached form screenshot."}}, + }, + { + ID: runID.String() + "-assistant", Role: "assistant", + Parts: []session.Part{{Type: session.PartText, Text: `{"summary":"Use a single-column form layout."}`}}, + }, + })) + Expect(child.Detail.Model).To(Equal("gemini-2.5-pro")) + Expect(child.Detail.Backend).To(Equal("gemini-api")) + Expect(child.Detail.Provider).To(Equal("google")) + Expect(child.Detail.StartedAt).To(PointTo(Equal(startedAt))) + Expect(child.Detail.EndedAt).To(PointTo(Equal(finishedAt))) + Expect(child.Detail.Prompt).To(MatchJSON(`{ + "name":"structured-ui-review", + "outputSchema":{"type":"object"} + }`)) + Expect(child.Detail.StructuredOutput).To(Equal(map[string]any{ + "summary": "Use a single-column form layout.", + })) + }) + + It("prefers persisted structured output and ignores non-schema JSON text", func() { + runID := uuid.MustParse("293b06b4-f6b7-4f69-a531-7499bd5a473a") + overview := database.SessionOverview{ID: uuid.New(), Source: "captain"} + detail, err := sessionFromPromptRun(overview, database.PromptRun{ + ID: runID, + RenderedSpec: map[string]any{ + "outputSchema": map[string]any{"type": "object"}, + }, + ResultText: `{"source":"text"}`, + ResultJSON: map[string]any{"source": "stored"}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(detail.StructuredOutput).To(Equal(map[string]any{"source": "stored"})) + Expect(detail.Messages).To(HaveLen(1)) + Expect(detail.Messages[0].Parts[0].Text).To(Equal(`{"source":"text"}`)) + + detail, err = sessionFromPromptRun(overview, database.PromptRun{ + ID: runID, ResultText: `{"source":"plain-text-prompt"}`, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(detail.StructuredOutput).To(BeNil()) + }) + + It("renders every match sequentially and preserves metadata-only sessions", func() { + result := SessionGetResult{ + Sessions: []SessionGetItem{ + { + CaptainID: "055781c7-360a-4eb2-80be-452b3937fcfe", + ProviderSessionID: "ad4c854e-cde6-4b99-99f3-667bf74112e3", + Host: "MacBook-Pro.local", + DetailAvailable: true, + Summary: SessionRecord{ + ID: "ad4c854e-cde6-4b99-99f3-667bf74112e3", Source: "claude", Project: "flanksource", + }, + Detail: &session.Session{ID: "ad4c854e-cde6-4b99-99f3-667bf74112e3", Source: "claude"}, + }, + { + CaptainID: "7ca78c55-e280-50ff-a19a-9f355a6fc55e", + ProviderSessionID: "ad4c854e-cde6-4b99-99f3-667bf74112e3", + Host: "local", + Summary: SessionRecord{ + ID: "ad4c854e-cde6-4b99-99f3-667bf74112e3", Source: "gavel", Project: "xero-cli", + }, + }, + }, + Total: 2, + } + + plain := result.Pretty().String() + Expect(plain).To(ContainSubstring("055781c7-360a-4eb2-80be-452b3937fcfe")) + Expect(plain).To(ContainSubstring("7ca78c55-e280-50ff-a19a-9f355a6fc55e")) + Expect(plain).To(ContainSubstring("Transcript: unavailable")) + Expect(plain).To(MatchRegexp("055781c7[\\s\\S]*7ca78c55")) + + markdown := result.Pretty().Markdown() + Expect(markdown).To(MatchRegexp("055781c7[\\s\\S]*7ca78c55")) + + html := result.Pretty().HTML() + Expect(html).NotTo(MatchRegexp(`text-gray-(600|700)`)) + Expect(html).To(ContainSubstring("text-muted")) + formattedHTML, err := clicky.Format(result, clicky.FormatOptions{Format: "html"}) + Expect(err).NotTo(HaveOccurred()) + Expect(formattedHTML).NotTo(MatchRegexp(`text-gray-(600|700)`)) + Expect(formattedHTML).To(ContainSubstring("text-muted")) + + wire, err := json.Marshal(result) + Expect(err).NotTo(HaveOccurred()) + Expect(wire).To(MatchJSON(`{ + "sessions": [ + { + "captainId": "055781c7-360a-4eb2-80be-452b3937fcfe", + "providerSessionId": "ad4c854e-cde6-4b99-99f3-667bf74112e3", + "host": "MacBook-Pro.local", + "detailAvailable": true, + "summary": {"key":"","id":"ad4c854e-cde6-4b99-99f3-667bf74112e3","source":"claude","project":"flanksource","toolCalls":0,"messages":0,"detailAvailable":false}, + "detail": {"id":"ad4c854e-cde6-4b99-99f3-667bf74112e3","source":"claude","git":{},"usage":{"inputTokens":0,"outputTokens":0},"cost":{"inputTokens":0,"outputTokens":0,"totalTokens":0,"inputCost":0,"outputCost":0},"capabilities":{},"files":{},"approvals":{"approved":0,"denied":0}} + }, + { + "captainId": "7ca78c55-e280-50ff-a19a-9f355a6fc55e", + "providerSessionId": "ad4c854e-cde6-4b99-99f3-667bf74112e3", + "host": "local", + "detailAvailable": false, + "summary": {"key":"","id":"ad4c854e-cde6-4b99-99f3-667bf74112e3","source":"gavel","project":"xero-cli","toolCalls":0,"messages":0,"detailAvailable":false} + } + ], + "total": 2 + }`)) + }) +}) + +type sessionGetOverviewStore struct { + identity []database.SessionOverview + thread []database.SessionOverview + list []database.SessionListSummary + listFilter database.SessionListFilter + listFilters []database.SessionListFilter + listPages map[string]database.SessionListPage + identities []string + threadRoots []uuid.UUID + listCalls int + promptRuns map[uuid.UUID][]database.PromptRun +} + +func (s *sessionGetOverviewStore) ListSessionSummaries(_ context.Context, filter database.SessionListFilter) (database.SessionListPage, error) { + s.listFilter = filter + s.listFilters = append(s.listFilters, filter) + if s.listPages != nil { + return s.listPages[filter.Cursor], nil + } + return database.SessionListPage{Rows: s.list, Total: int64(len(s.list))}, nil +} + +func (s *sessionGetOverviewStore) ListSessionOverviewsByIdentity(_ context.Context, identity string) ([]database.SessionOverview, error) { + s.identities = append(s.identities, identity) + return s.identity, nil +} + +func (s *sessionGetOverviewStore) ListSessionOverviews(context.Context, database.SessionOverviewFilter) ([]database.SessionOverview, error) { + s.listCalls++ + return nil, nil +} + +func (s *sessionGetOverviewStore) ListThreadSessionOverviews(_ context.Context, rootID uuid.UUID) ([]database.SessionOverview, error) { + s.threadRoots = append(s.threadRoots, rootID) + return s.thread, nil +} + +func (s *sessionGetOverviewStore) ListPromptRuns(_ context.Context, filter database.PromptRunFilter) ([]database.PromptRun, error) { + if filter.SessionID == nil { + return nil, nil + } + return s.promptRuns[*filter.SessionID], nil +} diff --git a/pkg/cli/session_list_render.go b/pkg/cli/session_list_render.go new file mode 100644 index 0000000..dfd2b2e --- /dev/null +++ b/pkg/cli/session_list_render.go @@ -0,0 +1,112 @@ +package cli + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/flanksource/clicky/api" +) + +const sessionIDDisplayWidth = 12 + +// sessionInitialPromptStyle collapses the prompt to a single line but sets no +// width cap: the column runs to whatever the terminal leaves it and the renderer +// truncates there, rather than at a width picked here that the terminal knows +// nothing about. +const sessionInitialPromptStyle = "max-lines-[1] truncate-suffix" + +var _ api.TableProvider = SessionRecord{} + +// Columns keeps the historical session list compact while leaving the full +// SessionRecord available to JSON and detail consumers. +func (SessionRecord) Columns() []api.ColumnDef { + return []api.ColumnDef{ + api.Column("age").Label("Age").Build(), + api.Column("source").Label("Agent").MaxWidth(12).Build(), + api.Column("project").Label("Project").MaxWidth(20).Build(), + api.Column("session").Label("Session").MaxWidth(sessionIDDisplayWidth).Build(), + api.Column("model").Label("Model").MaxWidth(24).Build(), + api.Column("initial_prompt").Label("Initial Prompt").Style(sessionInitialPromptStyle).Build(), + api.Column("tokens").Label("Tokens").Build(), + } +} + +func (r SessionRecord) Row() map[string]any { + row := map[string]any{ + "source": r.Source, + "session": sessionListID(r.ID), + "model": r.Model, + "initial_prompt": sessionInitialPromptCell(r.InitialPrompt), + } + if project := sessionProjectName(r); project != "" { + row["project"] = project + } + if age := sessionAge(r); age != nil { + row["age"] = age + } + if usage := sessionUsageCell(r); usage != nil { + row["tokens"] = usage + } + return row +} + +func sessionListID(id string) string { + if len(id) > sessionIDDisplayWidth { + return id[:sessionIDDisplayWidth] + } + return id +} + +func sessionInitialPromptCell(prompt string) api.Textable { + return api.Text{}.Append(prompt, sessionInitialPromptStyle) +} + +func sessionProjectName(r SessionRecord) string { + project := strings.TrimSpace(r.Project) + if project == "" { + project = strings.TrimSpace(r.CWD) + } + if project == "" { + return "" + } + return filepath.Base(filepath.Clean(project)) +} + +func sessionAge(r SessionRecord) api.Textable { + observedAt := sessionSortTime(r) + if observedAt.IsZero() { + return nil + } + return api.TimeAgo(&observedAt) +} + +// sessionUsageCell renders the one compact usage value shared by session list +// and ps tables: "27.4M $13.61". +func sessionUsageCell(r SessionRecord) api.Textable { + total := sessionTokenTotal(r) + if total == 0 && r.CostUSD <= 0 { + return nil + } + t := api.Text{} + if total > 0 { + t = t.Add(api.HumanNumber(total, "text-muted")) + } + if r.CostUSD > 0 { + if total > 0 { + t = t.Space() + } + t = t.Append(fmt.Sprintf("$%.2f", r.CostUSD), "text-green-600") + } + return t +} + +func sessionTokenTotal(r SessionRecord) int64 { + if r.Tokens == nil { + return 0 + } + if r.Tokens.TotalTokens > 0 { + return int64(r.Tokens.TotalTokens) + } + return int64(r.Tokens.InputTokens + r.Tokens.OutputTokens + r.Tokens.CacheReadTokens + r.Tokens.CacheCreationTokens) +} diff --git a/pkg/cli/session_list_render_test.go b/pkg/cli/session_list_render_test.go new file mode 100644 index 0000000..875fcd2 --- /dev/null +++ b/pkg/cli/session_list_render_test.go @@ -0,0 +1,100 @@ +package cli + +import ( + "strings" + "testing" + "time" + + "github.com/flanksource/clicky/api" +) + +func TestSessionRecordTableProviderIsCompact(t *testing.T) { + endedAt := time.Now().Add(-2 * time.Hour) + longTail := strings.Repeat("x", 120) + secondLine := "second line must not render" + record := SessionRecord{ + ID: "6522fe00-9a7c-4cee-a205-123456789abc", + Source: "codex", + Project: "/work/flanksource/captain", + Model: "gpt-5.6-sol", + InitialPrompt: "Improve the sessions table " + longTail + "\n" + secondLine, + EndedAt: &endedAt, + Tokens: &SessionTokensWire{TotalTokens: 12_345}, + CostUSD: 1.25, + } + + columns := record.Columns() + wantColumns := []string{"age", "source", "project", "session", "model", "initial_prompt", "tokens"} + if len(columns) != len(wantColumns) { + t.Fatalf("columns = %+v", columns) + } + for i, want := range wantColumns { + if columns[i].Name != want { + t.Fatalf("column %d = %q, want %q", i, columns[i].Name, want) + } + } + // The prompt column declares no width: it runs to whatever the terminal + // leaves it, and the renderer truncates there. + if columns[5].MaxWidth != 0 { + t.Fatalf("initial prompt max width = %d, want no cap", columns[5].MaxWidth) + } + if !strings.Contains(columns[5].Style, "max-lines-[1]") { + t.Fatalf("initial prompt style = %q", columns[5].Style) + } + + row := record.Row() + if columns[3].MaxWidth != sessionIDDisplayWidth { + t.Fatalf("session max width = %d, want %d", columns[3].MaxWidth, sessionIDDisplayWidth) + } + if row["session"] != "6522fe00-9a7" { + t.Fatalf("session = %q", row["session"]) + } + if row["project"] != "captain" { + t.Fatalf("project = %q", row["project"]) + } + usage, ok := row["tokens"].(api.Textable) + if !ok || !strings.Contains(usage.String(), "12.3K") || !strings.Contains(usage.String(), "$1.25") { + t.Fatalf("tokens = %#v", row["tokens"]) + } + if _, ok := row["age"].(api.Textable); !ok { + t.Fatalf("age = %#v", row["age"]) + } + + table := api.NewTableFrom([]SessionRecord{record}) + prompt := table.Rows[0]["initial_prompt"].String() + if strings.Contains(prompt, "\n") || strings.Contains(prompt, secondLine) { + t.Fatalf("initial prompt is not one line: %q", prompt) + } + // The cell keeps its full width -- only the extra line is dropped. Capping + // it here would truncate at a width the terminal knows nothing about. + if !strings.Contains(prompt, longTail) { + t.Fatalf("initial prompt was width-truncated before rendering: %q", prompt) + } + if !strings.HasPrefix(prompt, "Improve the sessions table") || !strings.HasSuffix(prompt, "…") { + t.Fatalf("initial prompt = %q", prompt) + } + + rendered := table.ANSI() + if !strings.Contains(rendered, "6522fe00-9a7") || !strings.Contains(rendered, "$1.25") { + t.Fatalf("ansi output missing compact identity/usage: %q", rendered) + } + if strings.Contains(rendered, longTail) || strings.Contains(rendered, secondLine) { + t.Fatalf("ansi output did not truncate the initial prompt: %q", rendered) + } + if lines := strings.Count(strings.Trim(rendered, "\n"), "\n") + 1; lines != 5 { + t.Fatalf("ansi output should be border/header/separator/row/border, got %d lines: %q", lines, rendered) + } +} + +func TestSessionRecordTableProviderHandlesSparseRows(t *testing.T) { + row := (SessionRecord{ID: "short", CWD: "/work/project"}).Row() + if row["session"] != "short" || row["project"] != "project" { + t.Fatalf("row = %+v", row) + } + if _, ok := row["age"]; ok { + t.Fatalf("zero age should be omitted: %+v", row) + } + if _, ok := row["tokens"]; ok { + t.Fatalf("zero usage should be omitted: %+v", row) + } +} diff --git a/pkg/cli/session_live.go b/pkg/cli/session_live.go index c117028..3fca76f 100644 --- a/pkg/cli/session_live.go +++ b/pkg/cli/session_live.go @@ -8,13 +8,21 @@ import ( "time" "github.com/flanksource/captain/pkg/claude" - rpchttp "github.com/flanksource/clicky/rpc/http" + "github.com/flanksource/captain/pkg/database" ) -var discoverSessionProcesses = discoverAgentProcesses - const defaultSessionLiveLimit = 25 +type SessionDatabaseStatusWire struct { + Source string `json:"source,omitempty"` + DSN string `json:"dsn,omitempty"` + Coverage string `json:"coverage"` + ReadAt time.Time `json:"readAt"` + LatestSampledAt *time.Time `json:"latestSampledAt,omitempty"` + LatestHeartbeatAt *time.Time `json:"latestHeartbeatAt,omitempty"` + EarliestLeaseExpiry *time.Time `json:"earliestLeaseExpiry,omitempty"` +} + func RunSessionLive(ctx context.Context, opts SessionLiveOptions) (SessionLiveResult, error) { source, err := normalizeSessionSource(opts.Source) if err != nil { @@ -25,159 +33,126 @@ func RunSessionLive(ctx context.Context, opts SessionLiveOptions) (SessionLiveRe return SessionLiveResult{}, err } - scope := "current" - if opts.All { - scope = "all" - } + scope, projectRoot, _ := resolveSessionScope(cwd, opts.All, opts.Project) limit := opts.Limit if limit <= 0 && !opts.Full { limit = defaultSessionLiveLimit } - records, err := discoverLiveSessionRecords(ctx, cwd, opts.All, source, limit, opts.Full) + db, err := captainDB(ctx) if err != nil { return SessionLiveResult{}, err } - - stopEnrich := rpchttp.Track(ctx, "enrich") - processes, _ := discoverSessionProcesses() - if !opts.All { - processes = filterAgentProcessesByProject(processes, sessionProjectRoot(cwd)) + query := sessionRecordQuery{ + Source: source, ProjectRoot: projectRoot, Query: opts.Query, LiveOnly: true, + Limit: limit, Cursor: opts.Cursor, } - records = enrichSessionsWithLive(records, processes) - filtered := make([]SessionRecord, 0, len(records)) - for _, record := range records { - if sessionMatchesQuery(record, opts.Query) { - filtered = append(filtered, record) - } - } - sortSessionRecords(filtered) - total := len(filtered) - summary := summarizeSessionDashboard(filtered) - stopEnrich() - if !opts.Full && limit > 0 && len(filtered) > limit { - filtered = filtered[:limit] - } - - return SessionLiveResult{ - Sessions: filtered, - Total: total, - Source: source, - Scope: scope, - Summary: summary, - }, nil -} - -func discoverLiveSessionRecords(ctx context.Context, cwd string, searchAll bool, source string, limit int, full bool) ([]SessionRecord, error) { - if full { - list, err := RunSessionList(ctx, SessionListOptions{ - Source: source, - All: searchAll, - Query: "", - Limit: 0, - }) - if err != nil { - return nil, err - } - return list.Sessions, nil + var page sessionRecordPage + if opts.Full { + page, err = dbAllSessionRecords(ctx, db, query) + } else { + page, err = dbSessionRecords(ctx, db, query) } - - candidates, err := discoverLiveSessionCandidates(ctx, cwd, searchAll, source, limit) if err != nil { - return nil, err + return SessionLiveResult{}, err } - records := make([]SessionRecord, 0, len(candidates)) - for _, candidate := range candidates { - records = append(records, candidate.record) + enrichLiveSessionSurfaces(page.Records) + coverage := "page" + if opts.Full { + coverage = "all" } - sortSessionRecords(records) - return records, nil + return buildSessionLiveResult(sessionLiveResultOptions{ + Page: page, Source: source, Scope: scope, Project: projectResultValue(scope, projectRoot), + ReadAt: time.Now().UTC(), DatabaseCoverage: coverage, + }), nil } -func sessionProjectRoot(cwd string) string { - if cwd == "" { - return "" - } - projectInfo := claude.FindProjectInfo(cwd) - if projectInfo.Root != "" { - return projectInfo.Root - } - return cwd +type sessionLiveResultOptions struct { + Page sessionRecordPage + Source string + Scope string + Project string + ReadAt time.Time + DatabaseCoverage string } -func filterAgentProcessesByProject(processes []agentProcess, projectRoot string) []agentProcess { - if projectRoot == "" { - return processes +func buildSessionLiveResult(options sessionLiveResultOptions) SessionLiveResult { + coverage := options.DatabaseCoverage + if coverage == "" { + coverage = "page" } - filtered := make([]agentProcess, 0, len(processes)) - for _, proc := range processes { - if sessionRecordMatchesProject(SessionRecord{CWD: proc.CWD}, projectRoot) { - filtered = append(filtered, proc) - } + databaseStatus := sessionDatabaseStatus(options.Page.Records, options.ReadAt) + databaseStatus.Coverage = coverage + return SessionLiveResult{ + Sessions: options.Page.Records, Total: options.Page.Total, Source: options.Source, + Scope: options.Scope, Project: options.Project, + Summary: summarizeSessionDashboard(options.Page.Records), Database: databaseStatus, + NextCursor: options.Page.NextCursor, } - return filtered } -func enrichSessionsWithLive(records []SessionRecord, processes []agentProcess) []SessionRecord { - out := make([]SessionRecord, len(records)) - copy(out, records) - matched := make(map[int]bool) - for i := range out { - idx := bestLiveMatch(out[i], processes, matched) - if idx < 0 { - out[i].Health = deriveSessionHealth(out[i]) +func enrichLiveSessionSurfaces(records []SessionRecord) { + detected := false + for i := range records { + if records[i].Live == nil || records[i].Live.PID <= 0 { continue } - matched[idx] = true - out[i].Live = processes[idx].wire() - if out[i].CWD == "" { - out[i].CWD = processes[idx].CWD + if records[i].Live.Surface != nil { + detected = true } - if out[i].StartedAt == nil && processes[idx].StartedAt != nil { - out[i].StartedAt = processes[idx].StartedAt + if surface := discoverProcessSurface(records[i].Live.PID); surface != nil { + records[i].Live.Surface = surface + detected = true } - out[i].Health = deriveSessionHealth(out[i]) } - for i, proc := range processes { - if matched[i] { + if !detected { + return + } + surfaces, err := discoverCmuxSurfaces() + if err != nil { + return + } + for i := range records { + if records[i].Live == nil { continue } - record := SessionRecord{ - Key: fmt.Sprintf("live-%s-%d", proc.Source, proc.PID), - ID: fmt.Sprintf("pid:%d", proc.PID), - Source: proc.Source, - CWD: proc.CWD, - StartedAt: proc.StartedAt, - DetailAvailable: false, - Live: proc.wire(), + if surface := records[i].Live.Surface; surface != nil && surface.SurfaceID != "" { + enrichCmuxSurface(surface, surfaces) } - record.Health = deriveSessionHealth(record) - out = append(out, record) } - sortSessionRecords(out) - return out } -func bestLiveMatch(record SessionRecord, processes []agentProcess, matched map[int]bool) int { - if record.Source == "" { - return -1 - } - best := -1 - for i, proc := range processes { - if matched[i] || proc.Source != record.Source { +func sessionDatabaseStatus(records []SessionRecord, readAt time.Time) SessionDatabaseStatusWire { + dsn, source := captainDatabaseIdentity() + status := SessionDatabaseStatusWire{Source: source, DSN: database.MaskDSN(dsn), ReadAt: readAt} + for _, record := range records { + if record.Live == nil { continue } - if record.CWD != "" && proc.CWD != "" && samePath(record.CWD, proc.CWD) { - return i + if sampledAt := record.Live.SampledAt; sampledAt != nil && + (status.LatestSampledAt == nil || sampledAt.After(*status.LatestSampledAt)) { + status.LatestSampledAt = sampledAt } - if best < 0 && record.CWD == "" { - best = i + if heartbeatAt := record.Live.LastHeartbeatAt; heartbeatAt != nil && + (status.LatestHeartbeatAt == nil || heartbeatAt.After(*status.LatestHeartbeatAt)) { + status.LatestHeartbeatAt = heartbeatAt + } + if expiresAt := record.Live.LeaseExpiresAt; expiresAt != nil && + (status.EarliestLeaseExpiry == nil || expiresAt.Before(*status.EarliestLeaseExpiry)) { + status.EarliestLeaseExpiry = expiresAt } } - return best + return status } -func samePath(a, b string) bool { - return strings.TrimRight(a, "/") == strings.TrimRight(b, "/") +func sessionProjectRoot(cwd string) string { + if cwd == "" { + return "" + } + projectInfo := claude.FindProjectInfo(cwd) + if projectInfo.Root != "" { + return projectInfo.Root + } + return cwd } func deriveSessionHealth(record SessionRecord) []SessionHealthWire { @@ -298,6 +273,11 @@ func liveMatchesQuery(live *SessionLiveWire, query string) bool { live.Status, live.CWD, live.Command, + live.SessionID, + } + values = append(values, live.AgentIDs...) + if live.Surface != nil { + values = append(values, live.Surface.SurfaceID, live.Surface.TabID, live.Surface.WorkspaceID, live.Surface.AgentKind) } for _, value := range values { if strings.Contains(strings.ToLower(value), query) { diff --git a/pkg/cli/session_live_db_test.go b/pkg/cli/session_live_db_test.go new file mode 100644 index 0000000..3f138fc --- /dev/null +++ b/pkg/cli/session_live_db_test.go @@ -0,0 +1,107 @@ +package cli + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/monitor" +) + +func TestRunSessionLiveReadsStoredProcessStatusWithoutPolling(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + db := withTestCaptainDB(t) + project := filepath.Join(home, "work", "project") + if err := os.MkdirAll(project, 0o755); err != nil { + t.Fatal(err) + } + t.Chdir(project) + + stored, err := db.CreateOrGetSession(t.Context(), database.CreateSessionInput{ + ProviderSessionID: "sess-stored", Source: "codex", CWD: project, + }) + if err != nil { + t.Fatalf("create stored session: %v", err) + } + started := time.Now().UTC().Add(-time.Minute).Truncate(time.Second) + sampledAt := started.Add(30 * time.Second) + expiresAt := sampledAt.Add(time.Minute) + if err := db.UpsertSessionProcess(t.Context(), database.SessionProcessInput{ + SessionID: stored.ID, HostID: captainHostID(), BootID: "boot", PID: 24680, + ProcessStartedAt: started, SampledAt: sampledAt, + Status: "sleeping", CWD: project, Source: "codex", + }); err != nil { + t.Fatalf("upsert stored process: %v", err) + } + if err := db.Gorm().WithContext(t.Context()).Exec(` + UPDATE captain_session_processes + SET lease_owner = ?, lease_token = ?, lease_expires_at = ? + WHERE session_id = ?`, "captain-serve", "6522fe00-9a7c-4cee-a205-123456789abc", expiresAt, stored.ID).Error; err != nil { + t.Fatalf("set stored process lease: %v", err) + } + + discoveryCalls := 0 + monitorDiscoverProcesses = func() ([]monitor.Process, error) { + discoveryCalls++ + return nil, nil + } + + result, err := RunSessionLive(context.Background(), SessionLiveOptions{Source: "all", Limit: 10}) + if err != nil { + t.Fatalf("RunSessionLive: %v", err) + } + if discoveryCalls != 0 { + t.Fatalf("process discovery calls = %d, want 0", discoveryCalls) + } + if result.Total != 1 || len(result.Sessions) != 1 { + t.Fatalf("live sessions = %+v", result) + } + if result.Database.ReadAt.IsZero() { + t.Fatal("database read time is missing") + } + if result.Database.LatestSampledAt == nil || !result.Database.LatestSampledAt.Equal(sampledAt) { + t.Fatalf("database latest sample = %v, want %v", result.Database.LatestSampledAt, sampledAt) + } + if result.Database.LatestHeartbeatAt == nil || !result.Database.LatestHeartbeatAt.Equal(sampledAt) { + t.Fatalf("database latest heartbeat = %v, want %v", result.Database.LatestHeartbeatAt, sampledAt) + } + if result.Database.EarliestLeaseExpiry == nil || !result.Database.EarliestLeaseExpiry.Equal(expiresAt) { + t.Fatalf("database lease expiry = %v, want %v", result.Database.EarliestLeaseExpiry, expiresAt) + } + live := result.Sessions[0].Live + if live == nil || live.PID != 24680 || live.Status != "sleeping" { + t.Fatalf("stored process status = %+v", live) + } + if live.SampledAt == nil || !live.SampledAt.Equal(sampledAt) { + t.Fatalf("stored sample time = %v, want %v", live.SampledAt, sampledAt) + } + if live.LastHeartbeatAt == nil || !live.LastHeartbeatAt.Equal(sampledAt) { + t.Fatalf("stored heartbeat = %v, want %v", live.LastHeartbeatAt, sampledAt) + } + if live.LeaseOwner != "captain-serve" { + t.Fatalf("stored lease owner = %q", live.LeaseOwner) + } + if live.LeaseExpiresAt == nil || !live.LeaseExpiresAt.Equal(expiresAt) { + t.Fatalf("stored lease expiry = %v, want %v", live.LeaseExpiresAt, expiresAt) + } +} + +func refreshTestSessionDB(t *testing.T) { + t.Helper() + if _, err := freshenSessionDB(t.Context()); err != nil { + t.Fatalf("refresh test session database: %v", err) + } +} + +func hasHealth(signals []SessionHealthWire, kind, severity string) bool { + for _, signal := range signals { + if signal.Kind == kind && signal.Severity == severity { + return true + } + } + return false +} diff --git a/pkg/cli/session_live_render.go b/pkg/cli/session_live_render.go new file mode 100644 index 0000000..9179417 --- /dev/null +++ b/pkg/cli/session_live_render.go @@ -0,0 +1,108 @@ +package cli + +import ( + "strconv" + "strings" + "time" + + "github.com/flanksource/clicky" + "github.com/flanksource/clicky/api" +) + +type sessionLiveRow struct { + SessionRecord +} + +type sessionLiveConfiguration struct { + mode string + model string + effort string +} + +func (c sessionLiveConfiguration) Pretty() api.Text { + return api.Text{}. + Append(sessionLiveConfigValue(c.mode), "text-cyan-600"). + Append(" · ", "text-gray-400"). + Append(sessionLiveConfigValue(c.model), "text-purple-600 font-medium"). + Append(" · ", "text-gray-400"). + Append(sessionLiveConfigValue(c.effort), "text-amber-600") +} + +func (c sessionLiveConfiguration) String() string { return c.Pretty().String() } +func (c sessionLiveConfiguration) ANSI() string { return c.Pretty().ANSI() } +func (c sessionLiveConfiguration) HTML() string { return c.Pretty().HTML() } +func (c sessionLiveConfiguration) Markdown() string { return c.String() } + +func (r SessionLiveResult) Pretty() api.Text { + diagnostics := map[string]any{ + "Database": r.Database.Source, + "DSN": r.Database.DSN, + "Read": api.Human(r.Database.ReadAt, "text-muted"), + "Latest sample": sessionLiveTime(r.Database.LatestSampledAt), + "Latest heartbeat": sessionLiveTime(r.Database.LatestHeartbeatAt), + "Lease expiry": sessionLiveTime(r.Database.EarliestLeaseExpiry), + "Live sessions": r.Total, + } + rows := make([]sessionLiveRow, len(r.Sessions)) + for i := range r.Sessions { + rows[i] = sessionLiveRow{SessionRecord: r.Sessions[i]} + } + return api.Text{}. + Add(clicky.Map(diagnostics)). + NewLine(). + Add(api.NewTableFrom(rows)) +} + +func (sessionLiveRow) Columns() []api.ColumnDef { + return []api.ColumnDef{ + api.Column("configuration").Label("Agent").MaxWidth(32).Build(), + api.Column("project").Label("Project").MaxWidth(16).Build(), + api.Column("session").Label("Session").MaxWidth(sessionIDDisplayWidth).Build(), + api.Column("pid").Label("PID").MaxWidth(20).Build(), + api.Column("status").Label("Status").MaxWidth(8).Build(), + api.Column("title").Label("Title").Build(), + } +} + +func (r sessionLiveRow) Row() map[string]any { + title := strings.TrimSpace(r.Title) + if prompt := strings.TrimSpace(r.InitialPrompt); prompt != "" && (r.Source == "codex" || title == "") { + title, _, _ = strings.Cut(prompt, "\n") + title = strings.TrimSpace(title) + } + row := map[string]any{ + "configuration": sessionLiveConfiguration{mode: r.Backend, model: r.Model, effort: r.ReasoningEffort}, + "project": sessionProjectName(r.SessionRecord), + "session": sessionListID(r.ID), + "title": title, + } + if r.Live == nil { + return row + } + identity := make([]string, 0, 2) + if r.Live.PID > 0 { + identity = append(identity, strconv.Itoa(r.Live.PID)) + } + if r.Live.Surface != nil && r.Live.Surface.SurfaceRef != "" { + identity = append(identity, r.Live.Surface.SurfaceRef) + } + if len(identity) > 0 { + row["pid"] = strings.Join(identity, " ") + } + row["status"] = r.Live.Status + return row +} + +func sessionLiveConfigValue(value string) string { + if value = strings.TrimSpace(value); value != "" { + return value + } + return "—" +} + +func sessionLiveTime(value *time.Time) api.Textable { + if value == nil { + return api.Text{}.Append("—", "text-muted") + } + return api.Human(*value, "text-muted") +} diff --git a/pkg/cli/session_live_render_test.go b/pkg/cli/session_live_render_test.go new file mode 100644 index 0000000..2567ae6 --- /dev/null +++ b/pkg/cli/session_live_render_test.go @@ -0,0 +1,184 @@ +package cli + +import ( + "strings" + "testing" + "time" + + "github.com/flanksource/captain/pkg/cmux" + "github.com/flanksource/captain/pkg/database" + clickyapi "github.com/flanksource/clicky/api" +) + +func TestSessionLiveResultPrettyReportsDatabaseDiagnostics(t *testing.T) { + sampledAt := time.Date(2026, time.July, 13, 15, 0, 0, 0, time.UTC) + expiresAt := sampledAt.Add(time.Minute) + result := SessionLiveResult{ + Sessions: []SessionRecord{{ + ID: "6522fe00-9a7c-4cee-a205-123456789abc", Source: "codex", Project: "/work/captain", + Backend: "codex-cmux", Model: "gpt-5.6-sol", ReasoningEffort: "high", + InitialPrompt: "Diagnose the monitor lock without polling", + Live: &SessionLiveWire{ + PID: 24680, Status: "sleeping", SampledAt: &sampledAt, LastHeartbeatAt: &sampledAt, + LeaseOwner: "captain-serve", LeaseExpiresAt: &expiresAt, + }, + }}, + Total: 1, + Database: SessionDatabaseStatusWire{ + Source: "captain embedded database", DSN: "postgres://localhost/captain", ReadAt: sampledAt, + LatestSampledAt: &sampledAt, LatestHeartbeatAt: &sampledAt, EarliestLeaseExpiry: &expiresAt, + }, + } + + rendered := result.Pretty().String() + for _, expected := range []string{ + "captain embedded database", "Latest sample", "Latest heartbeat", "Lease expiry", + "6522fe00", "24680", "sleeping", + } { + if !strings.Contains(rendered, expected) { + t.Fatalf("pretty output missing %q: %s", expected, rendered) + } + } +} + +func TestSessionLiveResultPrettyKeepsUnavailableDatabaseDiagnosticsVisible(t *testing.T) { + rendered := (SessionLiveResult{Database: SessionDatabaseStatusWire{ReadAt: time.Now()}}).Pretty().String() + for _, expected := range []string{"Latest sample", "Latest heartbeat", "Lease expiry", "—"} { + if !strings.Contains(rendered, expected) { + t.Fatalf("pretty output missing %q: %s", expected, rendered) + } + } +} + +func TestSessionLiveRowColumnsIncludeModelConfigurationAndTitle(t *testing.T) { + row := sessionLiveRow{SessionRecord: SessionRecord{ + Backend: "codex-cli", Model: "gpt-5.5", ReasoningEffort: "medium", + Title: "Stored title", InitialPrompt: "Fallback prompt", + }} + wantColumns := []string{ + "configuration", "project", "session", "pid", "status", "title", + } + columns := row.Columns() + if len(columns) != len(wantColumns) { + t.Fatalf("columns = %+v", columns) + } + if columns[0].Label != "Agent" { + t.Fatalf("configuration label = %q, want Agent", columns[0].Label) + } + for i, want := range wantColumns { + if columns[i].Name != want { + t.Fatalf("column %d = %q, want %q", i, columns[i].Name, want) + } + } + wantWidths := []int{32, 16, sessionIDDisplayWidth, 20, 8, 0} + for i, want := range wantWidths { + if columns[i].MaxWidth != want { + t.Fatalf("column %q width = %d, want %d", columns[i].Name, columns[i].MaxWidth, want) + } + } + if columns[5].MaxWidth != 0 || columns[5].Style != "" { + t.Fatalf("title column = %+v", columns[5]) + } + values := row.Row() + configuration := values["configuration"].(interface { + String() string + HTML() string + Markdown() string + }) + if configuration.String() != "codex-cli · gpt-5.5 · medium" { + t.Fatalf("configuration = %q", configuration.String()) + } + if configuration.Markdown() != configuration.String() { + t.Fatalf("configuration markdown = %q, want plain text", configuration.Markdown()) + } + for _, style := range []string{"text-cyan-600", "text-purple-600", "text-amber-600"} { + if !strings.Contains(configuration.HTML(), style) { + t.Fatalf("configuration HTML missing %q: %s", style, configuration.HTML()) + } + } + ansiTable := clickyapi.NewTableFrom([]sessionLiveRow{row}).ANSI() + if !strings.Contains(ansiTable, values["configuration"].(interface{ ANSI() string }).ANSI()) { + t.Fatalf("table ANSI dropped configuration colors: %q", ansiTable) + } + if title := values["title"]; title != "Stored title" { + t.Fatalf("title = %q, want stored title", title) + } + withID := sessionLiveRow{SessionRecord: SessionRecord{ID: "6522fe00-9a7c-4cee-a205-123456789abc"}}.Row() + if withID["session"] != "6522fe00-9a7" { + t.Fatalf("session = %q, want wider copyable prefix", withID["session"]) + } +} + +func TestSessionLiveRowPIDIncludesCmuxSurfaceRef(t *testing.T) { + row := sessionLiveRow{SessionRecord: SessionRecord{Live: &SessionLiveWire{ + PID: 24680, + Surface: &CmuxSurface{ + SurfaceRef: "surface:383", + }, + }}}.Row() + if got := row["pid"]; got != "24680 surface:383" { + t.Fatalf("pid = %q, want process and cmux surface refs", got) + } +} + +func TestSessionLiveRowCodexTitleUsesFullFirstPromptLine(t *testing.T) { + firstLine := "Review dependency installation versions and identify duplication using jscpd" + row := sessionLiveRow{SessionRecord: SessionRecord{ + Source: "codex", + Title: "Review dependency installation versions and identify duplication…", + InitialPrompt: firstLine + "\n\nDetails that do not belong in the title", + }}.Row() + if got := row["title"]; got != firstLine { + t.Fatalf("title = %q, want full first prompt line", got) + } +} + +func TestEnrichLiveSessionSurfacesResolvesShortRef(t *testing.T) { + const ( + pid = 24680 + surfaceID = "4F846CB1-2EE5-4359-8D5B-A0F6F3837952" + ) + defer stubPSDiscovery(t, nil, nil, func(gotPID int) *CmuxSurface { + if gotPID != pid { + t.Fatalf("surface lookup pid = %d, want %d", gotPID, pid) + } + return &CmuxSurface{SurfaceID: surfaceID} + }, map[string]cmux.Surface{ + surfaceID: {ID: surfaceID, Ref: "surface:383"}, + })() + + records := []SessionRecord{{Live: &SessionLiveWire{PID: pid}}} + enrichLiveSessionSurfaces(records) + if got := records[0].Live.Surface; got == nil || got.SurfaceRef != "surface:383" { + t.Fatalf("surface = %+v, want short cmux ref", got) + } +} + +func TestSessionLiveRowTitleFallsBackToInitialPrompt(t *testing.T) { + prompt := "Fallback prompt " + strings.Repeat("x", 80) + liveRow := sessionLiveRow{SessionRecord: SessionRecord{InitialPrompt: prompt}} + row := liveRow.Row() + title := row["title"] + if title != prompt { + t.Fatalf("title = %q, want full fallback prompt", title) + } + if rendered := clickyapi.NewTableFrom([]sessionLiveRow{liveRow}).String(); !strings.Contains(rendered, prompt) { + t.Fatalf("rendered table truncated title: %s", rendered) + } +} + +func TestSessionLiveRowKeepsEmptyModelConfigurationColumnsVisible(t *testing.T) { + row := (sessionLiveRow{}).Row() + configuration := row["configuration"].(interface{ String() string }).String() + if configuration != "— · — · —" { + t.Fatalf("configuration = %q, want missing-value markers", configuration) + } +} + +func TestRecordFromOverviewIncludesDatabaseBackend(t *testing.T) { + backend := "claude-cmux" + record := recordFromOverview(database.SessionOverview{Backend: &backend}) + if record.Backend != backend { + t.Fatalf("backend = %q, want %q", record.Backend, backend) + } +} diff --git a/pkg/cli/session_live_summary.go b/pkg/cli/session_live_summary.go deleted file mode 100644 index f82d257..0000000 --- a/pkg/cli/session_live_summary.go +++ /dev/null @@ -1,90 +0,0 @@ -package cli - -import ( - "context" - "os" - "sort" - - "github.com/flanksource/captain/pkg/ai/history" - "github.com/flanksource/captain/pkg/claude" - rpchttp "github.com/flanksource/clicky/rpc/http" -) - -type liveSessionFile struct { - source string - path string - modUnix int64 -} - -func discoverLiveSessionCandidates(ctx context.Context, cwd string, searchAll bool, source string, limit int) ([]sessionCandidate, error) { - stopFind := rpchttp.Track(ctx, "find") - files, err := discoverLiveSessionFiles(cwd, searchAll, source) - stopFind() - if err != nil { - return nil, err - } - sort.Slice(files, func(i, j int) bool { - if files[i].modUnix == files[j].modUnix { - return files[i].path > files[j].path - } - return files[i].modUnix > files[j].modUnix - }) - if limit > 0 && len(files) > limit { - files = files[:limit] - } - - refs := make([]sessionFileRef, len(files)) - for i, file := range files { - refs[i] = sessionFileRef{source: file.source, path: file.path} - } - return summarizeSessionRefs(ctx, refs), nil -} - -func discoverLiveSessionFiles(cwd string, searchAll bool, source string) ([]liveSessionFile, error) { - var files []liveSessionFile - if source == "all" || source == "claude" { - claudeFiles, err := claude.FindSessionFiles(claude.GetProjectsDir(), cwd, searchAll) - if err != nil { - return nil, err - } - files = appendSessionFiles(files, "claude", claudeFiles) - } - if source == "all" || source == "codex" { - codexFiles, err := history.FindCodexSessionFiles() - if err != nil { - return nil, err - } - matchRoot := cwd - if cwd != "" { - projectInfo := claude.FindProjectInfo(cwd) - if projectInfo.Root != "" { - matchRoot = projectInfo.Root - } - } - for _, file := range codexFiles { - if !searchAll { - meta, err := history.ReadCodexSessionMeta(file) - if err != nil || meta == nil || !codexMetaMatchesProject(meta, matchRoot) { - continue - } - } - files = appendSessionFiles(files, "codex", []string{file}) - } - } - return files, nil -} - -func appendSessionFiles(out []liveSessionFile, source string, paths []string) []liveSessionFile { - for _, path := range paths { - info, err := os.Stat(path) - if err != nil || info.IsDir() { - continue - } - out = append(out, liveSessionFile{ - source: source, - path: path, - modUnix: info.ModTime().UnixNano(), - }) - } - return out -} diff --git a/pkg/cli/session_openfiles.go b/pkg/cli/session_openfiles.go new file mode 100644 index 0000000..5a96774 --- /dev/null +++ b/pkg/cli/session_openfiles.go @@ -0,0 +1,100 @@ +package cli + +import ( + "bytes" + "os/exec" + "path/filepath" + "regexp" + "runtime" + "strconv" + "strings" + + "github.com/flanksource/captain/pkg/ai/history" +) + +// codexRolloutIDPattern matches the trailing UUID in a codex rollout filename +// (rollout--.jsonl). The timestamp segment never forms a +// UUID group, so the first match is always the session id. +var codexRolloutIDPattern = regexp.MustCompile(`[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}`) + +// processOpenSessionFiles maps each pid to the session transcript files it holds +// open, via a single `lsof` pass. Enrichment is best-effort: when some pids have +// already exited lsof exits non-zero but still streams the rest, which is parsed. +func processOpenSessionFiles(pids []int) map[int][]string { + if runtime.GOOS == "windows" { + return map[int][]string{} + } + var list []string + for _, pid := range pids { + if pid > 0 { + list = append(list, strconv.Itoa(pid)) + } + } + if len(list) == 0 { + return map[int][]string{} + } + out, _ := exec.Command("lsof", "-n", "-P", "-w", "-F", "pn", "-p", strings.Join(list, ",")).Output() + return parseLsofOpenFiles(out) +} + +// parseLsofOpenFiles parses `lsof -F pn` output into pid → open session-file +// paths, keeping only claude/codex transcript files. +func parseLsofOpenFiles(out []byte) map[int][]string { + files := make(map[int][]string) + current := 0 + for _, raw := range bytes.Split(out, []byte{'\n'}) { + line := strings.TrimSpace(string(raw)) + if line == "" { + continue + } + switch line[0] { + case 'p': + if pid, err := strconv.Atoi(line[1:]); err == nil { + current = pid + } + case 'n': + if current <= 0 { + continue + } + path := line[1:] + if source, _, _ := classifyOpenSessionFile(path); source != "" { + files[current] = append(files[current], path) + } + } + } + return files +} + +// classifyOpenSessionFile identifies a session transcript by path, returning its +// source ("claude"|"codex"), session/agent id, and kind ("root"|"subagent"| +// "codex"). A non-transcript path yields empty strings. +func classifyOpenSessionFile(path string) (source, id, kind string) { + if !strings.HasSuffix(path, ".jsonl") { + return "", "", "" + } + base := filepath.Base(path) + if history.IsCodexSession(path) { + return "codex", codexRolloutID(base), "codex" + } + if strings.Contains(path, "/.claude/projects/") { + if strings.HasPrefix(base, "agent-") || strings.Contains(path, "/subagents/") { + return "claude", claudeAgentID(base), "subagent" + } + return "claude", sessionIDFromFile(path), "root" + } + return "", "", "" +} + +// codexRolloutID extracts the session UUID from a codex rollout filename. +func codexRolloutID(base string) string { + if m := codexRolloutIDPattern.FindString(base); m != "" { + return m + } + return strings.TrimSuffix(base, filepath.Ext(base)) +} + +// claudeAgentID extracts the agent id from a sub-agent transcript filename +// ("agent-.jsonl" → ""). +func claudeAgentID(base string) string { + return strings.TrimPrefix(strings.TrimSuffix(base, ".jsonl"), "agent-") +} diff --git a/pkg/cli/session_openfiles_test.go b/pkg/cli/session_openfiles_test.go new file mode 100644 index 0000000..4a06d49 --- /dev/null +++ b/pkg/cli/session_openfiles_test.go @@ -0,0 +1,75 @@ +package cli + +import ( + "reflect" + "testing" +) + +func TestClassifyOpenSessionFile(t *testing.T) { + cases := []struct { + name string + path string + wantSource string + wantID string + wantKind string + }{ + { + name: "codex rollout", + path: "/Users/moshe/.codex/sessions/2026/07/09/rollout-2026-07-09T08-45-27-019f4568-d9a2-76c0-9ab1-2e465c2d5e35.jsonl", + wantSource: "codex", + wantID: "019f4568-d9a2-76c0-9ab1-2e465c2d5e35", + wantKind: "codex", + }, + { + name: "claude root", + path: "/Users/moshe/.claude/projects/-Users-moshe-work/b3c450b5-aab4-4834-9420-33a1839e15b8.jsonl", + wantSource: "claude", + wantID: "b3c450b5-aab4-4834-9420-33a1839e15b8", + wantKind: "root", + }, + { + name: "claude subagent", + path: "/Users/moshe/.claude/projects/-Users-moshe-work/b3c450b5/subagents/agent-ae31f6111faba5d7b.jsonl", + wantSource: "claude", + wantID: "ae31f6111faba5d7b", + wantKind: "subagent", + }, + {name: "non-jsonl", path: "/Users/moshe/.codex/logs_2.sqlite"}, + {name: "unrelated jsonl", path: "/Users/moshe/notes/todo.jsonl"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + source, id, kind := classifyOpenSessionFile(tc.path) + if source != tc.wantSource || id != tc.wantID || kind != tc.wantKind { + t.Fatalf("classify(%q) = (%q,%q,%q), want (%q,%q,%q)", + tc.path, source, id, kind, tc.wantSource, tc.wantID, tc.wantKind) + } + }) + } +} + +func TestParseLsofOpenFiles(t *testing.T) { + // Two codex rollouts on pid 84650 (main + sub-agent), one claude root on + // pid 77310; sqlite/dir/socket fds must be filtered out. + out := []byte("p84650\n" + + "n/Users/moshe/.codex/logs_2.sqlite\n" + + "n/Users/moshe/.codex/sessions/2026/07/09/rollout-2026-07-09T08-45-27-019f4568-d9a2-76c0-9ab1-2e465c2d5e35.jsonl\n" + + "n/Users/moshe/.codex/sessions/2026/07/09/rollout-2026-07-09T08-45-27-019f4568-da0d-7b82-b9f7-f24a2348a699.jsonl\n" + + "p77310\n" + + "n/Users/moshe/go/src/project\n" + + "n/Users/moshe/.claude/projects/-Users-moshe-work/b3c450b5-aab4-4834-9420-33a1839e15b8.jsonl\n") + + files := parseLsofOpenFiles(out) + + wantCodex := []string{ + "/Users/moshe/.codex/sessions/2026/07/09/rollout-2026-07-09T08-45-27-019f4568-d9a2-76c0-9ab1-2e465c2d5e35.jsonl", + "/Users/moshe/.codex/sessions/2026/07/09/rollout-2026-07-09T08-45-27-019f4568-da0d-7b82-b9f7-f24a2348a699.jsonl", + } + if !reflect.DeepEqual(files[84650], wantCodex) { + t.Fatalf("pid 84650 files = %v", files[84650]) + } + wantClaude := []string{"/Users/moshe/.claude/projects/-Users-moshe-work/b3c450b5-aab4-4834-9420-33a1839e15b8.jsonl"} + if !reflect.DeepEqual(files[77310], wantClaude) { + t.Fatalf("pid 77310 files = %v", files[77310]) + } +} diff --git a/pkg/cli/session_process.go b/pkg/cli/session_process.go index 0bdf3cf..b191457 100644 --- a/pkg/cli/session_process.go +++ b/pkg/cli/session_process.go @@ -20,6 +20,11 @@ type agentProcess struct { StartedAt *time.Time CWD string Command string + SessionID string + AgentIDs []string + LastActivity *time.Time + SessionFile string + Surface *CmuxSurface } func (p agentProcess) wire() *SessionLiveWire { @@ -32,9 +37,48 @@ func (p agentProcess) wire() *SessionLiveWire { StartedAt: p.StartedAt, CWD: p.CWD, Command: p.Command, + SessionID: p.SessionID, + AgentIDs: p.AgentIDs, + LastActivity: p.LastActivity, + SessionFile: p.SessionFile, + Surface: p.Surface, } } +// discoverSessionProcesses is indirected so ps tests can fake process lists. +var discoverSessionProcesses = discoverAgentProcesses + +func filterAgentProcessesByProject(processes []agentProcess, projectRoot string) []agentProcess { + if projectRoot == "" { + return processes + } + filtered := make([]agentProcess, 0, len(processes)) + for _, proc := range processes { + if sessionRecordMatchesProject(SessionRecord{CWD: proc.CWD}, projectRoot) { + filtered = append(filtered, proc) + } + } + return filtered +} + +// parseClaudeSessionIDFromCommand extracts the session id claude was launched +// with, from its "--session-id " / "--resume " argv (either the +// space-separated or "=" form). Returns "" when absent. +func parseClaudeSessionIDFromCommand(command string) string { + fields := strings.Fields(command) + for i, field := range fields { + for _, flag := range []string{"--session-id", "--resume"} { + if field == flag && i+1 < len(fields) { + return fields[i+1] + } + if value, ok := strings.CutPrefix(field, flag+"="); ok { + return value + } + } + } + return "" +} + func discoverAgentProcesses() ([]agentProcess, error) { if runtime.GOOS == "windows" { return nil, nil @@ -109,6 +153,11 @@ func processSource(command string) string { strings.Contains(lower, "codex-linux") || strings.Contains(lower, "codex-win") || commandNameMatches(lower, "codex") { + // mcp-server / app-server are codex's tool/IPC servers, not interactive + // sessions — they never hold a rollout transcript open. + if commandNameMatches(lower, "mcp-server") || commandNameMatches(lower, "app-server") { + return "" + } return "codex" } return "" diff --git a/pkg/cli/session_process_test.go b/pkg/cli/session_process_test.go index 2af20a8..dc77afe 100644 --- a/pkg/cli/session_process_test.go +++ b/pkg/cli/session_process_test.go @@ -26,6 +26,20 @@ func TestParseAgentProcessLineIgnoresCaptain(t *testing.T) { } } +func TestParseClaudeSessionIDFromCommand(t *testing.T) { + cases := map[string]string{ + "/Users/moshe/.local/bin/claude --session-id e4352f60-b395-4b43-996d-0bfc92a92ee0 --settings {}": "e4352f60-b395-4b43-996d-0bfc92a92ee0", + "claude --resume 50a80579-0108-464b-a555-a3620542f3f9": "50a80579-0108-464b-a555-a3620542f3f9", + "claude --session-id=abc123 -p": "abc123", + "node /path/codex.js mcp-server": "", + } + for command, want := range cases { + if got := parseClaudeSessionIDFromCommand(command); got != want { + t.Fatalf("parseClaudeSessionIDFromCommand(%q) = %q, want %q", command, got, want) + } + } +} + func TestParseLsofCWDs(t *testing.T) { out := []byte("p123\nn/Users/moshe/project-a\np456\nn/Users/moshe/project-b\n") cwds := parseLsofCWDs(out) diff --git a/pkg/cli/session_ps.go b/pkg/cli/session_ps.go new file mode 100644 index 0000000..b05b5e8 --- /dev/null +++ b/pkg/cli/session_ps.go @@ -0,0 +1,436 @@ +package cli + +import ( + "context" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "time" + + "github.com/flanksource/captain/pkg/claude" + "github.com/flanksource/captain/pkg/cmux" + "github.com/flanksource/clicky/api" + rpchttp "github.com/flanksource/clicky/rpc/http" +) + +type PSOptions struct { + Source string `flag:"source" help:"Filter source: all, claude, codex" default:"all"` + All bool `flag:"all" help:"Include agents from all projects" short:"a"` + Project string `flag:"project" help:"Restrict to an explicit project path"` + Query string `flag:"q" help:"Search session id, cwd, pid, surface, or health"` +} + +type PSResult struct { + Source string `json:"source" pretty:"label=Source"` + Scope string `json:"scope" pretty:"label=Scope"` + Project string `json:"project,omitempty" pretty:"label=Project"` + Total int `json:"total" pretty:"label=Total"` + Live int `json:"live" pretty:"label=Live"` + Active int `json:"active" pretty:"label=Active"` + Alerts int `json:"alerts,omitempty" pretty:"label=Alerts"` + Tokens string `json:"tokens,omitempty" pretty:"label=Tokens"` + Cost string `json:"cost,omitempty" pretty:"label=Cost"` + Sessions []PSRow `json:"sessions"` +} + +// RunPS lists currently-active agent sessions. Unlike `sessions live` (which +// starts from on-disk history and overlays live processes), `ps` starts from the +// live process listing, resolves each process's session-id/agent-ids/CMUX +// surface/last-activity from the OS (ps + lsof + environ), then augments each +// from the session cache/DB — so it reports only sessions with a live process. +func RunPS(ctx context.Context, opts PSOptions) (PSResult, error) { + source, err := normalizeSessionSource(opts.Source) + if err != nil { + return PSResult{}, err + } + cwd, err := os.Getwd() + if err != nil { + return PSResult{}, err + } + scope, projectRoot, _ := resolveSessionScope(cwd, opts.All, opts.Project) + + stopDiscover := rpchttp.Track(ctx, "discover") + processes, err := discoverSessionProcesses() + stopDiscover() + if err != nil { + return PSResult{}, err + } + processes = filterProcessesBySource(processes, source) + if projectRoot != "" { + processes = filterAgentProcessesByProject(processes, projectRoot) + } + + stopInspect := rpchttp.Track(ctx, "inspect") + enrichProcessesFromOS(processes) + enrichSurfacesFromCmux(processes) + stopInspect() + + stopAugment := rpchttp.Track(ctx, "augment") + records := make([]SessionRecord, 0, len(processes)) + for i := range processes { + record := psRecord(ctx, processes[i]) + if sessionMatchesQuery(record, opts.Query) { + records = append(records, record) + } + } + stopAugment() + + sortPSRecords(records) + summary := summarizeSessionDashboard(records) + rows := make([]PSRow, len(records)) + for i, record := range records { + rows[i] = PSRow{record} + } + + return PSResult{ + Source: source, + Scope: scope, + Project: projectResultValue(scope, projectRoot), + Total: len(records), + Live: summary.LiveSessions, + Active: summary.ActiveSessions, + Alerts: summary.AlertSessions, + Tokens: psSummaryTokens(summary), + Cost: psSummaryCost(summary), + Sessions: rows, + }, nil +} + +func psSummaryTokens(summary SessionDashboardWire) string { + if summary.TotalTokens <= 0 { + return "" + } + return api.HumanNumber(int64(summary.TotalTokens)).String() +} + +func psSummaryCost(summary SessionDashboardWire) string { + if summary.CostUSD <= 0 { + return "" + } + return fmt.Sprintf("$%.2f", summary.CostUSD) +} + +func filterProcessesBySource(processes []agentProcess, source string) []agentProcess { + if source == "all" { + return processes + } + filtered := make([]agentProcess, 0, len(processes)) + for _, proc := range processes { + if proc.Source == source { + filtered = append(filtered, proc) + } + } + return filtered +} + +// Indirection points so RunPS can be tested without live processes. +var ( + discoverProcessSurface = processSurface + discoverOpenSessionFiles = processOpenSessionFiles + discoverCmuxSurfaces = cmux.Surfaces +) + +// enrichProcessesFromOS resolves each process's CMUX surface (from its +// environment) and session/agent identity + last activity (from the transcript +// files it holds open via a single lsof pass). +func enrichProcessesFromOS(processes []agentProcess) { + openFiles := discoverOpenSessionFiles(processIDs(processes)) + for i := range processes { + processes[i].Surface = discoverProcessSurface(processes[i].PID) + enrichProcessFromOpenFiles(&processes[i], openFiles[processes[i].PID]) + } +} + +// enrichSurfacesFromCmux joins each process's CMUX surface id to cmux's own tree +// to attach the authoritative surface title and workspace name. Best-effort: +// when cmux is not running the processes keep their env-derived surface only. +func enrichSurfacesFromCmux(processes []agentProcess) { + surfaces, err := discoverCmuxSurfaces() + if err != nil || len(surfaces) == 0 { + return + } + for i := range processes { + s := processes[i].Surface + if s == nil || s.SurfaceID == "" { + continue + } + enrichCmuxSurface(s, surfaces) + } +} + +func enrichCmuxSurface(surface *CmuxSurface, surfaces map[string]cmux.Surface) { + if info, ok := surfaces[surface.SurfaceID]; ok { + surface.SurfaceRef = info.Ref + surface.Title = info.Title + surface.Workspace = info.Workspace + } +} + +type openTranscript struct { + id string + kind string + path string + mod time.Time +} + +// enrichProcessFromOpenFiles fills SessionID/AgentIDs/SessionFile/LastActivity +// from the process's open transcripts. The primary session is the claude root +// (or, for codex, the most-recently-written rollout); every other open +// transcript of the same source is a sub-agent. Last activity is the newest +// write across all of them. +// +// Claude holds foreign transcript fds open (inherited across launches) and often +// does not hold its own current transcript open, so for claude the open set is +// restricted to the process's own project and, when that leaves nothing, the +// session is resolved by cwd instead (mirroring `sessions live`). +func enrichProcessFromOpenFiles(p *agentProcess, paths []string) { + opens := collectOpenTranscripts(p.Source, paths) + if p.Source == "claude" { + resolveClaudeSession(p, filterClaudeOwnTranscripts(opens, p.CWD)) + return + } + if len(opens) == 0 { + return + } + selectPrimaryTranscript(p, opens) +} + +// resolveClaudeSession sets the primary session by priority — (1) the explicit +// argv "--session-id" when its transcript exists, (2) the process's own-cwd open +// transcript, (3) the newest transcript under the cwd — then records any own-cwd +// sub-agent transcripts. The argv id leads because it is set explicitly by the +// launcher, whereas open fds are unreliable (claude inherits foreign transcript +// fds — the source of the original mis-attribution). +func resolveClaudeSession(p *agentProcess, ownOpens []openTranscript) { + if id := parseClaudeSessionIDFromCommand(p.Command); id != "" { + if path := locateClaudeTranscript(id, p.CWD); path != "" { + setPrimaryFromFile(p, id, path) + } + } + if p.SessionFile == "" && len(ownOpens) > 0 { + idx, newest := selectPrimary(ownOpens) + o := ownOpens[idx] + p.SessionID, p.SessionFile = o.id, o.path + la := newest + p.LastActivity = &la + } + if p.SessionFile == "" { + resolveClaudeSessionByCwd(p) + } + collectClaudeAgentIDs(p, ownOpens) +} + +// locateClaudeTranscript finds the root transcript file for a session id, +// preferring one under the process's own cwd project, then any project. +func locateClaudeTranscript(id, cwd string) string { + projectsDir := claude.GetProjectsDir() + matches, err := filepath.Glob(filepath.Join(projectsDir, "*", id+".jsonl")) + if err != nil || len(matches) == 0 { + return "" + } + if cwd != "" { + suffix := strings.ToLower(claude.NormalizePath(cwd)) + for _, m := range matches { + if strings.HasSuffix(strings.ToLower(claudeProjectName(m)), suffix) { + return m + } + } + } + return matches[0] +} + +func setPrimaryFromFile(p *agentProcess, id, path string) { + p.SessionID = id + p.SessionFile = path + if info, err := os.Stat(path); err == nil { + mod := info.ModTime() + p.LastActivity = &mod + } +} + +// collectClaudeAgentIDs records every own-cwd open transcript that is not the +// primary session as a sub-agent. +func collectClaudeAgentIDs(p *agentProcess, ownOpens []openTranscript) { + seen := map[string]bool{p.SessionID: true} + for _, o := range ownOpens { + if seen[o.id] { + continue + } + seen[o.id] = true + p.AgentIDs = append(p.AgentIDs, o.id) + } +} + +func collectOpenTranscripts(source string, paths []string) []openTranscript { + var opens []openTranscript + for _, path := range paths { + src, id, kind := classifyOpenSessionFile(path) + if src != source || id == "" { + continue + } + info, err := os.Stat(path) + if err != nil { + continue + } + opens = append(opens, openTranscript{id: id, kind: kind, path: path, mod: info.ModTime()}) + } + return opens +} + +// filterClaudeOwnTranscripts keeps only transcripts under the project directory +// matching the process's cwd, dropping foreign inherited fds. +func filterClaudeOwnTranscripts(opens []openTranscript, cwd string) []openTranscript { + if cwd == "" { + return opens + } + suffix := strings.ToLower(claude.NormalizePath(cwd)) + kept := opens[:0] + for _, o := range opens { + if strings.HasSuffix(strings.ToLower(claudeProjectName(o.path)), suffix) { + kept = append(kept, o) + } + } + return kept +} + +// claudeProjectName returns the directory segment of a claude +// transcript path (~/.claude/projects//...). +func claudeProjectName(path string) string { + rel, err := filepath.Rel(claude.GetProjectsDir(), path) + if err != nil { + return "" + } + rel = filepath.ToSlash(rel) + if i := strings.IndexByte(rel, '/'); i >= 0 { + return rel[:i] + } + return rel +} + +// resolveClaudeSessionByCwd resolves a claude process that holds no own +// transcript open: it picks the newest transcript under the process's cwd +// project (like `sessions live`), falling back to the "--session-id" argv. +func resolveClaudeSessionByCwd(p *agentProcess) { + files, err := claude.FindSessionFiles(claude.GetProjectsDir(), p.CWD, false) + if err == nil { + var newestPath string + var newest time.Time + for _, f := range files { + info, err := os.Stat(f) + if err != nil { + continue + } + if newestPath == "" || info.ModTime().After(newest) { + newestPath, newest = f, info.ModTime() + } + } + if newestPath != "" { + p.SessionID = sessionIDFromFile(newestPath) + p.SessionFile = newestPath + mod := newest + p.LastActivity = &mod + return + } + } + p.SessionID = parseClaudeSessionIDFromCommand(p.Command) +} + +// selectPrimaryTranscript picks the primary session among a process's open +// transcripts and records the rest as sub-agents. +func selectPrimaryTranscript(p *agentProcess, opens []openTranscript) { + idx, newest := selectPrimary(opens) + p.SessionID = opens[idx].id + p.SessionFile = opens[idx].path + activity := newest + p.LastActivity = &activity + collectClaudeAgentIDs(p, opens) +} + +// selectPrimary returns the index of the primary transcript (claude root +// preferred, else most recent) and the newest write time across all of them. +func selectPrimary(opens []openTranscript) (int, time.Time) { + primary := 0 + newest := opens[0].mod + for i, o := range opens { + if o.mod.After(newest) { + newest = o.mod + } + if isPreferredPrimary(o, opens[primary]) { + primary = i + } + } + return primary, newest +} + +// isPreferredPrimary reports whether candidate should replace current as the +// process's primary transcript: a claude root always wins over a sub-agent, +// otherwise the more-recently-written transcript wins. +func isPreferredPrimary(candidate, current openTranscript) bool { + if candidate.kind == "root" && current.kind != "root" { + return true + } + if candidate.kind != "root" && current.kind == "root" { + return false + } + return candidate.mod.After(current.mod) +} + +// psRecord builds the session record for a live process, augmenting it with the +// database summary (tokens, cost, context, model) when its session is known, +// and falling back to a minimal synthetic record otherwise. +func psRecord(ctx context.Context, proc agentProcess) SessionRecord { + if proc.SessionID != "" { + if db, err := captainDB(ctx); err == nil { + if overview, err := db.GetSessionOverviewByIdentity(ctx, proc.SessionID); err == nil { + record := recordFromOverview(*overview) + applyLiveProcess(&record, proc) + return record + } + } + } + record := SessionRecord{ + Key: fmt.Sprintf("live-%s-%d", proc.Source, proc.PID), + ID: psRecordID(proc), + Source: proc.Source, + CWD: proc.CWD, + StartedAt: proc.StartedAt, + } + applyLiveProcess(&record, proc) + return record +} + +func psRecordID(proc agentProcess) string { + if proc.SessionID != "" { + return proc.SessionID + } + return fmt.Sprintf("pid:%d", proc.PID) +} + +func applyLiveProcess(record *SessionRecord, proc agentProcess) { + record.Live = proc.wire() + if record.CWD == "" { + record.CWD = proc.CWD + } + if record.StartedAt == nil { + record.StartedAt = proc.StartedAt + } + if record.EndedAt == nil { + record.EndedAt = proc.LastActivity + } + record.Health = deriveSessionHealth(*record) +} + +func sortPSRecords(records []SessionRecord) { + sort.Slice(records, func(i, j int) bool { + return psActivity(records[i]).After(psActivity(records[j])) + }) +} + +func psActivity(record SessionRecord) time.Time { + if record.Live != nil && record.Live.LastActivity != nil { + return *record.Live.LastActivity + } + return sessionSortTime(record) +} diff --git a/pkg/cli/session_ps_render.go b/pkg/cli/session_ps_render.go new file mode 100644 index 0000000..a5ce3c7 --- /dev/null +++ b/pkg/cli/session_ps_render.go @@ -0,0 +1,221 @@ +package cli + +import ( + "fmt" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/flanksource/clicky/api" + "github.com/flanksource/clicky/api/icons" +) + +// PSRow is a live-session table row. It embeds SessionRecord (so JSON keeps the +// SessionRecord shape via field promotion) and implements clicky's TableProvider +// so each attribute renders in its own column — one value per cell. +type PSRow struct { + SessionRecord +} + +func (r PSRow) Columns() []api.ColumnDef { + return []api.ColumnDef{ + api.Column("status").Label("").Build(), + api.Column("agent").Label("Agent").MaxWidth(26).Build(), + api.Column("title").Label("Title").MaxWidth(32).Build(), + api.Column("project").Label("Project").MaxWidth(20).Build(), + api.Column("session").Label("Session").MaxWidth(12).Build(), + api.Column("pid").Label("PID").Build(), + api.Column("cmux").Label("Cmux").MaxWidth(10).Build(), + api.Column("usage").Label("Usage").Build(), + api.Column("activity").Label("Activity").Build(), + } +} + +func (r PSRow) Row() map[string]any { + row := map[string]any{ + "status": psStatusIcon(r), + "agent": psAgentCell(r), + "title": psTitle(r), + "session": sessionListID(psSessionID(r)), + "pid": psPID(r), + } + if r.CWD != "" { + row["project"] = filepath.Base(r.CWD) + } + if r.Live != nil && r.Live.Surface != nil && r.Live.Surface.SurfaceID != "" { + row["cmux"] = shortSessionID(r.Live.Surface.SurfaceID) + } + if usage := psUsageCell(r); usage != nil { + row["usage"] = usage + } + if r.Live != nil && r.Live.LastActivity != nil { + row["activity"] = api.Human(time.Since(*r.Live.LastActivity), "text-muted") + } + return row +} + +// psAgentCell merges source, model, and sub-agent count into one cell: +// "codex gpt-5.6-sol +1". +func psAgentCell(r PSRow) api.Text { + t := psSourceText(r.Source) + if r.Model != "" { + t = t.Space().Append(r.Model, "text-muted") + } + if n := psAgentCount(r); n > 0 { + t = t.Append(fmt.Sprintf(" +%d", n), "text-blue-500") + } + return t +} + +// psUsageCell merges token total and cost into one cell: "27.4M $13.61". +// Returns nil when the session has neither (a fresh/synthetic process). +func psUsageCell(r PSRow) api.Textable { + return sessionUsageCell(r.SessionRecord) +} + +// RowDetail expands the verbose fields that don't belong in the scannable table: +// full identifiers, working dir, command, the full CMUX surface, sub-agent ids, +// context headroom, and any health signals. +func (r PSRow) RowDetail() api.Textable { + items := []api.KeyValuePair{ + api.KeyValue("Session", psSessionID(r)), + api.KeyValue("CWD", r.CWD), + } + if r.Live != nil { + items = append(items, api.KeyValue("Command", compactSessionCommand(r.Live.Command))) + if s := r.Live.Surface; s != nil { + items = append(items, + api.KeyValue("Workspace", s.Workspace), + api.KeyValue("Surface", s.SurfaceID), + api.KeyValue("Tab", s.TabID), + api.KeyValue("Panel", s.PanelID), + api.KeyValue("Port", intOrBlank(s.Port)), + api.KeyValue("Claude PID", intOrBlank(s.ClaudePID)), + api.KeyValue("Socket", s.SocketPath), + ) + } + } + if r.Context != nil && r.Context.FreePercent > 0 { + items = append(items, api.KeyValue("Context", fmt.Sprintf("%d%% free", r.Context.FreePercent))) + } + + t := api.Text{}.Add(api.DescriptionList{Items: items}) + if agents := psAgentIDs(r); len(agents) > 0 { + t = t.NewLine().Append("Sub-agents: ", "text-muted").Add(api.CompactList(agents)) + } + for _, h := range r.Health { + t = t.NewLine().Add(psHealthIcon(h.Severity)).Space().Append(h.Message, psHealthStyle(h.Severity)) + } + return t +} + +// psStatusIcon is green when the process is alive and healthy, and escalates to +// warning/error on stopped/zombie processes or health signals. +func psStatusIcon(r PSRow) api.Textable { + severity := "" + for _, h := range r.Health { + if h.Severity == "critical" { + severity = "critical" + break + } + if h.Severity == "warning" { + severity = "warning" + } + } + status := "" + if r.Live != nil { + status = strings.ToLower(r.Live.Status) + } + switch { + case severity == "critical" || status == "zombie": + return icons.Error + case severity == "warning" || status == "stopped": + return icons.Warning + case status == "active" || status == "sleeping": + return icons.Success + default: + return icons.Info + } +} + +// psTitle prefers the authoritative cmux surface title, falling back to the +// transcript slug when the process has no known cmux surface. +func psTitle(r PSRow) string { + if r.Live != nil && r.Live.Surface != nil && r.Live.Surface.Title != "" { + return r.Live.Surface.Title + } + return r.Slug +} + +func psSourceText(source string) api.Text { + switch source { + case "claude": + return api.Text{Content: "claude", Style: "text-violet-500"} + case "codex": + return api.Text{Content: "codex", Style: "text-cyan-600"} + default: + return api.Text{Content: source} + } +} + +func psHealthIcon(severity string) api.Textable { + switch severity { + case "critical": + return icons.Error + case "warning": + return icons.Warning + default: + return icons.Info + } +} + +func psHealthStyle(severity string) string { + switch severity { + case "critical": + return "text-red-500" + case "warning": + return "text-yellow-600" + default: + return "text-muted" + } +} + +// psPID renders the pid as a string so clicky doesn't humanize it (a numeric +// cell value like 33300 would otherwise render as "33.3K"). +func psPID(r PSRow) string { + if r.Live != nil && r.Live.PID > 0 { + return strconv.Itoa(r.Live.PID) + } + return "" +} + +func intOrBlank(n int) string { + if n == 0 { + return "" + } + return strconv.Itoa(n) +} + +// psSessionID prefers the live-resolved session id, falling back to the record +// id (blank for the "pid:" synthetic id of a process with no transcript). +func psSessionID(r PSRow) string { + if r.Live != nil && r.Live.SessionID != "" { + return r.Live.SessionID + } + if strings.HasPrefix(r.ID, "pid:") { + return "" + } + return r.ID +} + +func psAgentIDs(r PSRow) []string { + if r.Live == nil { + return nil + } + return r.Live.AgentIDs +} + +func psAgentCount(r PSRow) int { + return len(psAgentIDs(r)) +} diff --git a/pkg/cli/session_ps_test.go b/pkg/cli/session_ps_test.go new file mode 100644 index 0000000..bdaf069 --- /dev/null +++ b/pkg/cli/session_ps_test.go @@ -0,0 +1,271 @@ +package cli + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/flanksource/captain/pkg/claude" + "github.com/flanksource/captain/pkg/cmux" +) + +// TestRunPSStartsFromProcessesAndAugments verifies the process-first flow: a +// live claude process is enriched with its CMUX surface + open transcript +// (session id, sub-agent id, last activity) and augmented from the on-disk +// session, producing a single live record. +func TestRunPSStartsFromProcessesAndAugments(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + project := filepath.Join(home, "work", "project") + if err := os.MkdirAll(project, 0o755); err != nil { + t.Fatal(err) + } + t.Chdir(project) + + root := filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-ps.jsonl") + writeJSONL(t, root, + map[string]any{ + "type": "user", "sessionId": "sess-ps", "timestamp": "2026-06-01T10:00:00Z", "cwd": project, + "message": map[string]any{"role": "user", "content": []any{map[string]any{"type": "text", "text": "hi"}}}, + }, + map[string]any{ + "type": "assistant", "sessionId": "sess-ps", "timestamp": "2026-06-01T10:00:02Z", "cwd": project, + "message": map[string]any{"role": "assistant", "model": "claude-sonnet-4", + "content": []any{map[string]any{"type": "text", "text": "ok"}}}, + }, + ) + subagent := filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-ps", "subagents", "agent-abc123.jsonl") + writeJSONL(t, subagent, map[string]any{ + "type": "user", "sessionId": "sess-ps", "timestamp": "2026-06-01T10:00:01Z", "cwd": project, + "message": map[string]any{"role": "user", "content": []any{map[string]any{"type": "text", "text": "sub"}}}, + }) + + surface := &CmuxSurface{SurfaceID: "SURFACE-1", TabID: "TAB-1", AgentKind: "claude", Port: 9150} + started := time.Date(2026, 6, 1, 9, 59, 0, 0, time.UTC) + + restore := stubPSDiscovery(t, + func() ([]agentProcess, error) { + return []agentProcess{{ + Source: "claude", PID: 4242, Status: "active", Active: true, + CWD: project, StartedAt: &started, + Command: "claude --session-id sess-ps", + }}, nil + }, + func(pids []int) map[int][]string { return map[int][]string{4242: {root, subagent}} }, + func(pid int) *CmuxSurface { return surface }, + map[string]cmux.Surface{"SURFACE-1": {ID: "SURFACE-1", Title: "implement-captain-ps", Workspace: "gavel-claude"}}, + ) + defer restore() + + result, err := RunPS(context.Background(), PSOptions{Source: "claude"}) + if err != nil { + t.Fatalf("RunPS: %v", err) + } + if result.Total != 1 || len(result.Sessions) != 1 { + t.Fatalf("sessions = %+v", result) + } + rec := result.Sessions[0] + if rec.Live == nil { + t.Fatal("expected live process") + } + if rec.Live.PID != 4242 || rec.Live.SessionID != "sess-ps" { + t.Fatalf("live = %+v", rec.Live) + } + if rec.Live.SessionFile != root { + t.Fatalf("primary session file = %q, want %q", rec.Live.SessionFile, root) + } + if len(rec.Live.AgentIDs) != 1 || rec.Live.AgentIDs[0] != "abc123" { + t.Fatalf("agent ids = %v", rec.Live.AgentIDs) + } + if rec.Live.Surface == nil || rec.Live.Surface.SurfaceID != "SURFACE-1" { + t.Fatalf("surface = %+v", rec.Live.Surface) + } + if rec.Live.Surface.Title != "implement-captain-ps" || rec.Live.Surface.Workspace != "gavel-claude" { + t.Fatalf("cmux enrichment = %+v", rec.Live.Surface) + } + if got := psTitle(rec); got != "implement-captain-ps" { + t.Fatalf("title = %q, want cmux surface title", got) + } + if rec.Live.LastActivity == nil { + t.Fatal("expected last activity") + } + if rec.ID != "sess-ps" { + t.Fatalf("record id = %q", rec.ID) + } + if result.Live != 1 || result.Active != 1 { + t.Fatalf("summary = live %d active %d", result.Live, result.Active) + } +} + +// TestRunPSSynthesizesWhenNoTranscript covers a codex process with no resolvable +// open transcript: it still appears as a live record keyed by pid, with the CMUX +// surface attached. +func TestRunPSSynthesizesWhenNoTranscript(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Chdir(home) + + surface := &CmuxSurface{WorkspaceID: "WS-9"} + restore := stubPSDiscovery(t, + func() ([]agentProcess, error) { + return []agentProcess{{Source: "codex", PID: 777, Status: "sleeping", Active: true, Command: "codex"}}, nil + }, + func(pids []int) map[int][]string { return map[int][]string{} }, + func(pid int) *CmuxSurface { return surface }, + nil, + ) + defer restore() + + result, err := RunPS(context.Background(), PSOptions{Source: "codex", All: true}) + if err != nil { + t.Fatalf("RunPS: %v", err) + } + if result.Total != 1 { + t.Fatalf("sessions = %+v", result) + } + rec := result.Sessions[0] + if rec.ID != "pid:777" { + t.Fatalf("record id = %q, want pid:777", rec.ID) + } + if rec.Live == nil || rec.Live.Surface == nil || rec.Live.Surface.WorkspaceID != "WS-9" { + t.Fatalf("live = %+v", rec.Live) + } +} + +// TestRunPSRejectsForeignTranscript covers the claude mis-attribution fix: a +// claude process that holds only a FOREIGN project's transcript open (an +// inherited fd) must resolve to its OWN cwd's session, not the foreign one. +func TestRunPSRejectsForeignTranscript(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + projectA := filepath.Join(home, "work", "captain") + projectB := filepath.Join(home, "work", "xero-cli") + if err := os.MkdirAll(projectA, 0o755); err != nil { + t.Fatal(err) + } + t.Chdir(projectA) + + // The process's real session lives under projectA. + own := filepath.Join(home, ".claude", "projects", claude.NormalizePath(projectA), "own-session.jsonl") + writeJSONL(t, own, map[string]any{ + "type": "user", "sessionId": "own-session", "timestamp": "2026-06-01T10:00:00Z", "cwd": projectA, + "message": map[string]any{"role": "user", "content": []any{map[string]any{"type": "text", "text": "hi"}}}, + }) + // A foreign transcript the process merely holds open (inherited fd). + foreign := filepath.Join(home, ".claude", "projects", claude.NormalizePath(projectB), "foreign-session.jsonl") + writeJSONL(t, foreign, map[string]any{ + "type": "user", "sessionId": "foreign-session", "timestamp": "2026-06-01T09:00:00Z", "cwd": projectB, + "message": map[string]any{"role": "user", "content": []any{map[string]any{"type": "text", "text": "other"}}}, + }) + + restore := stubPSDiscovery(t, + func() ([]agentProcess, error) { + return []agentProcess{{Source: "claude", PID: 4243, Status: "active", Active: true, CWD: projectA, Command: "claude"}}, nil + }, + func(pids []int) map[int][]string { return map[int][]string{4243: {foreign}} }, + func(pid int) *CmuxSurface { return nil }, + nil, + ) + defer restore() + + result, err := RunPS(context.Background(), PSOptions{Source: "claude"}) + if err != nil { + t.Fatalf("RunPS: %v", err) + } + if result.Total != 1 { + t.Fatalf("sessions = %+v", result) + } + rec := result.Sessions[0] + if rec.Live == nil || rec.Live.SessionID != "own-session" { + t.Fatalf("resolved session = %+v (must reject foreign fd, use own cwd)", rec.Live) + } + if rec.Live.SessionFile != own { + t.Fatalf("session file = %q, want %q", rec.Live.SessionFile, own) + } +} + +// TestRunPSPrefersArgvSessionID covers the priority chain: a claude process +// whose argv carries --session-id must resolve to that session, even when a +// newer transcript exists under the same cwd (which cwd-newest would pick). +func TestRunPSPrefersArgvSessionID(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + project := filepath.Join(home, "work", "captain") + if err := os.MkdirAll(project, 0o755); err != nil { + t.Fatal(err) + } + t.Chdir(project) + + dir := filepath.Join(home, ".claude", "projects", claude.NormalizePath(project)) + argvSession := filepath.Join(dir, "argv-sid.jsonl") + newerSession := filepath.Join(dir, "newer-sid.jsonl") + writeJSONL(t, argvSession, map[string]any{ + "type": "user", "sessionId": "argv-sid", "timestamp": "2026-06-01T09:00:00Z", "cwd": project, + "message": map[string]any{"role": "user", "content": []any{map[string]any{"type": "text", "text": "argv"}}}, + }) + writeJSONL(t, newerSession, map[string]any{ + "type": "user", "sessionId": "newer-sid", "timestamp": "2026-06-01T11:00:00Z", "cwd": project, + "message": map[string]any{"role": "user", "content": []any{map[string]any{"type": "text", "text": "newer"}}}, + }) + // Force newer-sid to be the most-recently-modified transcript. + old := time.Date(2026, 6, 1, 9, 0, 0, 0, time.UTC) + recent := time.Date(2026, 6, 1, 11, 0, 0, 0, time.UTC) + if err := os.Chtimes(argvSession, old, old); err != nil { + t.Fatal(err) + } + if err := os.Chtimes(newerSession, recent, recent); err != nil { + t.Fatal(err) + } + + restore := stubPSDiscovery(t, + func() ([]agentProcess, error) { + return []agentProcess{{Source: "claude", PID: 4244, Status: "active", Active: true, CWD: project, + Command: "claude --session-id argv-sid"}}, nil + }, + func(pids []int) map[int][]string { return map[int][]string{} }, + func(pid int) *CmuxSurface { return nil }, + nil, + ) + defer restore() + + result, err := RunPS(context.Background(), PSOptions{Source: "claude"}) + if err != nil { + t.Fatalf("RunPS: %v", err) + } + if result.Total != 1 { + t.Fatalf("sessions = %+v", result) + } + rec := result.Sessions[0] + if rec.Live == nil || rec.Live.SessionID != "argv-sid" { + t.Fatalf("resolved = %+v, want argv-sid (argv id must beat cwd-newest)", rec.Live) + } + if rec.Live.SessionFile != argvSession { + t.Fatalf("session file = %q, want %q", rec.Live.SessionFile, argvSession) + } +} + +func stubPSDiscovery( + t *testing.T, + procs func() ([]agentProcess, error), + files func([]int) map[int][]string, + surface func(int) *CmuxSurface, + cmuxSurfaces map[string]cmux.Surface, +) func() { + t.Helper() + origProcs := discoverSessionProcesses + origFiles := discoverOpenSessionFiles + origSurface := discoverProcessSurface + origCmux := discoverCmuxSurfaces + discoverSessionProcesses = procs + discoverOpenSessionFiles = files + discoverProcessSurface = surface + discoverCmuxSurfaces = func() (map[string]cmux.Surface, error) { return cmuxSurfaces, nil } + return func() { + discoverSessionProcesses = origProcs + discoverOpenSessionFiles = origFiles + discoverProcessSurface = origSurface + discoverCmuxSurfaces = origCmux + } +} diff --git a/pkg/cli/session_record_db.go b/pkg/cli/session_record_db.go new file mode 100644 index 0000000..00fdae0 --- /dev/null +++ b/pkg/cli/session_record_db.go @@ -0,0 +1,279 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" +) + +type sessionOverviewStore interface { + ListSessionOverviewsByIdentity(context.Context, string) ([]database.SessionOverview, error) + ListThreadSessionOverviews(context.Context, uuid.UUID) ([]database.SessionOverview, error) +} + +type sessionListStore interface { + ListSessionSummaries(context.Context, database.SessionListFilter) (database.SessionListPage, error) +} + +// sessionRecordQuery narrows the DB-backed session record list. +type sessionRecordQuery struct { + Source string // "all", "claude", or "codex" + ProjectRoot string + Query string + LiveOnly bool + Limit int + Cursor string +} + +type sessionRecordPage struct { + Records []SessionRecord + Total int + NextCursor string +} + +// dbSessionRecords projects the bounded list-specific database shape onto the +// SessionRecord wire type. Transcript/detail metrics remain on the overview +// path used by session get. +func dbSessionRecords(ctx context.Context, db sessionListStore, q sessionRecordQuery) (sessionRecordPage, error) { + filter := database.SessionListFilter{ + ProjectRoot: q.ProjectRoot, + Query: q.Query, + RootsOnly: true, + LiveOnly: q.LiveOnly, + Limit: q.Limit, + Cursor: q.Cursor, + } + if q.Source != "" && q.Source != "all" { + filter.Source = q.Source + } + page, err := db.ListSessionSummaries(ctx, filter) + if err != nil { + return sessionRecordPage{}, err + } + records := make([]SessionRecord, len(page.Rows)) + for i := range page.Rows { + records[i] = recordFromOverview(overviewFromSummary(page.Rows[i])) + } + return sessionRecordPage{Records: records, Total: int(page.Total), NextCursor: page.NextCursor}, nil +} + +func dbAllSessionRecords(ctx context.Context, db sessionListStore, query sessionRecordQuery) (sessionRecordPage, error) { + query.Cursor = "" + result := sessionRecordPage{} + seen := map[string]struct{}{} + for { + page, err := dbSessionRecords(ctx, db, query) + if err != nil { + return sessionRecordPage{}, err + } + if result.Total == 0 { + result.Total = page.Total + } + result.Records = append(result.Records, page.Records...) + if page.NextCursor == "" { + return result, nil + } + if _, ok := seen[page.NextCursor]; ok { + return sessionRecordPage{}, fmt.Errorf("captain session pagination repeated cursor %q", page.NextCursor) + } + seen[page.NextCursor] = struct{}{} + query.Cursor = page.NextCursor + } +} + +// resolveOverviewsByIdentity resolves sessions by Captain UUID or +// provider-session-id prefix. +func resolveOverviewsByIdentity(ctx context.Context, db sessionOverviewStore, id string) ([]database.SessionOverview, error) { + overviews, err := db.ListSessionOverviewsByIdentity(ctx, id) + if err != nil { + return nil, err + } + if parsed, parseErr := uuid.Parse(id); parseErr == nil && len(overviews) == 1 && + overviews[0].ID == parsed && overviews[0].ParentSessionID == nil && overviews[0].RootSessionID == nil { + thread, threadErr := db.ListThreadSessionOverviews(ctx, parsed) + if threadErr != nil { + return nil, threadErr + } + if len(thread) > 1 { + return thread, nil + } + } + return overviews, nil +} + +// candidateFromOverview adapts a DB overview row to the transcript-parsing +// candidate shape used by the detail and plan readers. +func candidateFromOverview(overview database.SessionOverview) sessionCandidate { + path := stringOr(overview.HistoryFile, stringOr(overview.Path, "")) + return sessionCandidate{ + record: SessionRecord{ + ID: stringOr(overview.ProviderSessionID, overview.ID.String()), Source: overview.Source, DetailAvailable: true, + }, + path: path, + } +} + +// sessionOverviewMetadata is the monitor-owned projection stored in +// captain_sessions.metadata (see pkg/monitor sessionMetadata). +type sessionOverviewMetadata struct { + Model string `json:"model,omitempty"` + Provider string `json:"provider,omitempty"` + Files session.ChangedFiles `json:"files,omitempty"` + Approvals session.ApprovalStats `json:"approvals,omitempty"` + Plan *session.Plan `json:"plan,omitempty"` +} + +func overviewMetadata(overview database.SessionOverview) sessionOverviewMetadata { + var metadata sessionOverviewMetadata + if len(overview.Metadata) > 0 { + _ = json.Unmarshal(overview.Metadata, &metadata) + } + return metadata +} + +func overviewGitBranch(overview database.SessionOverview) string { + if len(overview.Git) == 0 { + return "" + } + var git struct { + Branch string `json:"branch"` + } + _ = json.Unmarshal(overview.Git, &git) + return git.Branch +} + +func stringOr(value *string, fallback string) string { + if value != nil && *value != "" { + return *value + } + return fallback +} + +func overviewFromSummary(summary database.SessionListSummary) database.SessionOverview { + return database.SessionOverview{ + ID: summary.ID, ProviderSessionID: summary.ProviderSessionID, Source: summary.Source, + Provider: summary.Provider, HostID: summary.HostID, ParentSessionID: summary.ParentSessionID, + RootSessionID: summary.RootSessionID, Path: summary.Path, HistoryFile: summary.HistoryFile, + Project: summary.Project, CWD: summary.CWD, Title: summary.Title, InitialPrompt: summary.InitialPrompt, + Slug: summary.Slug, CLIVersion: summary.CLIVersion, LifecycleStatus: summary.LifecycleStatus, + Git: summary.Git, Metadata: summary.Metadata, StartedAt: summary.StartedAt, EndedAt: summary.EndedAt, + LastActivityAt: summary.LastActivityAt, PID: summary.PID, ProcessStatus: summary.ProcessStatus, + ProcessCommand: summary.ProcessCommand, ProcessCWD: summary.ProcessCWD, + ProcessStartedAt: summary.ProcessStartedAt, ProcessSampledAt: summary.ProcessSampledAt, + LastHeartbeatAt: summary.LastHeartbeatAt, LeaseOwner: summary.LeaseOwner, + LeaseExpiresAt: summary.LeaseExpiresAt, ProcessActive: summary.ProcessActive, + CPUPercent: summary.CPUPercent, MemoryPercent: summary.MemoryPercent, Model: summary.Model, + Backend: summary.Backend, Effort: summary.Effort, ContextTokens: summary.ContextTokens, + ContextWindowTokens: summary.ContextWindowTokens, ContextFreePercent: summary.ContextFreePercent, + InputTokens: summary.InputTokens, OutputTokens: summary.OutputTokens, + CacheReadTokens: summary.CacheReadTokens, CacheWriteTokens: summary.CacheWriteTokens, + TotalTokens: summary.TotalTokens, CostUSD: summary.CostUSD, + MessageCount: summary.MessageCount, ToolCallCount: summary.ToolCallCount, + } +} + +// recordFromOverview projects one overview row to the SessionRecord wire shape. +func recordFromOverview(overview database.SessionOverview) SessionRecord { + metadata := overviewMetadata(overview) + path := stringOr(overview.HistoryFile, stringOr(overview.Path, "")) + id := stringOr(overview.ProviderSessionID, overview.ID.String()) + record := SessionRecord{ + Key: overview.ID.String(), + ID: id, + Source: overview.Source, + Project: stringOr(overview.Project, ""), + Slug: stringOr(overview.Slug, ""), + Title: stringOr(overview.Title, ""), + InitialPrompt: stringOr(overview.InitialPrompt, ""), + StartedAt: overview.StartedAt, + EndedAt: overview.LastActivityAt, + Model: stringOr(overview.Model, metadata.Model), + ReasoningEffort: stringOr(overview.Effort, ""), + Version: stringOr(overview.CLIVersion, ""), + GitBranch: overviewGitBranch(overview), + Provider: firstNonEmpty(overview.Provider, metadata.Provider), + Backend: stringOr(overview.Backend, ""), + LifecycleStatus: string(overview.LifecycleStatus), + CWD: stringOr(overview.CWD, ""), + ToolCalls: int(overview.ToolCallCount), + Messages: int(overview.MessageCount), + DetailAvailable: path != "" || overview.PromptRunCount > 0, + CostUSD: overview.CostUSD, + } + if record.EndedAt == nil { + record.EndedAt = overview.EndedAt + } + if overview.TotalTokens > 0 { + record.Tokens = &SessionTokensWire{ + InputTokens: int(overview.InputTokens), + OutputTokens: int(overview.OutputTokens), + CacheReadTokens: int(overview.CacheReadTokens), + CacheCreationTokens: int(overview.CacheWriteTokens), + TotalTokens: int(overview.TotalTokens), + } + } + if overview.ContextTokens != nil && *overview.ContextTokens > 0 { + context := &SessionContextWire{UsedTokens: int(*overview.ContextTokens)} + if overview.ContextWindowTokens != nil { + context.WindowTokens = int(*overview.ContextWindowTokens) + } + if overview.ContextFreePercent != nil { + context.FreePercent = *overview.ContextFreePercent + } + record.Context = context + } + if overview.ProcessActive { + record.Live = liveWireFromOverview(overview, id, path) + if record.CWD == "" { + record.CWD = stringOr(overview.ProcessCWD, "") + } + } + record.Health = deriveSessionHealth(record) + return record +} + +func liveWireFromOverview(overview database.SessionOverview, id, path string) *SessionLiveWire { + status := stringOr(overview.ProcessStatus, "") + _, active := processLiveStatus(status) + live := &SessionLiveWire{ + Status: status, + Active: active, + CWD: stringOr(overview.ProcessCWD, ""), + Command: stringOr(overview.ProcessCommand, ""), + SessionID: id, + SessionFile: path, + StartedAt: overview.ProcessStartedAt, + SampledAt: overview.ProcessSampledAt, + LastHeartbeatAt: overview.LastHeartbeatAt, + LeaseOwner: stringOr(overview.LeaseOwner, ""), + LeaseExpiresAt: overview.LeaseExpiresAt, + LastActivity: overview.LastActivityAt, + } + if overview.PID != nil { + live.PID = int(*overview.PID) + } + if overview.CPUPercent != nil { + live.CPUPercent = *overview.CPUPercent + } + if overview.MemoryPercent != nil { + live.MemoryPercent = *overview.MemoryPercent + } + return live +} + +// processLiveStatus mirrors the monitor's ps-stat classification: zombies and +// stopped processes are present but not active. +func processLiveStatus(status string) (string, bool) { + switch status { + case "zombie", "stopped", "exited": + return status, false + case "": + return status, false + default: + return status, true + } +} diff --git a/pkg/cli/session_record_db_test.go b/pkg/cli/session_record_db_test.go new file mode 100644 index 0000000..06b80ed --- /dev/null +++ b/pkg/cli/session_record_db_test.go @@ -0,0 +1,62 @@ +package cli + +import ( + "context" + "errors" + + "github.com/flanksource/captain/pkg/database" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type sessionOverviewStoreStub struct { + getErr error + listCalled bool +} + +func (s *sessionOverviewStoreStub) ListSessionOverviewsByIdentity(context.Context, string) ([]database.SessionOverview, error) { + return nil, s.getErr +} + +func (s *sessionOverviewStoreStub) ListSessionOverviews(context.Context, database.SessionOverviewFilter) ([]database.SessionOverview, error) { + s.listCalled = true + return nil, nil +} + +func (s *sessionOverviewStoreStub) ListThreadSessionOverviews(context.Context, uuid.UUID) ([]database.SessionOverview, error) { + return nil, nil +} + +var _ = Describe("session route identity", func() { + It("never scans all overviews when an identity is not found", func(ctx SpecContext) { + store := &sessionOverviewStoreStub{getErr: database.ErrSessionNotFound} + + _, err := resolveOverviewsByIdentity(ctx, store, "codex-22ea4efed82ed44e") + + Expect(errors.Is(err, database.ErrSessionNotFound)).To(BeTrue()) + Expect(store.listCalled).To(BeFalse()) + }) + + It("never scans all overviews when an identity is ambiguous", func(ctx SpecContext) { + store := &sessionOverviewStoreStub{getErr: &database.SessionConflictError{Identity: "ad4c854e"}} + + _, err := resolveOverviewsByIdentity(ctx, store, "ad4c854e") + + Expect(errors.Is(err, database.ErrSessionConflict)).To(BeTrue()) + Expect(store.listCalled).To(BeFalse()) + }) + + It("uses the Captain UUID as the list and route key", func() { + captainID := uuid.MustParse("055781c7-360a-4eb2-80be-452b3937fcfe") + providerID := "019f7c25-9adf-7901-add9-8c46693472fb" + path := "/home/acme/.codex/sessions/rollout.jsonl" + + record := recordFromOverview(database.SessionOverview{ + ID: captainID, ProviderSessionID: &providerID, Source: "codex", Path: &path, + }) + + Expect(record.Key).To(Equal(captainID.String())) + Expect(record.ID).To(Equal(providerID)) + }) +}) diff --git a/pkg/cli/session_store.go b/pkg/cli/session_store.go deleted file mode 100644 index bed2e92..0000000 --- a/pkg/cli/session_store.go +++ /dev/null @@ -1,507 +0,0 @@ -package cli - -import ( - "encoding/json" - "errors" - "fmt" - "net" - "os" - "path/filepath" - "strconv" - "strings" - "sync" - "syscall" - "time" - - "github.com/flanksource/captain/pkg/api" - "github.com/flanksource/captain/pkg/session" - commonsdb "github.com/flanksource/commons-db/db" - "gorm.io/gorm" - "gorm.io/gorm/clause" -) - -// StoredSession is one persisted transcript summary — a root session or a -// sub-agent (linked by ParentID) — invalidated by the file's ModUnix+Size. The -// rich fields are stored as Postgres jsonb. -type StoredSession struct { - Path string `gorm:"primaryKey;column:path"` - ID string `gorm:"index;column:id"` - ParentID *string `gorm:"index;column:parent_id"` - Source string `gorm:"column:source"` - IsAgent bool `gorm:"column:is_agent"` - AgentType string `gorm:"column:agent_type"` - AgentDesc string `gorm:"column:agent_desc"` - - ModUnix int64 `gorm:"column:mod_unix"` - Size int64 `gorm:"column:size"` - - Project string `gorm:"column:project"` - CWD string `gorm:"column:cwd"` - Model string `gorm:"column:model"` - - Git session.GitState `gorm:"serializer:json;type:jsonb;column:git"` - Provider ProviderInfo `gorm:"serializer:json;type:jsonb;column:provider"` - - StartedAt *time.Time `gorm:"column:started_at"` - EndedAt *time.Time `gorm:"column:ended_at"` - - Cost api.Cost `gorm:"serializer:json;type:jsonb;column:cost"` - Usage api.Usage `gorm:"serializer:json;type:jsonb;column:usage"` - Files session.ChangedFiles `gorm:"serializer:json;type:jsonb;column:files"` - Approvals session.ApprovalStats `gorm:"serializer:json;type:jsonb;column:approvals"` - - ToolCalls int `gorm:"column:tool_calls"` - MessageCount int `gorm:"column:message_count"` - ContextTokens int `gorm:"column:context_tokens"` - - Slug string `gorm:"column:slug"` - PlanPath string `gorm:"column:plan_path"` - PlanSlug string `gorm:"column:plan_slug"` - Plan *session.Plan `gorm:"serializer:json;type:jsonb;column:plan"` - - UpdatedAt time.Time -} - -func (StoredSession) TableName() string { return "captain_sessions" } - -// ProviderInfo is the provider block stored as jsonb. -type ProviderInfo struct { - Name string `json:"name,omitempty"` - Version string `json:"version,omitempty"` - ReasoningEffort string `json:"reasoningEffort,omitempty"` - Backend string `json:"backend,omitempty"` -} - -// StoredPrompt is the realized prompt for a captain-launched session, keyed by -// session id and written by the prompt-run path. -type StoredPrompt struct { - SessionID string `gorm:"primaryKey;column:session_id"` - RunID string `gorm:"column:run_id"` - Model string `gorm:"column:model"` - Backend string `gorm:"column:backend"` - Realized PromptRenderResult `gorm:"serializer:json;type:jsonb;column:realized"` - CreatedAt time.Time -} - -func (StoredPrompt) TableName() string { return "captain_session_prompts" } - -// sessionDB wraps the gorm handle; a nil *sessionDB means "no store — uncached". -type sessionDB struct{ gdb *gorm.DB } - -var ( - storeOnce sync.Once - store *sessionDB -) - -const ( - captainSessionEnvDSN = "CAPTAIN_SESSION_DB_URL" - gavelCacheEnvDSN = "GAVEL_GITHUB_CACHE_DSN" - - gavelDBModeDSN = "dsn" - gavelDBModeEmbedded = "embedded" -) - -// sessionStore returns the persistent session store, opening it once. It returns -// nil (and logs a single Warn) when the DB is unavailable, so callers degrade to -// uncached summarization. -func sessionStore() *sessionDB { - storeOnce.Do(func() { store = openSessionStore() }) - return store -} - -func openSessionStore() *sessionDB { - dsn, source, disabled, err := configuredSessionDSN() - if disabled { - return nil // explicitly disabled (tests, or users who opt out) - } - if err != nil { - log.Warnf("gavel session store unavailable: %v; falling back to captain session store", err) - } - if dsn == "" { - dir, err := sessionDBDir() - if err != nil { - // WORKAROUND(session-cache): user-approved degrade-to-uncached when the - // summary DB can't be opened; summarization still works, just uncached. - log.Warnf("session store unavailable: %v; continuing uncached", err) - return nil - } - embeddedDSN, _, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{DataDir: dir}) - if err != nil { - log.Warnf("session store unavailable: %v; continuing uncached", err) - return nil - } - dsn = embeddedDSN // shared daemon: leave running, don't call stop() - source = "captain embedded session DB" - } - log.Debugf("session store using %s", source) - gdb, _, err := commonsdb.SetupDB(dsn, "session-cache") - if err != nil { - log.Warnf("session store unavailable: %v; continuing uncached", err) - return nil - } - if err := gdb.AutoMigrate(&StoredSession{}, &StoredPrompt{}); err != nil { - log.Warnf("session store migrate failed: %v; continuing uncached", err) - return nil - } - return &sessionDB{gdb: gdb} -} - -func configuredSessionDSN() (dsn string, source string, disabled bool, err error) { - if dsn := strings.TrimSpace(os.Getenv(gavelCacheEnvDSN)); dsn != "" { - return dsn, gavelCacheEnvDSN, false, nil - } - - if dsn := os.Getenv(captainSessionEnvDSN); dsn != "" { - if dsn == "off" { - return "", "", true, nil - } - return dsn, captainSessionEnvDSN, false, nil - } - - dsn, source, err = gavelConfiguredSessionDSN() - return dsn, source, false, err -} - -// sessionDBDir is the embedded-postgres data directory (shared across processes). -func sessionDBDir() (string, error) { - cache, err := os.UserCacheDir() - if err != nil { - return "", err - } - return filepath.Join(cache, "captain", "session-db"), nil -} - -type gavelDBConfig struct { - Mode string `json:"mode"` - DSN string `json:"dsn,omitempty"` -} - -func gavelSessionDSN() (string, string, error) { - if dsn := strings.TrimSpace(os.Getenv(gavelCacheEnvDSN)); dsn != "" { - return dsn, gavelCacheEnvDSN, nil - } - return gavelConfiguredSessionDSN() -} - -func gavelConfiguredSessionDSN() (string, string, error) { - cfg, path, err := loadGavelDBConfig() - if err != nil { - return "", "", err - } - switch cfg.Mode { - case "": - return "", "", nil - case gavelDBModeDSN: - if strings.TrimSpace(cfg.DSN) == "" { - return "", "", fmt.Errorf("%s has mode=%s but empty dsn", path, gavelDBModeDSN) - } - return cfg.DSN, path, nil - case gavelDBModeEmbedded: - running, err := findRunningGavelEmbeddedPostgres() - if err != nil { - return "", "", err - } - if running != nil { - return gavelEmbeddedDSN(running.Port), path, nil - } - dataDir, err := gavelEmbeddedDataDir() - if err != nil { - return "", "", err - } - dsn, _, err := commonsdb.StartEmbedded(commonsdb.EmbeddedConfig{ - DataDir: dataDir, - Database: "gavel", - }) - if err != nil { - return "", "", err - } - return dsn, path, nil - default: - return "", "", fmt.Errorf("%s has unsupported mode %q", path, cfg.Mode) - } -} - -func loadGavelDBConfig() (gavelDBConfig, string, error) { - path, err := gavelDBConfigPath() - if err != nil { - return gavelDBConfig{}, "", err - } - b, err := os.ReadFile(path) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - return gavelDBConfig{}, path, nil - } - return gavelDBConfig{}, path, fmt.Errorf("read %s: %w", path, err) - } - var cfg gavelDBConfig - if err := json.Unmarshal(b, &cfg); err != nil { - return gavelDBConfig{}, path, fmt.Errorf("parse %s: %w", path, err) - } - return cfg, path, nil -} - -func gavelDBConfigPath() (string, error) { - dir, err := gavelStateDir() - if err != nil { - return "", err - } - return filepath.Join(dir, "db.json"), nil -} - -func gavelEmbeddedDataDir() (string, error) { - dir, err := gavelStateDir() - if err != nil { - return "", err - } - return filepath.Join(dir, "embedded-pg"), nil -} - -func gavelStateDir() (string, error) { - home, err := os.UserHomeDir() - if err != nil { - return "", fmt.Errorf("resolve home dir: %w", err) - } - dir := filepath.Join(home, ".config", "gavel") - if err := os.MkdirAll(dir, 0o755); err != nil { - return "", fmt.Errorf("create gavel state dir %s: %w", dir, err) - } - return dir, nil -} - -type runningGavelEmbeddedPostgres struct { - PID int - Port int -} - -func findRunningGavelEmbeddedPostgres() (*runningGavelEmbeddedPostgres, error) { - dataDir, err := gavelEmbeddedDataDir() - if err != nil { - return nil, err - } - pidPath := filepath.Join(dataDir, "data", "postmaster.pid") - raw, err := os.ReadFile(pidPath) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - return nil, nil - } - return nil, fmt.Errorf("read %s: %w", pidPath, err) - } - lines := strings.Split(string(raw), "\n") - const postmasterLinePort = 3 - if len(lines) <= postmasterLinePort { - return nil, fmt.Errorf("%s has %d lines, need >%d", pidPath, len(lines), postmasterLinePort) - } - pid, err := strconv.Atoi(strings.TrimSpace(lines[0])) - if err != nil || pid <= 0 { - return nil, fmt.Errorf("%s: invalid pid %q: %w", pidPath, lines[0], err) - } - port, err := strconv.Atoi(strings.TrimSpace(lines[postmasterLinePort])) - if err != nil || port <= 0 || port > 65535 { - return nil, fmt.Errorf("%s: invalid port %q: %w", pidPath, lines[postmasterLinePort], err) - } - if !processAlive(pid) || !tcpPortReachable("localhost", port) { - return nil, nil - } - return &runningGavelEmbeddedPostgres{PID: pid, Port: port}, nil -} - -func gavelEmbeddedDSN(port int) string { - return fmt.Sprintf("postgres://postgres:postgres@localhost:%d/gavel?sslmode=disable", port) -} - -func processAlive(pid int) bool { - if pid <= 0 { - return false - } - p, err := os.FindProcess(pid) - if err != nil { - return false - } - return p.Signal(syscall.Signal(0)) == nil -} - -func tcpPortReachable(host string, port int) bool { - conn, err := net.DialTimeout("tcp", net.JoinHostPort(host, strconv.Itoa(port)), 200*time.Millisecond) - if err != nil { - return false - } - _ = conn.Close() - return true -} - -// lookupFresh returns the stored row for path when it exists and its mtime+size -// still match the file (a fresh cache hit). -func (s *sessionDB) lookupFresh(path string, modUnix, size int64) (*StoredSession, bool) { - var row StoredSession - if err := s.gdb.Where("path = ?", path).First(&row).Error; err != nil { - return nil, false - } - if row.ModUnix != modUnix || row.Size != size { - return nil, false - } - return &row, true -} - -// upsertRows persists each Row (one transcript), stamping the file's mtime+size. -func (s *sessionDB) upsertRows(rows []session.Row) { - for _, r := range rows { - stored, ok := storedFromRow(r) - if !ok { - continue - } - if err := s.gdb.Clauses(clause.OnConflict{ - Columns: []clause.Column{{Name: "path"}}, - UpdateAll: true, - }).Create(&stored).Error; err != nil { - log.Warnf("session store upsert %s: %v", r.Path, err) - } - } -} - -// prompt returns the realized-prompt record for a session id, if any. -func (s *sessionDB) prompt(sessionID string) (*StoredPrompt, bool) { - var p StoredPrompt - if err := s.gdb.Where("session_id = ?", sessionID).First(&p).Error; err != nil { - return nil, false - } - return &p, true -} - -// upsertPrompt records the realized prompt for a launched session. -func (s *sessionDB) upsertPrompt(p StoredPrompt) { - if err := s.gdb.Clauses(clause.OnConflict{ - Columns: []clause.Column{{Name: "session_id"}}, - UpdateAll: true, - }).Create(&p).Error; err != nil { - log.Warnf("session store upsert prompt %s: %v", p.SessionID, err) - } -} - -// storedBase maps a session.Row to a StoredSession without the file identity -// (ModUnix/Size) — the fields shared by persistence and projection. -func storedBase(r session.Row) StoredSession { - var parent *string - if r.ParentID != "" { - p := r.ParentID - parent = &p - } - row := StoredSession{ - Path: r.Path, - ID: r.ID, - ParentID: parent, - Source: r.Source, - IsAgent: r.IsAgent, - AgentType: r.AgentType, - AgentDesc: r.AgentDesc, - Project: r.Project, - CWD: r.CWD, - Model: r.Model, - Git: r.Git, - Provider: ProviderInfo{Name: r.Provider, Version: r.Version, ReasoningEffort: r.ReasoningEffort}, - StartedAt: r.StartedAt, - EndedAt: r.EndedAt, - Cost: r.Cost, - Usage: r.Usage, - Files: r.Files, - Approvals: r.Approvals, - ToolCalls: r.ToolCalls, - MessageCount: r.Messages, - ContextTokens: r.ContextTokens, - } - if r.Plan != nil { - row.PlanPath = r.Plan.Path - row.PlanSlug = r.Plan.Slug - plan := *r.Plan - row.Plan = &plan - } - return row -} - -// storedFromRow maps a session.Row to a StoredSession, stamping the transcript -// file's mtime+size. Returns ok=false when the file can't be stat'd. -func storedFromRow(r session.Row) (StoredSession, bool) { - info, err := os.Stat(r.Path) - if err != nil { - return StoredSession{}, false - } - row := storedBase(r) - row.ModUnix = info.ModTime().UnixNano() - row.Size = info.Size() - return row, true -} - -const ( - claudeContextWindow = 1_000_000 - codexContextWindow = 200_000 -) - -// contextWindow returns the model context window for a source. -func contextWindow(source string) int { - if source == "codex" { - return codexContextWindow - } - return claudeContextWindow -} - -// freeContextPercent is 100 minus the used fraction of the window, clamped. -func freeContextPercent(used, window int) int { - if window <= 0 { - return 0 - } - if used < 0 { - used = 0 - } - free := 100 - int(float64(used)/float64(window)*100) - if free < 0 { - return 0 - } - if free > 100 { - return 100 - } - return free -} - -// recordFromRow projects a session.Row straight to a SessionRecord (miss path, -// where the store may be unavailable). -func recordFromRow(r session.Row) SessionRecord { - return storedBase(r).toRecord() -} - -// toRecord projects a StoredSession to the SessionRecord list/live wire shape. -func (row StoredSession) toRecord() SessionRecord { - rec := SessionRecord{ - Key: sessionRecordKey(row.Source, row.Path), - ID: row.ID, - Source: row.Source, - Model: row.Model, - ReasoningEffort: row.Provider.ReasoningEffort, - Version: row.Provider.Version, - Provider: row.Provider.Name, - GitBranch: row.Git.Branch, - CWD: row.CWD, - StartedAt: row.StartedAt, - EndedAt: row.EndedAt, - ToolCalls: row.ToolCalls, - Messages: row.MessageCount, - DetailAvailable: true, - CostUSD: row.Cost.Total(), - } - if u := row.Usage; u.TotalTokens() > 0 { - rec.Tokens = &SessionTokensWire{ - InputTokens: u.InputTokens, - OutputTokens: u.OutputTokens, - CacheReadTokens: u.CacheReadTokens, - CacheCreationTokens: u.CacheWriteTokens, - TotalTokens: u.TotalTokens(), - } - } - if row.ContextTokens > 0 { - window := contextWindow(row.Source) - rec.Context = &SessionContextWire{ - UsedTokens: row.ContextTokens, - WindowTokens: window, - FreePercent: freeContextPercent(row.ContextTokens, window), - } - } - return rec -} diff --git a/pkg/cli/session_store_test.go b/pkg/cli/session_store_test.go deleted file mode 100644 index 926697f..0000000 --- a/pkg/cli/session_store_test.go +++ /dev/null @@ -1,184 +0,0 @@ -package cli - -import ( - "os" - "path/filepath" - "testing" - - "github.com/flanksource/captain/pkg/api" - "github.com/flanksource/captain/pkg/session" - commonsdb "github.com/flanksource/commons-db/db" -) - -// TestRecordFromRow_ProjectsRichFields verifies the Row→SessionRecord projection -// carries git/provider/cost/tokens and derives the context-free percent. -func TestRecordFromRow_ProjectsRichFields(t *testing.T) { - r := session.Row{ - ID: "s1", - Source: "claude", - Model: "claude-opus-4", - Git: session.GitState{Branch: "main"}, - Provider: "anthropic", - Version: "1.2.3", - ReasoningEffort: "high", - Usage: api.Usage{InputTokens: 1000, OutputTokens: 500}, - Cost: api.Cost{InputCost: 0.015, OutputCost: 0.0375}, - ContextTokens: 940_000, // 94% of the 1M window → 6% free - ToolCalls: 3, - Messages: 2, - } - rec := recordFromRow(r) - - if rec.ID != "s1" || rec.GitBranch != "main" || rec.Provider != "anthropic" || rec.ReasoningEffort != "high" { - t.Fatalf("record meta = %+v", rec) - } - if rec.Context == nil || rec.Context.FreePercent != 6 || rec.Context.WindowTokens != 1_000_000 { - t.Fatalf("context = %+v, want 6%% free of 1M", rec.Context) - } - if rec.Tokens == nil || rec.Tokens.InputTokens != 1000 || rec.Tokens.OutputTokens != 500 { - t.Fatalf("tokens = %+v", rec.Tokens) - } - if rec.CostUSD == 0 { - t.Errorf("cost not projected") - } - if rec.ToolCalls != 3 || rec.Messages != 2 { - t.Errorf("counts = %d/%d, want 3/2", rec.ToolCalls, rec.Messages) - } -} - -func TestStoredBase_PersistsInlinePlan(t *testing.T) { - r := session.Row{ - ID: "codex-plan", - Source: "codex", - Plan: &session.Plan{ - Content: "- [x] inspect\n- [ ] test", - Explicit: true, - Events: []session.PlanEvent{{Kind: session.PlanWrite}}, - }, - } - - stored := storedBase(r) - - if stored.Plan == nil { - t.Fatal("stored plan is nil") - } - if stored.Plan.Content != r.Plan.Content || !stored.Plan.Explicit { - t.Fatalf("stored plan = %+v, want %+v", stored.Plan, r.Plan) - } -} - -func TestGavelSessionDSNPrefersEnv(t *testing.T) { - t.Setenv(gavelCacheEnvDSN, "postgres://env-dsn") - - dsn, source, err := gavelSessionDSN() - if err != nil { - t.Fatalf("gavelSessionDSN: %v", err) - } - if dsn != "postgres://env-dsn" || source != gavelCacheEnvDSN { - t.Fatalf("dsn/source = %q/%q", dsn, source) - } -} - -func TestConfiguredSessionDSNPrefersGavelEnvOverCaptainEnv(t *testing.T) { - t.Setenv(gavelCacheEnvDSN, "postgres://gavel-dsn") - t.Setenv(captainSessionEnvDSN, "postgres://captain-dsn") - - dsn, source, disabled, err := configuredSessionDSN() - if err != nil { - t.Fatalf("configuredSessionDSN: %v", err) - } - if disabled { - t.Fatal("configuredSessionDSN unexpectedly disabled") - } - if dsn != "postgres://gavel-dsn" || source != gavelCacheEnvDSN { - t.Fatalf("dsn/source = %q/%q", dsn, source) - } -} - -func TestConfiguredSessionDSNFallsBackToCaptainEnv(t *testing.T) { - t.Setenv(gavelCacheEnvDSN, "") - t.Setenv(captainSessionEnvDSN, "postgres://captain-dsn") - - dsn, source, disabled, err := configuredSessionDSN() - if err != nil { - t.Fatalf("configuredSessionDSN: %v", err) - } - if disabled { - t.Fatal("configuredSessionDSN unexpectedly disabled") - } - if dsn != "postgres://captain-dsn" || source != captainSessionEnvDSN { - t.Fatalf("dsn/source = %q/%q", dsn, source) - } -} - -func TestGavelSessionDSNFromDBConfig(t *testing.T) { - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv(gavelCacheEnvDSN, "") - dir := filepath.Join(home, ".config", "gavel") - if err := os.MkdirAll(dir, 0o755); err != nil { - t.Fatalf("mkdir: %v", err) - } - path := filepath.Join(dir, "db.json") - if err := os.WriteFile(path, []byte(`{"mode":"dsn","dsn":"postgres://configured-dsn"}`), 0o600); err != nil { - t.Fatalf("write config: %v", err) - } - - dsn, source, err := gavelSessionDSN() - if err != nil { - t.Fatalf("gavelSessionDSN: %v", err) - } - if dsn != "postgres://configured-dsn" || source != path { - t.Fatalf("dsn/source = %q/%q", dsn, source) - } -} - -// TestSessionStoreRoundTrip exercises the real gorm store: fresh miss inserts, -// unchanged hit is served from the row (parent-linked child included), a changed -// file invalidates, and a realized prompt round-trips. Gated on a Postgres DSN. -func TestSessionStoreRoundTrip(t *testing.T) { - dsn := os.Getenv("CAPTAIN_SESSION_DB_URL") - if dsn == "" || dsn == "off" { - t.Skip("set CAPTAIN_SESSION_DB_URL to a Postgres DSN to run the session store integration test") - } - gdb, _, err := commonsdb.SetupDB(dsn, "session-cache-test") - if err != nil { - t.Fatalf("open db: %v", err) - } - if err := gdb.AutoMigrate(&StoredSession{}, &StoredPrompt{}); err != nil { - t.Fatalf("migrate: %v", err) - } - gdb.Exec("TRUNCATE captain_sessions, captain_session_prompts") - st := &sessionDB{gdb: gdb} - - // A root session with one sub-agent transcript. - dir := t.TempDir() - rootPath := filepath.Join(dir, "root.jsonl") - writeJSONL(t, rootPath, map[string]any{ - "type": "assistant", "sessionId": "root", "uuid": "r1", "cwd": "/repo", "gitBranch": "main", - "message": map[string]any{"role": "assistant", "model": "claude-opus-4", - "usage": map[string]any{"input_tokens": 1000, "output_tokens": 500}, - "content": []any{map[string]any{"type": "text", "text": "hi"}}}, - }) - - rows, err := session.RowsFromFile(rootPath, "claude") - if err != nil { - t.Fatalf("rows: %v", err) - } - st.upsertRows(rows) - - info, _ := os.Stat(rootPath) - if _, ok := st.lookupFresh(rootPath, info.ModTime().UnixNano(), info.Size()); !ok { - t.Fatal("expected a fresh hit after upsert") - } - // A changed size invalidates. - if _, ok := st.lookupFresh(rootPath, info.ModTime().UnixNano(), info.Size()+1); ok { - t.Fatal("size change must invalidate the row") - } - - // Realized-prompt round-trip. - st.upsertPrompt(StoredPrompt{SessionID: "root", RunID: "run-1", Model: "claude-opus-4", Backend: "claude_cli"}) - if p, ok := st.prompt("root"); !ok || p.RunID != "run-1" { - t.Fatalf("prompt = %+v, ok=%v", p, ok) - } -} diff --git a/pkg/cli/session_throughput.go b/pkg/cli/session_throughput.go new file mode 100644 index 0000000..37da0f2 --- /dev/null +++ b/pkg/cli/session_throughput.go @@ -0,0 +1,221 @@ +package cli + +import ( + "context" + "os" + "sort" + "strings" + "time" +) + +const defaultSessionThroughputLimit = 500 + +type SessionThroughputOptions struct { + Source string + All bool + Project string + Query string + Limit int +} + +type SessionThroughputResult struct { + Groups []SessionThroughputGroup `json:"groups"` + Total int `json:"total"` + Skipped int `json:"skipped"` + Source string `json:"source"` + Scope string `json:"scope"` + Project string `json:"project,omitempty"` +} + +type SessionThroughputGroup struct { + Key string `json:"key"` + Source string `json:"source"` + Model string `json:"model"` + ReasoningEffort string `json:"reasoningEffort"` + Sessions int `json:"sessions"` + DurationSeconds float64 `json:"durationSeconds"` + InputTokens int `json:"inputTokens,omitempty"` + OutputTokens int `json:"outputTokens,omitempty"` + CacheReadTokens int `json:"cacheReadTokens,omitempty"` + CacheCreationTokens int `json:"cacheCreationTokens,omitempty"` + TotalTokens int `json:"totalTokens,omitempty"` + ContextTokens int `json:"contextTokens,omitempty"` + ContextSamples int `json:"contextSamples,omitempty"` + OutputTokensPerSecond float64 `json:"outputTokensPerSecond"` + TotalTokensPerSecond float64 `json:"totalTokensPerSecond"` + ContextTokensPerSecond float64 `json:"contextTokensPerSecond,omitempty"` + AvgContextUsedPercent float64 `json:"avgContextUsedPercent,omitempty"` + Points []SessionThroughputPoint `json:"points,omitempty"` +} + +type SessionThroughputPoint struct { + At time.Time `json:"at"` + SessionID string `json:"sessionId"` + OutputTokensPerSecond float64 `json:"outputTokensPerSecond"` + TotalTokensPerSecond float64 `json:"totalTokensPerSecond"` + ContextTokensPerSecond float64 `json:"contextTokensPerSecond,omitempty"` + ContextUsedPercent float64 `json:"contextUsedPercent,omitempty"` +} + +type sessionThroughputAccumulator struct { + group SessionThroughputGroup + contextPercentSum float64 +} + +func RunSessionThroughput(ctx context.Context, opts SessionThroughputOptions) (SessionThroughputResult, error) { + source, err := normalizeSessionSource(opts.Source) + if err != nil { + return SessionThroughputResult{}, err + } + cwd, err := os.Getwd() + if err != nil { + return SessionThroughputResult{}, err + } + limit := opts.Limit + if limit <= 0 { + limit = defaultSessionThroughputLimit + } + scope, projectRoot, _ := resolveSessionScope(cwd, opts.All, opts.Project) + + db, err := freshenSessionDB(ctx) + if err != nil { + return SessionThroughputResult{}, err + } + page, err := dbSessionRecords(ctx, db, sessionRecordQuery{ + Source: source, ProjectRoot: projectRoot, Query: opts.Query, Limit: limit, + }) + if err != nil { + return SessionThroughputResult{}, err + } + return buildSessionThroughputResult(sessionThroughputResultOptions{ + Page: page, Source: source, Scope: scope, Project: projectResultValue(scope, projectRoot), + }), nil +} + +type sessionThroughputResultOptions struct { + Page sessionRecordPage + Source string + Scope string + Project string +} + +func buildSessionThroughputResult(options sessionThroughputResultOptions) SessionThroughputResult { + groups, skipped := aggregateSessionThroughput(options.Page.Records) + return SessionThroughputResult{ + Groups: groups, Total: len(options.Page.Records), Skipped: skipped, + Source: options.Source, Scope: options.Scope, Project: options.Project, + } +} + +func aggregateSessionThroughput(records []SessionRecord) ([]SessionThroughputGroup, int) { + accs := map[string]*sessionThroughputAccumulator{} + skipped := 0 + + for _, record := range records { + duration := sessionDurationSeconds(record) + if duration <= 0 || record.Tokens == nil { + skipped++ + continue + } + totalTokens := record.Tokens.TotalTokens + if totalTokens == 0 { + totalTokens = record.Tokens.InputTokens + record.Tokens.OutputTokens + record.Tokens.CacheReadTokens + record.Tokens.CacheCreationTokens + } + if totalTokens <= 0 { + skipped++ + continue + } + + source := strings.TrimSpace(record.Source) + if source == "" { + source = "unknown" + } + model := strings.TrimSpace(record.Model) + if model == "" { + model = "unknown" + } + effort := strings.TrimSpace(record.ReasoningEffort) + if effort == "" { + effort = "default" + } + key := source + "|" + model + "|" + effort + + acc := accs[key] + if acc == nil { + acc = &sessionThroughputAccumulator{group: SessionThroughputGroup{ + Key: key, + Source: source, + Model: model, + ReasoningEffort: effort, + }} + accs[key] = acc + } + + group := &acc.group + group.Sessions++ + group.DurationSeconds += duration + group.InputTokens += record.Tokens.InputTokens + group.OutputTokens += record.Tokens.OutputTokens + group.CacheReadTokens += record.Tokens.CacheReadTokens + group.CacheCreationTokens += record.Tokens.CacheCreationTokens + group.TotalTokens += totalTokens + + point := SessionThroughputPoint{ + At: *record.EndedAt, + SessionID: record.ID, + OutputTokensPerSecond: float64(record.Tokens.OutputTokens) / duration, + TotalTokensPerSecond: float64(totalTokens) / duration, + } + if record.Context != nil && record.Context.UsedTokens > 0 { + group.ContextTokens += record.Context.UsedTokens + point.ContextTokensPerSecond = float64(record.Context.UsedTokens) / duration + if record.Context.WindowTokens > 0 { + percent := float64(record.Context.UsedTokens) / float64(record.Context.WindowTokens) * 100 + group.ContextSamples++ + acc.contextPercentSum += percent + point.ContextUsedPercent = percent + } + } + group.Points = append(group.Points, point) + } + + groups := make([]SessionThroughputGroup, 0, len(accs)) + for _, acc := range accs { + group := acc.group + if group.DurationSeconds > 0 { + group.OutputTokensPerSecond = float64(group.OutputTokens) / group.DurationSeconds + group.TotalTokensPerSecond = float64(group.TotalTokens) / group.DurationSeconds + if group.ContextTokens > 0 { + group.ContextTokensPerSecond = float64(group.ContextTokens) / group.DurationSeconds + } + } + if group.ContextSamples > 0 { + group.AvgContextUsedPercent = acc.contextPercentSum / float64(group.ContextSamples) + } + sort.Slice(group.Points, func(i, j int) bool { + return group.Points[i].At.Before(group.Points[j].At) + }) + groups = append(groups, group) + } + + sort.Slice(groups, func(i, j int) bool { + if groups[i].OutputTokensPerSecond != groups[j].OutputTokensPerSecond { + return groups[i].OutputTokensPerSecond > groups[j].OutputTokensPerSecond + } + if groups[i].Sessions != groups[j].Sessions { + return groups[i].Sessions > groups[j].Sessions + } + if groups[i].Model != groups[j].Model { + return groups[i].Model < groups[j].Model + } + return groups[i].ReasoningEffort < groups[j].ReasoningEffort + }) + return groups, skipped +} + +func sessionDurationSeconds(record SessionRecord) float64 { + if record.StartedAt == nil || record.EndedAt == nil || !record.EndedAt.After(*record.StartedAt) { + return 0 + } + return record.EndedAt.Sub(*record.StartedAt).Seconds() +} diff --git a/pkg/cli/session_throughput_test.go b/pkg/cli/session_throughput_test.go new file mode 100644 index 0000000..0b9f8e2 --- /dev/null +++ b/pkg/cli/session_throughput_test.go @@ -0,0 +1,190 @@ +package cli + +import ( + "context" + "math" + "os" + "path/filepath" + "testing" + "time" + + "github.com/flanksource/captain/pkg/claude" +) + +func TestAggregateSessionThroughputSeparatesModelEffort(t *testing.T) { + base := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + records := []SessionRecord{ + throughputRecord("a", "claude", "claude-sonnet-4-6", "low", base, 10*time.Second, 100, 50, 25, 0, 500, 1000), + throughputRecord("b", "claude", "claude-sonnet-4-6", "low", base.Add(time.Minute), 20*time.Second, 200, 100, 0, 0, 250, 1000), + throughputRecord("c", "claude", "claude-sonnet-4-6", "high", base.Add(2*time.Minute), 10*time.Second, 100, 200, 0, 0, 800, 1000), + } + + groups, skipped := aggregateSessionThroughput(records) + if skipped != 0 { + t.Fatalf("skipped = %d, want 0", skipped) + } + if len(groups) != 2 { + t.Fatalf("groups = %d, want 2", len(groups)) + } + + high := findThroughputGroup(t, groups, "claude|claude-sonnet-4-6|high") + low := findThroughputGroup(t, groups, "claude|claude-sonnet-4-6|low") + + assertFloat(t, high.OutputTokensPerSecond, 20) + assertFloat(t, high.TotalTokensPerSecond, 30) + assertFloat(t, high.ContextTokensPerSecond, 80) + assertFloat(t, high.AvgContextUsedPercent, 80) + + assertFloat(t, low.OutputTokensPerSecond, 5) + assertFloat(t, low.TotalTokensPerSecond, 475.0/30.0) + assertFloat(t, low.ContextTokensPerSecond, 25) + assertFloat(t, low.AvgContextUsedPercent, 37.5) + if low.Sessions != 2 { + t.Fatalf("low sessions = %d, want 2", low.Sessions) + } + if len(low.Points) != 2 || !low.Points[0].At.Before(low.Points[1].At) { + t.Fatalf("low points are not sorted oldest-first: %+v", low.Points) + } +} + +func TestAggregateSessionThroughputSkipsIncompleteSessions(t *testing.T) { + base := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + valid := throughputRecord("valid", "codex", "gpt-5", "", base, 10*time.Second, 10, 20, 0, 0, 0, 0) + noEnd := valid + noEnd.ID = "no-end" + noEnd.EndedAt = nil + noTokens := valid + noTokens.ID = "no-tokens" + noTokens.Tokens = nil + zeroDuration := valid + zeroDuration.ID = "zero-duration" + zeroDuration.EndedAt = zeroDuration.StartedAt + + groups, skipped := aggregateSessionThroughput([]SessionRecord{valid, noEnd, noTokens, zeroDuration}) + if skipped != 3 { + t.Fatalf("skipped = %d, want 3", skipped) + } + if len(groups) != 1 { + t.Fatalf("groups = %d, want 1", len(groups)) + } + group := groups[0] + if group.Key != "codex|gpt-5|default" { + t.Fatalf("group key = %q, want default effort key", group.Key) + } + assertFloat(t, group.OutputTokensPerSecond, 2) + assertFloat(t, group.TotalTokensPerSecond, 3) +} + +func TestRunSessionThroughputRestrictsExplicitProject(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + withTestCaptainDB(t) + project := filepath.Join(home, "work", "project") + otherProject := filepath.Join(home, "work", "other") + if err := os.MkdirAll(project, 0o755); err != nil { + t.Fatal(err) + } + if err := os.MkdirAll(otherProject, 0o755); err != nil { + t.Fatal(err) + } + t.Chdir(project) + markProjectRoot(t, project) + markProjectRoot(t, otherProject) + + writeThroughputClaudeSession(t, home, project, "sess-current", "2026-06-01T10:00:00Z", "2026-06-01T10:00:10Z") + writeThroughputClaudeSession(t, home, otherProject, "sess-other", "2026-06-01T10:01:00Z", "2026-06-01T10:01:10Z") + + result, err := RunSessionThroughput(context.Background(), SessionThroughputOptions{ + Source: "claude", + Project: otherProject, + Limit: 10, + }) + if err != nil { + t.Fatalf("RunSessionThroughput: %v", err) + } + if result.Scope != "project" || result.Project != otherProject { + t.Fatalf("scope/project = %q/%q, want project/%q", result.Scope, result.Project, otherProject) + } + if result.Total != 1 { + t.Fatalf("total = %d, want 1 (result=%+v)", result.Total, result) + } + group := findThroughputGroup(t, result.Groups, "claude|claude-sonnet-4|default") + if group.Sessions != 1 || len(group.Points) != 1 || group.Points[0].SessionID != "sess-other" { + t.Fatalf("throughput group = %+v", group) + } +} + +func writeThroughputClaudeSession(t *testing.T, home, project, id, started, ended string) { + t.Helper() + writeJSONL(t, filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), id+".jsonl"), + map[string]any{ + "type": "user", + "sessionId": id, + "timestamp": started, + "cwd": project, + "message": map[string]any{ + "role": "user", + "content": []any{map[string]any{"type": "text", "text": "measure"}}, + }, + }, + map[string]any{ + "type": "assistant", + "sessionId": id, + "timestamp": ended, + "cwd": project, + "message": map[string]any{ + "role": "assistant", + "model": "claude-sonnet-4", + "content": []any{map[string]any{"type": "text", "text": "done"}}, + "usage": map[string]any{ + "input_tokens": 100, + "output_tokens": 50, + }, + }, + }, + ) +} + +func findThroughputGroup(t *testing.T, groups []SessionThroughputGroup, key string) SessionThroughputGroup { + t.Helper() + for _, group := range groups { + if group.Key == key { + return group + } + } + t.Fatalf("group %q not found in %+v", key, groups) + return SessionThroughputGroup{} +} + +func throughputRecord(id, source, model, effort string, start time.Time, duration time.Duration, input, output, cacheRead, cacheCreate, contextUsed, contextWindow int) SessionRecord { + end := start.Add(duration) + rec := SessionRecord{ + ID: id, + Source: source, + Model: model, + ReasoningEffort: effort, + StartedAt: &start, + EndedAt: &end, + Tokens: &SessionTokensWire{ + InputTokens: input, + OutputTokens: output, + CacheReadTokens: cacheRead, + CacheCreationTokens: cacheCreate, + TotalTokens: input + output + cacheRead + cacheCreate, + }, + } + if contextUsed > 0 || contextWindow > 0 { + rec.Context = &SessionContextWire{ + UsedTokens: contextUsed, + WindowTokens: contextWindow, + } + } + return rec +} + +func assertFloat(t *testing.T, got, want float64) { + t.Helper() + if math.Abs(got-want) > 0.000001 { + t.Fatalf("got %f, want %f", got, want) + } +} diff --git a/pkg/cli/session_timing_ginkgo_test.go b/pkg/cli/session_timing_ginkgo_test.go new file mode 100644 index 0000000..e9019ab --- /dev/null +++ b/pkg/cli/session_timing_ginkgo_test.go @@ -0,0 +1,30 @@ +package cli + +import ( + "context" + + "github.com/flanksource/captain/pkg/database" + rpchttp "github.com/flanksource/clicky/rpc/http" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("session get timing", func() { + It("reports lookup, hydration, and prompt-run phases", func() { + sessionID := uuid.MustParse("055781c7-360a-4eb2-80be-452b3937fcfe") + store := &sessionGetOverviewStore{ + identity: []database.SessionOverview{{ID: sessionID, Source: "captain"}}, + } + ctx, timings := rpchttp.WithTimings(context.Background()) + + _, err := runSessionGet(ctx, store, SessionGetOptions{ID: sessionID.String()}) + + Expect(err).NotTo(HaveOccurred()) + Expect(timings.Header()).To(And( + ContainSubstring("lookup;dur="), + ContainSubstring("hydrate;dur="), + ContainSubstring("prompt_runs;dur="), + )) + }) +}) diff --git a/pkg/cli/sessions.go b/pkg/cli/sessions.go index 41850c6..b71d7d0 100644 --- a/pkg/cli/sessions.go +++ b/pkg/cli/sessions.go @@ -2,63 +2,72 @@ package cli import ( "context" - "crypto/sha256" - "encoding/hex" - "encoding/json" "fmt" "os" "path/filepath" - "sort" "strings" "time" - "github.com/flanksource/captain/pkg/ai/history" "github.com/flanksource/captain/pkg/claude" "github.com/flanksource/captain/pkg/session" - rpchttp "github.com/flanksource/clicky/rpc/http" ) type SessionListOptions struct { - Source string `flag:"source" help:"Filter source: all, claude, codex" default:"all"` - All bool `flag:"all" help:"Include sessions from all projects" short:"a"` - Query string `flag:"q" help:"Search session id, model, cwd, branch, or provider"` - Limit int `flag:"limit" help:"Maximum sessions to return; 0 means no limit" default:"100" short:"l"` + Source string `flag:"source" help:"Filter source: all, claude, codex" default:"all"` + All bool `flag:"all" help:"Include sessions from all projects" short:"a"` + Project string `flag:"project" help:"Restrict sessions to an explicit project path"` + Query string `flag:"q" help:"Search session id, model, cwd, branch, or provider"` + Limit int `flag:"limit" help:"Maximum sessions to return" default:"100" short:"l"` + Cursor string `flag:"cursor" help:"Continue after an opaque session-list cursor"` } type SessionGetOptions struct { - ID string `flag:"id" args:"true" help:"Session key or session id"` - Source string `flag:"source" help:"Restrict source: all, claude, codex" default:"all"` + ID string `flag:"id" args:"true" help:"Session id (full or unambiguous prefix)"` + Offset int `flag:"offset" help:"Skip this many transcript rows from the start"` + Limit int `flag:"limit" help:"Maximum transcript rows to return; 0 means all" default:"200" short:"l"` + Tail int `flag:"tail" help:"Return only the last N transcript rows (overrides offset/limit)"` } func (SessionGetOptions) GetName() string { return "get " } type SessionListResult struct { - Sessions []SessionRecord `json:"sessions"` - Total int `json:"total"` - Source string `json:"source"` - Scope string `json:"scope"` + Sessions []SessionRecord `json:"sessions"` + Total int `json:"total"` + Source string `json:"source"` + Scope string `json:"scope"` + Project string `json:"project,omitempty"` + NextCursor string `json:"nextCursor,omitempty"` } type SessionLiveOptions struct { - Source string `flag:"source" help:"Filter source: all, claude, codex" default:"all"` - All bool `flag:"all" help:"Include sessions from all projects" short:"a"` - Query string `flag:"q" help:"Search session id, model, cwd, branch, provider, pid, or health"` - Limit int `flag:"limit" help:"Maximum sessions to return" default:"25" short:"l"` - Full bool `flag:"full" help:"Parse all matching history exactly; ignores --limit"` + Source string `flag:"source" help:"Filter source: all, claude, codex" default:"all"` + All bool `flag:"all" help:"Include sessions from all projects" short:"a"` + Project string `flag:"project" help:"Restrict sessions to an explicit project path"` + Query string `flag:"q" help:"Search session id, model, cwd, branch, provider, pid, or health"` + Limit int `flag:"limit" help:"Maximum sessions to return" default:"25" short:"l"` + Full bool `flag:"full" help:"Parse all matching history exactly; ignores --limit"` + Cursor string `flag:"cursor" help:"Continue after an opaque session-list cursor"` } type SessionLiveResult struct { - Sessions []SessionRecord `json:"sessions"` - Total int `json:"total"` - Source string `json:"source"` - Scope string `json:"scope"` - Summary SessionDashboardWire `json:"summary"` + Sessions []SessionRecord `json:"sessions"` + Total int `json:"total"` + Source string `json:"source"` + Scope string `json:"scope"` + Project string `json:"project,omitempty"` + Summary SessionDashboardWire `json:"summary"` + Database SessionDatabaseStatusWire `json:"database"` + NextCursor string `json:"nextCursor,omitempty"` } type SessionRecord struct { Key string `json:"key"` ID string `json:"id"` Source string `json:"source"` + Project string `json:"project,omitempty"` + Slug string `json:"slug,omitempty"` + Title string `json:"title,omitempty"` + InitialPrompt string `json:"initialPrompt,omitempty"` StartedAt *time.Time `json:"startedAt,omitempty"` EndedAt *time.Time `json:"endedAt,omitempty"` Model string `json:"model,omitempty"` @@ -66,6 +75,8 @@ type SessionRecord struct { Version string `json:"version,omitempty"` GitBranch string `json:"gitBranch,omitempty"` Provider string `json:"provider,omitempty"` + Backend string `json:"backend,omitempty"` + LifecycleStatus string `json:"lifecycleStatus,omitempty"` CWD string `json:"cwd,omitempty"` ToolCalls int `json:"toolCalls"` Messages int `json:"messages"` @@ -92,14 +103,40 @@ type SessionContextWire struct { } type SessionLiveWire struct { - PID int `json:"pid,omitempty"` - Status string `json:"status,omitempty"` - Active bool `json:"active"` - CPUPercent float64 `json:"cpuPercent,omitempty"` - MemoryPercent float64 `json:"memoryPercent,omitempty"` - StartedAt *time.Time `json:"startedAt,omitempty"` - CWD string `json:"cwd,omitempty"` - Command string `json:"command,omitempty"` + PID int `json:"pid,omitempty"` + Status string `json:"status,omitempty"` + Active bool `json:"active"` + CPUPercent float64 `json:"cpuPercent,omitempty"` + MemoryPercent float64 `json:"memoryPercent,omitempty"` + StartedAt *time.Time `json:"startedAt,omitempty"` + SampledAt *time.Time `json:"sampledAt,omitempty"` + LastHeartbeatAt *time.Time `json:"lastHeartbeatAt,omitempty"` + LeaseOwner string `json:"leaseOwner,omitempty"` + LeaseExpiresAt *time.Time `json:"leaseExpiresAt,omitempty"` + CWD string `json:"cwd,omitempty"` + Command string `json:"command,omitempty"` + SessionID string `json:"sessionId,omitempty"` + AgentIDs []string `json:"agentIds,omitempty"` + LastActivity *time.Time `json:"lastActivity,omitempty"` + SessionFile string `json:"sessionFile,omitempty"` + Surface *CmuxSurface `json:"surface,omitempty"` +} + +// CmuxSurface identifies the cmux multiplexer surface hosting an agent process. +// SurfaceID/WorkspaceID/… are derived from the process's CMUX_* environment +// variables; Title/Workspace are the authoritative names joined from cmux itself. +type CmuxSurface struct { + SurfaceID string `json:"surfaceId,omitempty"` + SurfaceRef string `json:"surfaceRef,omitempty"` + WorkspaceID string `json:"workspaceId,omitempty"` + TabID string `json:"tabId,omitempty"` + PanelID string `json:"panelId,omitempty"` + Port int `json:"port,omitempty"` + AgentKind string `json:"agentKind,omitempty"` + SocketPath string `json:"socketPath,omitempty"` + ClaudePID int `json:"claudePid,omitempty"` + Title string `json:"title,omitempty"` + Workspace string `json:"workspace,omitempty"` } type SessionHealthWire struct { @@ -148,7 +185,7 @@ func (l SessionLiveWire) String() string { if l.PID == 0 && l.Status == "" && l.Command == "" { return "" } - parts := make([]string, 0, 4) + parts := make([]string, 0, 6) if l.PID > 0 { parts = append(parts, fmt.Sprintf("pid %d", l.PID)) } @@ -161,6 +198,12 @@ func (l SessionLiveWire) String() string { if l.MemoryPercent > 0 { parts = append(parts, fmt.Sprintf("%.1f%% mem", l.MemoryPercent)) } + if len(l.AgentIDs) > 0 { + parts = append(parts, fmt.Sprintf("%d agents", len(l.AgentIDs))) + } + if l.Surface != nil && l.Surface.SurfaceID != "" { + parts = append(parts, "cmux:"+shortSessionID(l.Surface.SurfaceID)) + } if command := compactSessionCommand(l.Command); command != "" { parts = append(parts, command) } @@ -215,84 +258,27 @@ func RunSessionList(ctx context.Context, opts SessionListOptions) (SessionListRe return SessionListResult{}, err } - candidates, err := discoverSessionCandidates(ctx, cwd, opts.All, source) + scope, projectRoot, _ := resolveSessionScope(cwd, opts.All, opts.Project) + db, err := freshenSessionDB(ctx) if err != nil { return SessionListResult{}, err } - records := make([]SessionRecord, 0, len(candidates)) - for _, candidate := range candidates { - if sessionMatchesQuery(candidate.record, opts.Query) { - records = append(records, candidate.record) - } - } - sortSessionRecords(records) - total := len(records) - if opts.Limit > 0 && len(records) > opts.Limit { - records = records[:opts.Limit] - } - - scope := "current" - if opts.All { - scope = "all" + page, err := dbSessionRecords(ctx, db, sessionRecordQuery{ + Source: source, ProjectRoot: projectRoot, Query: opts.Query, Limit: opts.Limit, Cursor: opts.Cursor, + }) + if err != nil { + return SessionListResult{}, err } return SessionListResult{ - Sessions: records, - Total: total, - Source: source, - Scope: scope, + Sessions: page.Records, + Total: page.Total, + Source: source, + Scope: scope, + Project: projectResultValue(scope, projectRoot), + NextCursor: page.NextCursor, }, nil } -// RunSessionGet returns the unified session model for a session id. The viewer -// consumes this model natively — one shape for both Claude and Codex, carrying -// the agent hierarchy, cost, changed files, plan events, approvals, and the -// message stream (with per-message provenance and the raw JSONL line). -func RunSessionGet(ctx context.Context, opts SessionGetOptions) (*session.Session, error) { - source, err := normalizeSessionSource(opts.Source) - if err != nil { - return nil, err - } - id := strings.TrimSpace(opts.ID) - if id == "" { - return nil, fmt.Errorf("id is required") - } - - stopFind := rpchttp.Track(ctx, "find") - candidate, found, err := findSessionCandidateByID(id, source) - stopFind() - if err != nil { - return nil, err - } else if found { - return buildUnifiedSession(ctx, candidate) - } - - candidates, err := discoverSessionCandidates(ctx, "", true, source) - if err != nil { - return nil, err - } - for _, candidate := range candidates { - if candidate.record.Key != id && candidate.record.ID != id && !strings.HasPrefix(candidate.record.ID, id) { - continue - } - return buildUnifiedSession(ctx, candidate) - } - return nil, fmt.Errorf("session %q not found", id) -} - -// buildUnifiedSession builds the pkg/session model for a resolved candidate: a -// Claude session (root + sub-agent transcripts, resolved by id) or a Codex -// session file. -func buildUnifiedSession(ctx context.Context, candidate sessionCandidate) (*session.Session, error) { - defer rpchttp.Track(ctx, "parse")() - s, err := buildSessionModel(candidate) - if err != nil { - return nil, err - } - persistSessionRows(candidate, s) - attachStoredPrompt(s) - return s, nil -} - func buildSessionModel(candidate sessionCandidate) (*session.Session, error) { switch candidate.record.Source { case "claude": @@ -328,45 +314,6 @@ func buildSessionModel(candidate sessionCandidate) (*session.Session, error) { } } -// persistSessionRows upserts the session's rows (root + sub-agents linked by -// parent) so the store's hierarchy is populated on a detailed read. -func persistSessionRows(candidate sessionCandidate, s *session.Session) { - st := sessionStore() - if st == nil { - return - } - switch candidate.record.Source { - case "claude": - parsed, err := claude.ParseSessions("", true, claude.Filter{SessionIDs: []string{s.ID}, IncludeAgents: true}) - if err != nil { - return - } - for _, ps := range parsed { - if ps.SessionID == s.ID { - st.upsertRows(session.Rows(ps)) - } - } - case "codex": - if r, ok := session.CodexRow(candidate.path); ok { - st.upsertRows([]session.Row{r}) - } - } -} - -// attachStoredPrompt attaches the realized prompt (for captain-launched -// sessions) from the store to the session model. -func attachStoredPrompt(s *session.Session) { - st := sessionStore() - if st == nil || s.ID == "" { - return - } - if p, ok := st.prompt(s.ID); ok { - if raw, err := json.Marshal(p.Realized); err == nil { - s.Prompt = raw - } - } -} - func normalizeSessionSource(source string) (string, error) { source = strings.ToLower(strings.TrimSpace(source)) switch source { @@ -379,171 +326,32 @@ func normalizeSessionSource(source string) (string, error) { } } -func discoverSessionCandidates(ctx context.Context, cwd string, searchAll bool, source string) ([]sessionCandidate, error) { - var candidates []sessionCandidate - if source == "all" || source == "claude" { - claudeSessions, err := discoverClaudeSessions(ctx, cwd, searchAll) - if err != nil { - return nil, err - } - candidates = append(candidates, claudeSessions...) - } - if source == "all" || source == "codex" { - codexSessions, err := discoverCodexSessions(ctx, cwd, searchAll) - if err != nil { - return nil, err - } - candidates = append(candidates, codexSessions...) - } - return candidates, nil -} - -func findSessionCandidateByID(id string, source string) (sessionCandidate, bool, error) { - id = strings.TrimSpace(id) - if id == "" { - return sessionCandidate{}, false, nil - } - if source == "all" || source == "claude" { - candidate, ok, err := findClaudeSessionCandidateByID(id) - if err != nil || ok { - return candidate, ok, err - } - } - if source == "all" || source == "codex" { - candidate, ok, err := findCodexSessionCandidateByID(id) - if err != nil || ok { - return candidate, ok, err - } - } - return sessionCandidate{}, false, nil -} - -func findClaudeSessionCandidateByID(id string) (sessionCandidate, bool, error) { - if hasPathOrGlobMeta(id) { - return sessionCandidate{}, false, nil - } - projectsDir := claude.GetProjectsDir() - if _, err := os.Stat(projectsDir); os.IsNotExist(err) { - return sessionCandidate{}, false, nil - } else if err != nil { - return sessionCandidate{}, false, err - } - - files, err := filepath.Glob(filepath.Join(projectsDir, "*", id+"*.jsonl")) - if err != nil { - return sessionCandidate{}, false, err - } - sort.Strings(files) - for _, file := range files { - sessionID := sessionIDFromFile(file) - if sessionID != id && !strings.HasPrefix(sessionID, id) { - continue - } - return sessionCandidate{ - record: minimalSessionRecord("claude", file, sessionID), - path: file, - }, true, nil - } - return sessionCandidate{}, false, nil -} - -func findCodexSessionCandidateByID(id string) (sessionCandidate, bool, error) { - files, err := history.FindCodexSessionFiles() - if err != nil { - return sessionCandidate{}, false, err - } - sort.Strings(files) - for _, file := range files { - key := sessionRecordKey("codex", file) - fileID := sessionIDFromFile(file) - if key == id || fileID == id || (fileID != "" && strings.HasPrefix(fileID, id)) { - return sessionCandidate{ - record: minimalSessionRecord("codex", file, fileID), - path: file, - }, true, nil - } - meta, err := history.ReadCodexSessionMeta(file) - if err != nil || meta == nil || meta.ID == "" { - continue - } - if meta.ID == id || strings.HasPrefix(meta.ID, id) { - record := minimalSessionRecord("codex", file, meta.ID) - record.CWD = meta.CWD - record.Provider = meta.ModelProvider - record.Version = meta.CLIVersion - record.GitBranch = meta.GitBranch - record.StartedAt = meta.StartedAt - return sessionCandidate{record: record, path: file}, true, nil - } - } - return sessionCandidate{}, false, nil -} - -func minimalSessionRecord(source, file, id string) SessionRecord { - return SessionRecord{ - Key: sessionRecordKey(source, file), - ID: id, - Source: source, - DetailAvailable: true, - } -} - -func hasPathOrGlobMeta(s string) bool { - return strings.ContainsAny(s, `/\*?[`) -} - -func discoverClaudeSessions(ctx context.Context, cwd string, searchAll bool) ([]sessionCandidate, error) { - stopFind := rpchttp.Track(ctx, "find") - files, err := claude.FindSessionFiles(claude.GetProjectsDir(), cwd, searchAll) - stopFind() - if err != nil { - return nil, err +func normalizeSessionProject(project string) string { + project = strings.TrimSpace(project) + if project == "" || strings.EqualFold(project, "all") { + return "" } - refs := make([]sessionFileRef, len(files)) - for i, file := range files { - refs[i] = sessionFileRef{source: "claude", path: file} + if abs, err := filepath.Abs(project); err == nil { + project = abs } - return summarizeSessionRefs(ctx, refs), nil + return filepath.Clean(project) } -func discoverCodexSessions(ctx context.Context, cwd string, searchAll bool) ([]sessionCandidate, error) { - stopFind := rpchttp.Track(ctx, "find") - files, err := history.FindCodexSessionFiles() - if err != nil { - stopFind() - return nil, err - } - matchRoot := cwd - if cwd != "" { - projectInfo := claude.FindProjectInfo(cwd) - if projectInfo.Root != "" { - matchRoot = projectInfo.Root - } +func resolveSessionScope(cwd string, all bool, project string) (scope string, projectRoot string, searchAll bool) { + if project = normalizeSessionProject(project); project != "" { + return "project", sessionProjectRoot(project), true } - refs := make([]sessionFileRef, 0, len(files)) - for _, file := range files { - if !searchAll { - meta, err := history.ReadCodexSessionMeta(file) - if err != nil || meta == nil || !codexMetaMatchesProject(meta, matchRoot) { - continue - } - } - refs = append(refs, sessionFileRef{source: "codex", path: file}) + if all { + return "all", "", true } - stopFind() - return summarizeSessionRefs(ctx, refs), nil + return "current", sessionProjectRoot(cwd), false } -func codexMetaMatchesProject(meta *history.CodexSessionInfo, projectRoot string) bool { - if meta == nil { - return false +func projectResultValue(scope, projectRoot string) string { + if scope == "project" { + return projectRoot } - return sessionRecordMatchesProject(SessionRecord{CWD: meta.CWD}, projectRoot) -} - -func sessionRecordKey(source, path string) string { - sum := sha256.Sum256([]byte(source + "\x00" + path)) - return source + "-" + hex.EncodeToString(sum[:])[:16] + return "" } func sessionMatchesQuery(record SessionRecord, query string) bool { @@ -555,6 +363,9 @@ func sessionMatchesQuery(record SessionRecord, query string) bool { record.Key, record.ID, record.Source, + record.Project, + record.Title, + record.InitialPrompt, record.Model, record.ReasoningEffort, record.Version, @@ -574,12 +385,6 @@ func sessionMatchesQuery(record SessionRecord, query string) bool { return false } -func sortSessionRecords(records []SessionRecord) { - sort.Slice(records, func(i, j int) bool { - return sessionSortTime(records[i]).After(sessionSortTime(records[j])) - }) -} - func sessionSortTime(record SessionRecord) time.Time { if record.EndedAt != nil { return *record.EndedAt diff --git a/pkg/cli/sessions_test.go b/pkg/cli/sessions_test.go index 14989b0..5512491 100644 --- a/pkg/cli/sessions_test.go +++ b/pkg/cli/sessions_test.go @@ -10,12 +10,14 @@ import ( "time" "github.com/flanksource/captain/pkg/claude" - rpchttp "github.com/flanksource/clicky/rpc/http" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/monitor" ) func TestRunSessionListAndGetClaude(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "project") if err := os.MkdirAll(project, 0o755); err != nil { t.Fatal(err) @@ -82,7 +84,7 @@ func TestRunSessionListAndGetClaude(t *testing.T) { if session.ID != "sess-claude" || session.Source != "claude" { t.Fatalf("session summary = %+v", session) } - if session.ToolCalls != 1 || session.Messages != 1 { + if session.ToolCalls != 1 || session.Messages != 3 { t.Fatalf("ToolCalls/Messages = %d/%d", session.ToolCalls, session.Messages) } if session.Key == "" { @@ -93,10 +95,14 @@ func TestRunSessionListAndGetClaude(t *testing.T) { if err != nil { t.Fatalf("RunSessionGet: %v", err) } - if detail.Source != "claude" || detail.ID != "sess-claude" { - t.Fatalf("session detail meta = %+v", detail) + if detail.Total != 1 || len(detail.Sessions) != 1 || detail.Sessions[0].Detail == nil { + t.Fatalf("session detail result = %+v", detail) } - entries := detail.ToReplayEntries() + parsed := detail.Sessions[0].Detail + if parsed.Source != "claude" || parsed.ID != "sess-claude" { + t.Fatalf("session detail meta = %+v", parsed) + } + entries := parsed.ToReplayEntries() if len(entries) != 2 { t.Fatalf("entries = %+v", entries) } @@ -114,6 +120,7 @@ func TestRunSessionListAndGetClaude(t *testing.T) { func TestRunSessionListCodexScope(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "project") if err := os.MkdirAll(project, 0o755); err != nil { t.Fatal(err) @@ -146,45 +153,83 @@ func TestRunSessionListCodexScope(t *testing.T) { } } +func TestSessionMatchesQueryIncludesIdentity(t *testing.T) { + record := SessionRecord{ + ID: "session-1", + Source: "codex", + Project: "captain", + Title: "Improve session identity", + InitialPrompt: "Show the initial user request in the listing", + } + for _, query := range []string{"captain", "identity", "initial user request"} { + if !sessionMatchesQuery(record, query) { + t.Fatalf("query %q did not match %+v", query, record) + } + } +} + func TestRunSessionGetUnknown(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) _, err := RunSessionGet(context.Background(), SessionGetOptions{ID: "missing"}) if err == nil || !strings.Contains(err.Error(), "not found") { t.Fatalf("err = %v", err) } } -func TestFindSessionCandidateByIDClaudeFilename(t *testing.T) { +func TestRunSessionGetReturnsAllProviderPrefixMatches(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + db := withTestCaptainDB(t) project := filepath.Join(home, "work", "project") - sessionFile := filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-fast.jsonl") - if err := os.MkdirAll(filepath.Dir(sessionFile), 0o755); err != nil { + if err := os.MkdirAll(project, 0o755); err != nil { t.Fatal(err) } - if err := os.WriteFile(sessionFile, []byte("{not-json"), 0o644); err != nil { - t.Fatal(err) + t.Chdir(project) + providerID := "ad4c854e-cde6-4b99-99f3-667bf74112e3" + sessionFile := filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), providerID+".jsonl") + writeJSONL(t, sessionFile, + map[string]any{ + "type": "assistant", "sessionId": providerID, "uuid": "a1", + "timestamp": "2026-07-14T10:00:00Z", "cwd": project, + "message": map[string]any{ + "role": "assistant", "model": "claude-opus-4-5", + "content": []any{map[string]any{"type": "text", "text": "canonical transcript"}}, + }, + }, + ) + for _, input := range []database.CreateSessionInput{ + {ProviderSessionID: providerID, Source: "claude", HostID: "test-host", Path: sessionFile, Project: "flanksource", CWD: project}, + {ProviderSessionID: providerID, Source: "gavel", Provider: "cmux", HostID: "local", Project: "xero-cli"}, + {ProviderSessionID: providerID, Source: "claude", Provider: "cmux", HostID: "local", Project: "xero-cli"}, + } { + if _, err := db.CreateOrGetSession(t.Context(), input); err != nil { + t.Fatalf("create duplicate session: %v", err) + } } - candidate, ok, err := findSessionCandidateByID("sess-fast", "all") + result, err := RunSessionGet(t.Context(), SessionGetOptions{ID: "ad4c854e"}) if err != nil { - t.Fatalf("findSessionCandidateByID: %v", err) + t.Fatalf("RunSessionGet: %v", err) } - if !ok { - t.Fatal("expected candidate") + if result.Total != 3 || len(result.Sessions) != 3 { + t.Fatalf("result = %+v", result) } - if candidate.path != sessionFile { - t.Fatalf("candidate.path = %q, want %q", candidate.path, sessionFile) + if !result.Sessions[0].DetailAvailable || result.Sessions[0].Detail == nil { + t.Fatalf("first match should contain parsed transcript: %+v", result.Sessions[0]) } - if candidate.record.ID != "sess-fast" || candidate.record.Source != "claude" { - t.Fatalf("candidate.record = %+v", candidate.record) + for _, item := range result.Sessions[1:] { + if item.DetailAvailable || item.Detail != nil { + t.Fatalf("metadata-only match = %+v", item) + } } } func TestRunSessionLiveEnrichesSummaryWithProcessHealth(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "project") if err := os.MkdirAll(project, 0o755); err != nil { t.Fatal(err) @@ -222,13 +267,11 @@ func TestRunSessionLiveEnrichesSummaryWithProcessHealth(t *testing.T) { ) started := time.Date(2026, 6, 1, 9, 59, 0, 0, time.UTC) - orig := discoverSessionProcesses - discoverSessionProcesses = func() ([]agentProcess, error) { - return []agentProcess{{ + monitorDiscoverProcesses = func() ([]monitor.Process, error) { + return []monitor.Process{{ Source: "claude", PID: 12345, Status: "active", - Active: true, CPUPercent: 2.5, MemoryPercent: 1.25, StartedAt: &started, @@ -236,7 +279,7 @@ func TestRunSessionLiveEnrichesSummaryWithProcessHealth(t *testing.T) { Command: "claude", }}, nil } - t.Cleanup(func() { discoverSessionProcesses = orig }) + refreshTestSessionDB(t) result, err := RunSessionLive(context.Background(), SessionLiveOptions{Source: "claude", Limit: 10}) if err != nil { @@ -263,9 +306,63 @@ func TestRunSessionLiveEnrichesSummaryWithProcessHealth(t *testing.T) { } } +func TestRunSessionLiveExcludesEndedSessions(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + db := withTestCaptainDB(t) + project := filepath.Join(home, "work", "project") + endedCWD := filepath.Join(project, "ended") + if err := os.MkdirAll(endedCWD, 0o755); err != nil { + t.Fatal(err) + } + t.Chdir(project) + + live, err := db.CreateOrGetSession(t.Context(), database.CreateSessionInput{ + ProviderSessionID: "sess-live", Source: "claude", CWD: project, + }) + if err != nil { + t.Fatalf("create live session: %v", err) + } + ended, err := db.CreateOrGetSession(t.Context(), database.CreateSessionInput{ + ProviderSessionID: "sess-ended", Source: "claude", CWD: endedCWD, + }) + if err != nil { + t.Fatalf("create ended session: %v", err) + } + + started := time.Now().UTC().Add(-time.Minute).Truncate(time.Second) + hostID := captainHostID() + for _, process := range []database.SessionProcessInput{ + {SessionID: live.ID, HostID: hostID, BootID: "boot", PID: 12345, ProcessStartedAt: started, Status: "active", CWD: project, Source: "claude"}, + {SessionID: ended.ID, HostID: hostID, BootID: "boot", PID: 54321, ProcessStartedAt: started, Status: "active", CWD: endedCWD, Source: "claude"}, + } { + if err := db.UpsertSessionProcess(t.Context(), process); err != nil { + t.Fatalf("upsert process: %v", err) + } + } + if _, err := db.EndVanishedProcesses(t.Context(), hostID, []int64{12345}); err != nil { + t.Fatalf("end historical process: %v", err) + } + monitorDiscoverProcesses = func() ([]monitor.Process, error) { + return []monitor.Process{{ + Source: "claude", PID: 12345, Status: "active", StartedAt: &started, + CWD: project, Command: "claude", + }}, nil + } + + result, err := RunSessionLive(context.Background(), SessionLiveOptions{Source: "all", All: true, Limit: 10}) + if err != nil { + t.Fatalf("RunSessionLive: %v", err) + } + if result.Total != 1 || len(result.Sessions) != 1 || result.Sessions[0].ID != "sess-live" { + t.Fatalf("live sessions = %+v", result) + } +} + func TestRunSessionLiveScopesUnmatchedProcessesToCurrentProject(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "project") otherProject := filepath.Join(home, "work", "other") if err := os.MkdirAll(project, 0o755); err != nil { @@ -277,14 +374,13 @@ func TestRunSessionLiveScopesUnmatchedProcessesToCurrentProject(t *testing.T) { t.Chdir(project) markProjectRoot(t, project) - orig := discoverSessionProcesses - discoverSessionProcesses = func() ([]agentProcess, error) { - return []agentProcess{ - {Source: "codex", PID: 100, Status: "sleeping", Active: true, CWD: project, Command: "codex"}, - {Source: "claude", PID: 200, Status: "sleeping", Active: true, CWD: otherProject, Command: "claude"}, + monitorDiscoverProcesses = func() ([]monitor.Process, error) { + return []monitor.Process{ + {Source: "codex", PID: 100, Status: "sleeping", CWD: project, Command: "codex"}, + {Source: "claude", PID: 200, Status: "sleeping", CWD: otherProject, Command: "claude"}, }, nil } - t.Cleanup(func() { discoverSessionProcesses = orig }) + refreshTestSessionDB(t) result, err := RunSessionLive(context.Background(), SessionLiveOptions{Source: "all", Limit: 10}) if err != nil { @@ -306,42 +402,70 @@ func TestRunSessionLiveScopesUnmatchedProcessesToCurrentProject(t *testing.T) { } } -func TestRunSessionLiveRecordsPhaseTimings(t *testing.T) { +func TestRunSessionLiveRestrictsExplicitProject(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) + withTestCaptainDB(t) project := filepath.Join(home, "work", "project") + otherProject := filepath.Join(home, "work", "other") if err := os.MkdirAll(project, 0o755); err != nil { t.Fatal(err) } + if err := os.MkdirAll(otherProject, 0o755); err != nil { + t.Fatal(err) + } t.Chdir(project) + markProjectRoot(t, project) + markProjectRoot(t, otherProject) - writeJSONL(t, filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-timing.jsonl"), + writeJSONL(t, filepath.Join(home, ".claude", "projects", claude.NormalizePath(project), "sess-current.jsonl"), map[string]any{ - "type": "user", - "sessionId": "sess-timing", + "type": "assistant", + "sessionId": "sess-current", "timestamp": "2026-06-01T10:00:00Z", "cwd": project, "message": map[string]any{ - "role": "user", - "content": []any{map[string]any{"type": "text", "text": "hi"}}, + "role": "assistant", + "model": "claude-sonnet-4", + "content": []any{map[string]any{"type": "text", "text": "current"}}, }, }, ) + writeJSONL(t, filepath.Join(home, ".claude", "projects", claude.NormalizePath(otherProject), "sess-other.jsonl"), + map[string]any{ + "type": "assistant", + "sessionId": "sess-other", + "timestamp": "2026-06-01T10:00:01Z", + "cwd": otherProject, + "message": map[string]any{ + "role": "assistant", + "model": "claude-sonnet-4", + "content": []any{map[string]any{"type": "text", "text": "other"}}, + }, + }, + ) + monitorDiscoverProcesses = func() ([]monitor.Process, error) { + return []monitor.Process{ + {Source: "claude", PID: 301, Status: "sleeping", CWD: project, Command: "claude --resume sess-current"}, + {Source: "claude", PID: 302, Status: "sleeping", CWD: otherProject, Command: "claude --resume sess-other"}, + }, nil + } + refreshTestSessionDB(t) - orig := discoverSessionProcesses - discoverSessionProcesses = func() ([]agentProcess, error) { return nil, nil } - t.Cleanup(func() { discoverSessionProcesses = orig }) - - ctx, timings := rpchttp.WithTimings(context.Background()) - if _, err := RunSessionLive(ctx, SessionLiveOptions{Source: "claude", Limit: 10}); err != nil { + result, err := RunSessionLive(context.Background(), SessionLiveOptions{ + Source: "claude", + All: true, + Project: otherProject, + Limit: 10, + }) + if err != nil { t.Fatalf("RunSessionLive: %v", err) } - - header := timings.Header() - for _, phase := range []string{"find", "parse", "enrich"} { - if !strings.Contains(header, phase+";dur=") { - t.Fatalf("Header() = %q, missing phase %q", header, phase) - } + if result.Scope != "project" || result.Project != otherProject { + t.Fatalf("scope/project = %q/%q, want project/%q", result.Scope, result.Project, otherProject) + } + if result.Total != 1 || len(result.Sessions) != 1 || result.Sessions[0].ID != "sess-other" { + t.Fatalf("project-scoped sessions = %+v", result) } } @@ -358,15 +482,6 @@ func markProjectRoot(t *testing.T, dir string) { } } -func hasHealth(signals []SessionHealthWire, kind, severity string) bool { - for _, signal := range signals { - if signal.Kind == kind && signal.Severity == severity { - return true - } - } - return false -} - func writeCodexSession(t *testing.T, path, id, cwd string) { t.Helper() writeJSONL(t, path, diff --git a/pkg/cli/srt.go b/pkg/cli/srt.go index c364c17..986d030 100644 --- a/pkg/cli/srt.go +++ b/pkg/cli/srt.go @@ -272,7 +272,7 @@ func RunSRTGenerate(opts SRTGenerateOptions) (any, error) { } analysis := AnalyzeToolUseLegacy(tu, projectRoot) for _, p := range analysis.WritePaths { - addDir(writeDirs, p, "") + addDir(writeDirs, DisplayPath(p), "") } for _, d := range analysis.Domains { domains[d] = true diff --git a/pkg/cli/stdin.go b/pkg/cli/stdin.go index ee0077b..356a100 100644 --- a/pkg/cli/stdin.go +++ b/pkg/cli/stdin.go @@ -137,7 +137,7 @@ func runHistoryFromReader(data []byte, opts HistoryOptions) (any, error) { if !opts.Since.IsZero() { filter.Since = &opts.Since } - if len(opts.Categories) == 0 && opts.TextFilter == "" { + if len(opts.Categories) == 0 && opts.TextFilter == "" && !usesDefaultHiddenHistoryTools(opts) { filter.Limit = opts.Limit } toolUses := claude.FilterToolUses(parsed.ToolUses, filter) @@ -146,7 +146,7 @@ func runHistoryFromReader(data []byte, opts HistoryOptions) (any, error) { return runHistorySummary(toolUses, opts, classifier, nil) } - tl := claude.ToolUsesToTools(toolUses) + tl := collapseRepeatedTitles(claude.ToolUsesToTools(toolUses)) return runHistorySingle(tl, opts, classifier, nil) } @@ -188,6 +188,11 @@ func codexToClaudeToolUses(uses []history.ToolUse) []claude.ToolUse { Source: source, Model: cu.Model, ReasoningEffort: cu.ReasoningEffort, + InputTokens: cu.InputTokens + cu.CacheReadTokens, + OutputTokens: cu.OutputTokens, + AgentID: cu.AgentID, + AgentType: cu.AgentType, + AgentDesc: cu.AgentDesc, Response: cu.Response, } } diff --git a/pkg/cli/stdin_claude_command_test.go b/pkg/cli/stdin_claude_command_test.go new file mode 100644 index 0000000..ca56667 --- /dev/null +++ b/pkg/cli/stdin_claude_command_test.go @@ -0,0 +1,73 @@ +package cli + +import ( + "github.com/flanksource/captain/pkg/claude" + "github.com/flanksource/captain/pkg/session" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// claudeGoalTranscript is the exact three-record shape Claude writes for a +// `/goal` interaction: a goal_status attachment, the slash-command wrapper as a +// text-only user message, and its captured stdout. +const claudeGoalTranscript = `{"type":"attachment","uuid":"uuid-goal","sessionId":"s1","timestamp":"2026-07-14T12:16:09.113Z","cwd":"/repo","attachment":{"type":"goal_status","met":false,"sentinel":true,"condition":"ship the docker build"}} +{"type":"user","uuid":"uuid-cmd","sessionId":"s1","timestamp":"2026-07-14T12:16:09.200Z","cwd":"/repo","message":{"role":"user","content":"/goal\n goal\n ship the docker build"}} +{"type":"user","uuid":"uuid-out","sessionId":"s1","timestamp":"2026-07-14T12:16:09.300Z","cwd":"/repo","message":{"role":"user","content":"Goal set: ship the docker build"}}` + +func historyToolNames(uses []claude.ToolUse) []string { + names := make([]string, 0, len(uses)) + for _, u := range uses { + names = append(names, u.Tool) + } + return names +} + +var _ = Describe("Claude /goal transcript history from a reader", func() { + It("detects the input as Claude JSONL and structures the command records", func() { + res, err := parseFromReader([]byte(claudeGoalTranscript)) + Expect(err).NotTo(HaveOccurred()) + Expect(res.Format).To(Equal(claude.FormatClaudeJSONL)) + + Expect(historyToolNames(res.ToolUses)).To(Equal([]string{"GoalStatus", "ClaudeCommand", "ClaudeCommand"})) + + for _, u := range res.ToolUses { + Expect(u.Tool).NotTo(Equal("User")) + if text, ok := u.Input["text"].(string); ok { + Expect(text).NotTo(ContainSubstring("")) + Expect(text).NotTo(ContainSubstring("")) + } + } + }) + + It("renders structured GoalStatus/ClaudeCommand rows and no raw wrapper rows under a normal limit", func() { + out, err := runHistoryFromReader([]byte(claudeGoalTranscript), HistoryOptions{Limit: 12}) + Expect(err).NotTo(HaveOccurred()) + + result, ok := out.(session.HistoryResult) + Expect(ok).To(BeTrue()) + + rowTools := make([]string, 0, len(result.Results)) + var goalSummary, cmdSummary, outSummary string + for _, row := range result.Results { + rowTools = append(rowTools, row.Tool) + switch { + case row.Tool == "GoalStatus": + goalSummary = row.Summary + case row.Tool == "ClaudeCommand" && cmdSummary == "": + cmdSummary = row.Summary + case row.Tool == "ClaudeCommand": + outSummary = row.Summary + } + Expect(row.Summary).NotTo(ContainSubstring("")) + Expect(row.Summary).NotTo(ContainSubstring("")) + } + + Expect(rowTools).To(ContainElement("GoalStatus")) + Expect(rowTools).To(ContainElement("ClaudeCommand")) + Expect(rowTools).NotTo(ContainElement("User")) + + Expect(goalSummary).To(ContainSubstring("ship the docker build")) + Expect(cmdSummary).To(ContainSubstring("/goal")) + Expect(outSummary).To(ContainSubstring("Goal set: ship the docker build")) + }) +}) diff --git a/pkg/cli/stdin_test.go b/pkg/cli/stdin_test.go index c38fe4c..cb7d972 100644 --- a/pkg/cli/stdin_test.go +++ b/pkg/cli/stdin_test.go @@ -42,6 +42,55 @@ func TestParseFromReader_CodexJSONL(t *testing.T) { assert.Equal(t, "codex", result.ToolUses[0].Source) } +func TestParseFromReader_CodexWaitWithContentOutput(t *testing.T) { + data := []byte(`{"timestamp":"2026-07-13T09:31:13.359Z","type":"session_meta","payload":{"id":"019f5a68-c638-7fe2-b0d0-c1d8a3c4de55","cwd":"/repo"}} +{"timestamp":"2026-07-13T09:31:13.362Z","type":"turn_context","payload":{"turn_id":"019f5ad0-fc23-7b92-8b13-a9f50d903704","model":"gpt-5.6-sol","effort":"max"}} +{"timestamp":"2026-07-13T09:41:33.611Z","type":"response_item","payload":{"type":"function_call","name":"wait","arguments":"{\"cell_id\":\"214\",\"yield_time_ms\":20000,\"max_tokens\":5000}","call_id":"call-wait"}} +{"timestamp":"2026-07-13T09:41:41.017Z","type":"response_item","payload":{"type":"function_call_output","call_id":"call-wait","output":[{"type":"input_text","text":"Script completed\nWall time 7.4 seconds\nOutput:\n"},{"type":"input_text","text":"evaluation failed\n"},{"type":"input_text","text":"exit=1"}]}} +`) + + result, err := parseFromReader(data) + require.NoError(t, err) + assert.Equal(t, claude.FormatCodexJSONL, result.Format) + require.Len(t, result.ToolUses, 1) + use := result.ToolUses[0] + assert.Equal(t, "Wait", use.Tool) + assert.Equal(t, "214", use.Input["cell_id"]) + assert.Equal(t, float64(20000), use.Input["yield_time_ms"]) + assert.Equal(t, float64(5000), use.Input["max_tokens"]) + assert.Equal(t, "evaluation failed\nexit=1", use.Response) + assert.Equal(t, "call-wait", use.ToolUseID) + + historyResult, err := runHistoryFromReader(data, HistoryOptions{}) + require.NoError(t, err) + rows := historyResult.(session.HistoryResult) + require.Len(t, rows.Results, 1) + assert.Equal(t, "Wait", rows.Results[0].Tool) +} + +func TestParseFromReader_CodexUserShellCommand(t *testing.T) { + data := []byte(`{"timestamp":"2026-07-13T06:13:59Z","type":"session_meta","payload":{"id":"sess-shell","cwd":"/repo"}} +{"timestamp":"2026-07-13T06:14:01Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"\n\ngavel proc restart\n\n\nExit code: 1\nDuration: 2.9909 seconds\nOutput:\nsignal: killed\nKill sent but port 8088 is still bound\n\n"}]}} +`) + + result, err := parseFromReader(data) + require.NoError(t, err) + assert.Equal(t, claude.FormatCodexJSONL, result.Format) + require.Len(t, result.ToolUses, 1) + use := result.ToolUses[0] + assert.Equal(t, "UserShellCommand", use.Tool) + assert.Equal(t, "gavel proc restart", use.Input["command"]) + assert.Equal(t, 1, use.Input["exit_code"]) + assert.Contains(t, use.Input["output"], "Kill sent but port 8088 is still bound") + assert.Contains(t, use.Response, "signal: killed") + + historyResult, err := runHistoryFromReader(data, HistoryOptions{}) + require.NoError(t, err) + rows := historyResult.(session.HistoryResult) + require.Len(t, rows.Results, 1) + assert.Equal(t, "UserShellCommand", rows.Results[0].Tool) +} + func TestParseFromReader_ClaudeStreamJSON(t *testing.T) { data := []byte(`{"type":"system","subtype":"init","cwd":"/tmp","session_id":"sess-1","model":"claude-sonnet-4-20250514","tools":["Bash","Read"]} {"type":"user","message":{"role":"user","content":[{"type":"text","text":"list files"}]},"session_id":"sess-1","uuid":"msg-1"} @@ -53,18 +102,19 @@ func TestParseFromReader_ClaudeStreamJSON(t *testing.T) { require.NoError(t, err) assert.Equal(t, claude.FormatClaudeStreamJSON, result.Format) assert.Nil(t, result.CLIOut) - // Stream-json now surfaces system/init and result/* as synthetic tools - // alongside real tool_use blocks. - require.Len(t, result.ToolUses, 4) + // Stream-json now surfaces chat, system/init, and result/* rows alongside + // real tool_use blocks. + require.Len(t, result.ToolUses, 5) names := []string{ result.ToolUses[0].Tool, result.ToolUses[1].Tool, result.ToolUses[2].Tool, result.ToolUses[3].Tool, + result.ToolUses[4].Tool, } - assert.Equal(t, []string{"SessionInit", "Bash", "Read", "Result"}, names) - assert.Equal(t, "ls -la", result.ToolUses[1].Input["command"]) - assert.Equal(t, "/tmp/foo.go", result.ToolUses[2].Input["file_path"]) + assert.Equal(t, []string{"SessionInit", "User", "Bash", "Read", "Result"}, names) + assert.Equal(t, "ls -la", result.ToolUses[2].Input["command"]) + assert.Equal(t, "/tmp/foo.go", result.ToolUses[3].Input["file_path"]) } func TestParseFromReader_ClaudeCLI(t *testing.T) { @@ -271,3 +321,68 @@ func TestRunHistoryFromReader_CategoryAliases(t *testing.T) { assert.Equal(t, "plan", histResult.Results[0].Category) assert.NotContains(t, histResult.Results[0].Summary, `{"plan"`) } + +func TestRunHistoryFromReader_CodexChatRowsAndChatCategoryFilter(t *testing.T) { + data := []byte(`{"timestamp":"2026-07-08T11:19:57.028Z","type":"session_meta","payload":{"id":"sess-rollout","cwd":"/repo","cli_version":"0.143.0","model_provider":"openai"}} +{"timestamp":"2026-07-08T11:19:57.028Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1"}} +{"timestamp":"2026-07-08T11:19:58.758Z","type":"turn_context","payload":{"model":"gpt-5.5","effort":"high"}} +{"timestamp":"2026-07-08T11:19:58.760Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}} +{"timestamp":"2026-07-08T11:19:58.761Z","type":"event_msg","payload":{"type":"user_message","message":"hi"}} +{"timestamp":"2026-07-08T11:20:00.403Z","type":"event_msg","payload":{"type":"agent_message","message":"hello"}} +{"timestamp":"2026-07-08T11:20:00.404Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}} +{"timestamp":"2026-07-08T11:20:00.430Z","type":"event_msg","payload":{"type":"token_count","turn_id":"turn-1","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":4,"output_tokens":2,"total_tokens":12},"total_token_usage":{"input_tokens":10,"cached_input_tokens":4,"output_tokens":2,"total_tokens":12},"model_context_window":100}}} +{"timestamp":"2026-07-08T11:20:00.435Z","type":"event_msg","payload":{"type":"task_complete","turn_id":"turn-1","duration_ms":3519}} +`) + + result, err := runHistoryFromReader(data, HistoryOptions{}) + require.NoError(t, err) + histResult := result.(session.HistoryResult) + require.Len(t, histResult.Results, 4) + assert.Equal(t, []string{"TaskStarted", "User", "Assistant", "TaskComplete"}, []string{ + histResult.Results[0].Tool, + histResult.Results[1].Tool, + histResult.Results[2].Tool, + histResult.Results[3].Tool, + }) + for _, row := range histResult.Results { + assert.Equal(t, "chat", row.Category) + } + + filtered, err := runHistoryFromReader(data, HistoryOptions{Categories: []string{"!chat"}}) + require.NoError(t, err) + assert.Empty(t, filtered.(session.HistoryResult).Results) + + tokens, err := runHistoryFromReader(data, HistoryOptions{Tools: []string{"TokenCount"}}) + require.NoError(t, err) + tokenResult := tokens.(session.HistoryResult) + require.Len(t, tokenResult.Results, 1) + assert.Equal(t, "TokenCount", tokenResult.Results[0].Tool) +} + +func TestRunHistoryFromReader_CodexEnvelopeRendersSummary(t *testing.T) { + data := []byte(`{"timestamp":"2026-07-08T11:19:57.028Z","type":"session_meta","payload":{"id":"sess-env","cwd":"/repo","cli_version":"0.143.0","model_provider":"openai"}} +{"timestamp":"2026-07-08T11:20:00.403Z","type":"event_msg","payload":{"type":"agent_message","message":"{\"endStatus\":\"completed\",\"plan\":{\"content\":\"\",\"path\":\"/Users/moshe/.codex/plans/x.md\",\"status\":\"new\"},\"questions\":[],\"summary\":\"Authored the review-banner plan.\"}"}} +`) + + result, err := runHistoryFromReader(data, HistoryOptions{}) + require.NoError(t, err) + histResult := result.(session.HistoryResult) + require.Len(t, histResult.Results, 1) + assert.Equal(t, "Assistant", histResult.Results[0].Tool) + assert.Contains(t, histResult.Results[0].Summary, "Authored the review-banner plan.") + assert.NotContains(t, histResult.Results[0].Summary, "endStatus") + assert.NotContains(t, histResult.Results[0].Summary, `{"plan"`) +} + +func TestRunHistoryFromReader_HidesCodexTokenCountBeforeLimit(t *testing.T) { + data := []byte(`{"timestamp":"2026-07-08T11:19:57.028Z","type":"session_meta","payload":{"id":"sess-rollout","cwd":"/repo","cli_version":"0.143.0","model_provider":"openai"}} +{"timestamp":"2026-07-08T11:20:00.430Z","type":"event_msg","payload":{"type":"token_count","turn_id":"turn-1","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":4,"output_tokens":2,"total_tokens":12},"total_token_usage":{"input_tokens":10,"cached_input_tokens":4,"output_tokens":2,"total_tokens":12},"model_context_window":100}}} +{"timestamp":"2026-07-08T11:20:00.435Z","type":"event_msg","payload":{"type":"task_started","turn_id":"turn-1"}} +`) + + result, err := runHistoryFromReader(data, HistoryOptions{Limit: 1}) + require.NoError(t, err) + histResult := result.(session.HistoryResult) + require.Len(t, histResult.Results, 1) + assert.Equal(t, "TaskStarted", histResult.Results[0].Tool) +} diff --git a/pkg/cli/structured_output.go b/pkg/cli/structured_output.go new file mode 100644 index 0000000..a496d6b --- /dev/null +++ b/pkg/cli/structured_output.go @@ -0,0 +1,38 @@ +package cli + +import ( + "encoding/json" + "fmt" +) + +func structuredOutputMap(value any) (map[string]any, error) { + if value == nil { + return nil, nil + } + if output, ok := value.(map[string]any); ok { + return output, nil + } + raw, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("encode structured output: %w", err) + } + if string(raw) == "null" { + return nil, nil + } + var output map[string]any + if err := json.Unmarshal(raw, &output); err != nil { + return nil, fmt.Errorf("decode structured output object: %w", err) + } + return output, nil +} + +func structuredOutputText(text string, output map[string]any) (string, error) { + if text != "" || output == nil { + return text, nil + } + raw, err := json.Marshal(output) + if err != nil { + return "", fmt.Errorf("encode structured output text: %w", err) + } + return string(raw), nil +} diff --git a/pkg/cli/summary.go b/pkg/cli/summary.go index 50454cd..abed6da 100644 --- a/pkg/cli/summary.go +++ b/pkg/cli/summary.go @@ -77,23 +77,22 @@ func BuildSummary(toolUses []claude.ToolUse, classifier *bash.CategoryClassifier denied++ } - category := classifier.ClassifyToolWithPath(tu.Tool, tu.FilePath()) - if category == bash.CategoryOther && tu.Tool == "Bash" { - if rawCmd, ok := tu.Input["command"].(string); ok { - category = classifier.ClassifyBash(rawCmd) - } - } - categories[string(category)]++ - categoryTokens[string(category)] += tuTokens + category := classifyToolUse(tu, classifier) + categories[category]++ + categoryTokens[category] += tuTokens analysis := AnalyzeToolUseLegacy(tu, tu.ProjectRoot) + // Analysis paths are absolute; this table is read by a human, so aggregate + // under the display form. for _, p := range analysis.ReadPaths { - paths[p]++ - pathTokens[p] += tuTokens + d := DisplayPath(p) + paths[d]++ + pathTokens[d] += tuTokens } for _, p := range analysis.WritePaths { - paths[p]++ - pathTokens[p] += tuTokens + d := DisplayPath(p) + paths[d]++ + pathTokens[d] += tuTokens } for _, d := range analysis.Domains { domains[d]++ diff --git a/pkg/cli/webapp/dist/.gitkeep b/pkg/cli/webapp/dist/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/pkg/cli/webapp/dist/index.html b/pkg/cli/webapp/dist/index.html index ff4c3d1..a48eaf9 100644 --- a/pkg/cli/webapp/dist/index.html +++ b/pkg/cli/webapp/dist/index.html @@ -4,8 +4,8 @@ Captain - - + +
diff --git a/pkg/cli/webapp/package.json b/pkg/cli/webapp/package.json index 95d4b9e..50da037 100644 --- a/pkg/cli/webapp/package.json +++ b/pkg/cli/webapp/package.json @@ -6,36 +6,38 @@ "scripts": { "dev": "vite", "build": "tsc -b && vite build", - "preview": "vite preview" + "preview": "vite preview", + "test": "vitest run" }, "dependencies": { "@ai-sdk/react": "^3.0.201", "@flanksource/clicky-ui": "link:../../../../clicky-ui/packages/ui", - "@radix-ui/react-compose-refs": "^1.1.2", - "@radix-ui/react-slot": "^1.1.1", "@shikijs/langs": "^1.24.0", "@shikijs/themes": "^1.24.0", "@shikijs/transformers": "^1.24.0", "@tanstack/react-query": "^5.80.7", "ai": "^6.0.199", - "class-variance-authority": "^0.7.1", - "clsx": "^2.1.1", "marked": "^15.0.0", + "monaco-editor": "0.48.0", "react": "^18.3.1", "react-dom": "^18.3.1", "react-rnd": "^10.5.3", "shiki": "^1.24.0", "streamdown": "^2.5.0", - "tailwind-merge": "^2.6.0" + "yaml": "^2.8.3" }, "devDependencies": { + "@testing-library/jest-dom": "^6.6.3", + "@testing-library/react": "^16.1.0", "@tailwindcss/vite": "^4.2.2", "@types/react": "^19.2.17", "@types/react-dom": "^19.2.3", "@vitejs/plugin-react": "^5.0.4", + "jsdom": "^26.0.0", "tailwindcss": "^4.2.2", "typescript": "~5.9.3", - "vite": "^7.1.7" + "vite": "^7.1.7", + "vitest": "^3.2.6" }, "pnpm": { "overrides": { diff --git a/pkg/cli/webapp/pnpm-lock.yaml b/pkg/cli/webapp/pnpm-lock.yaml index 4badd60..8311846 100644 --- a/pkg/cli/webapp/pnpm-lock.yaml +++ b/pkg/cli/webapp/pnpm-lock.yaml @@ -18,12 +18,6 @@ importers: '@flanksource/clicky-ui': specifier: link:../../../../clicky-ui/packages/ui version: link:../../../../clicky-ui/packages/ui - '@radix-ui/react-compose-refs': - specifier: ^1.1.2 - version: 1.1.3(@types/react@19.2.17)(react@18.3.1) - '@radix-ui/react-slot': - specifier: ^1.1.1 - version: 1.3.0(@types/react@19.2.17)(react@18.3.1) '@shikijs/langs': specifier: ^1.24.0 version: 1.29.2 @@ -39,15 +33,12 @@ importers: ai: specifier: ^6.0.199 version: 6.0.214(zod@4.4.3) - class-variance-authority: - specifier: ^0.7.1 - version: 0.7.1 - clsx: - specifier: ^2.1.1 - version: 2.1.1 marked: specifier: ^15.0.0 version: 15.0.12 + monaco-editor: + specifier: 0.48.0 + version: 0.48.0 react: specifier: ^18.3.1 version: 18.3.1 @@ -63,13 +54,19 @@ importers: streamdown: specifier: ^2.5.0 version: 2.5.0(react-dom@18.3.1(react@18.3.1))(react@18.3.1) - tailwind-merge: - specifier: ^2.6.0 - version: 2.6.1 + yaml: + specifier: ^2.8.3 + version: 2.9.0 devDependencies: '@tailwindcss/vite': specifier: ^4.2.2 - version: 4.3.1(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)) + version: 4.3.1(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0)) + '@testing-library/jest-dom': + specifier: ^6.6.3 + version: 6.9.1 + '@testing-library/react': + specifier: ^16.1.0 + version: 16.3.2(@testing-library/dom@10.4.1)(@types/react-dom@19.2.3(@types/react@19.2.17))(@types/react@19.2.17)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) '@types/react': specifier: 19.2.17 version: 19.2.17 @@ -78,7 +75,10 @@ importers: version: 19.2.3(@types/react@19.2.17) '@vitejs/plugin-react': specifier: ^5.0.4 - version: 5.2.0(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)) + version: 5.2.0(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0)) + jsdom: + specifier: ^26.0.0 + version: 26.1.0 tailwindcss: specifier: ^4.2.2 version: 4.3.1 @@ -87,10 +87,16 @@ importers: version: 5.9.3 vite: specifier: ^7.1.7 - version: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0) + version: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0) + vitest: + specifier: ^3.2.6 + version: 3.2.7(@types/debug@4.1.13)(jiti@2.7.0)(jsdom@26.1.0)(lightningcss@1.32.0)(yaml@2.9.0) packages: + '@adobe/css-tools@4.5.0': + resolution: {integrity: sha512-6OzddxPio9UiWTCemp4N8cYLV2ZN1ncRnV1cVGtve7dhPOtRkleRyx32GQCYSwDYgaHU3USMm84tNsvKzRCa1Q==} + '@ai-sdk/gateway@3.0.139': resolution: {integrity: sha512-RFpxyh5j9g7ZMfKxhiwCcF7bGU872o3JvoiIqZbHOM4qkR4RCzeKJF+k+zHomcJYGeUabEbeMXTKNASzz4Toxw==} engines: {node: '>=18'} @@ -116,6 +122,9 @@ packages: '@antfu/install-pkg@1.1.0': resolution: {integrity: sha512-MGQsmw10ZyI+EJo45CdSER4zEb+p31LpDAFp2Z3gkSd1yqVZGi0Ebx++YTEMonJy4oChEMLsxZ64j8FH6sSqtQ==} + '@asamuzakjp/css-color@3.2.0': + resolution: {integrity: sha512-K1A6z8tS3XsmCMM86xoWdn7Fkdn9m6RSVtocUrJYIwZnFVkng/PvkEoWtOWmP+Scc6saYWHWZYbndEEXxl24jw==} + '@babel/code-frame@7.29.7': resolution: {integrity: sha512-Aup7aUOfpbAUg2ROOJN6Iw5f9DMBlzu0mIkm/malLQFN/YQgO48wCj0Kxa3sEHJvPVFg7siR+qRInwXd2qhQKw==} engines: {node: '>=6.9.0'} @@ -187,6 +196,10 @@ packages: peerDependencies: '@babel/core': ^7.0.0-0 + '@babel/runtime@7.29.7': + resolution: {integrity: sha512-Nq8OhGWiZIZGV6hLHoyAKLLcJihP/xFeBMGJoUrxTX2psI8dCifzLhZISFb+VWS3wFMRDmCGw5R+dOySCqPLhw==} + engines: {node: '>=6.9.0'} + '@babel/template@7.29.7': resolution: {integrity: sha512-puq+Gf35oI24FeN11LkoUQFqv9uwNeWpxXZi/Ji3rRIoKAzKnxRaZ+Gkj0vKS9ZCiTESfng1N9LyOyXvo+m+Gg==} engines: {node: '>=6.9.0'} @@ -205,6 +218,34 @@ packages: '@chevrotain/types@11.1.2': resolution: {integrity: sha512-U+HFai5+zmJCkK86QsaJtoITlboZHBqrVketcO2ROv865xfCMSFpELQoz1GkX5GzME8pTa+3kbKrZHQtI0gdbw==} + '@csstools/color-helpers@5.1.0': + resolution: {integrity: sha512-S11EXWJyy0Mz5SYvRmY8nJYTFFd1LCNV+7cXyAgQtOOuzb4EsgfqDufL+9esx72/eLhsRdGZwaldu/h+E4t4BA==} + engines: {node: '>=18'} + + '@csstools/css-calc@2.1.4': + resolution: {integrity: sha512-3N8oaj+0juUw/1H3YwmDDJXCgTB1gKU6Hc/bB502u9zR0q2vd786XJH9QfrKIEgFlZmhZiq6epXl4rHqhzsIgQ==} + engines: {node: '>=18'} + peerDependencies: + '@csstools/css-parser-algorithms': ^3.0.5 + '@csstools/css-tokenizer': ^3.0.4 + + '@csstools/css-color-parser@3.1.0': + resolution: {integrity: sha512-nbtKwh3a6xNVIp/VRuXV64yTKnb1IjTAEEh3irzS+HkKjAOYLTGNb9pmVNntZ8iVBHcWDA2Dof0QtPgFI1BaTA==} + engines: {node: '>=18'} + peerDependencies: + '@csstools/css-parser-algorithms': ^3.0.5 + '@csstools/css-tokenizer': ^3.0.4 + + '@csstools/css-parser-algorithms@3.0.5': + resolution: {integrity: sha512-DaDeUkXZKjdGhgYaHNJTV9pV7Y9B3b644jCLs9Upc3VeNGg6LWARAT6O+Q+/COo+2gg/bM5rhpMAtf70WqfBdQ==} + engines: {node: '>=18'} + peerDependencies: + '@csstools/css-tokenizer': ^3.0.4 + + '@csstools/css-tokenizer@3.0.4': + resolution: {integrity: sha512-Vd/9EVDiu6PPJt9yAh6roZP6El1xHrdvIVGjyBsHR0RYwNHgL7FJPyIIW4fANJNG6FtyZfvlRPpFI4ZM/lubvw==} + engines: {node: '>=18'} + '@esbuild/aix-ppc64@0.28.1': resolution: {integrity: sha512-Svl7tq8k/08+p6CXPpRjQ1fKX+1odH/BQbb48fV6fj3CWHhsoIOoY87w1oHXm0qEpkIK3ZfVgp0hed3XBXzXMQ==} engines: {node: '>=18'} @@ -390,24 +431,6 @@ packages: resolution: {integrity: sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q==} engines: {node: '>=8.0.0'} - '@radix-ui/react-compose-refs@1.1.3': - resolution: {integrity: sha512-rYOP8OMnuuPMQF1uhPVlGNcCDlkokKqGFE3JcxFViIkAXP7EvFWUliJAstrapypaBLJNHbZL6jGhbVDGTwmVhA==} - peerDependencies: - '@types/react': 19.2.17 - react: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc - peerDependenciesMeta: - '@types/react': - optional: true - - '@radix-ui/react-slot@1.3.0': - resolution: {integrity: sha512-MojKku4U/miO8Av4Dkb+ctMAQx7JmY96LmtDQlAarCRtd7rN52QCSzBF+XAvr5S6coSVj9HEPBgHAHKEJVk/WA==} - peerDependencies: - '@types/react': 19.2.17 - react: ^16.8 || ^17.0 || ^18.0 || ^19.0 || ^19.0.0-rc - peerDependenciesMeta: - '@types/react': - optional: true - '@rolldown/pluginutils@1.0.0-rc.3': resolution: {integrity: sha512-eybk3TjzzzV97Dlj5c+XrBFW57eTNhzod66y9HrBlzJ6NsCrWCp/2kaPS3K9wJmurBC0Tdw4yPjXKZqlznim3Q==} @@ -678,6 +701,32 @@ packages: peerDependencies: react: ^18 || ^19 + '@testing-library/dom@10.4.1': + resolution: {integrity: sha512-o4PXJQidqJl82ckFaXUeoAW+XysPLauYI43Abki5hABd853iMhitooc6znOnczgbTYmEP6U6/y1ZyKAIsvMKGg==} + engines: {node: '>=18'} + + '@testing-library/jest-dom@6.9.1': + resolution: {integrity: sha512-zIcONa+hVtVSSep9UT3jZ5rizo2BsxgyDYU7WFD5eICBE7no3881HGeb/QkGfsJs6JTkY1aQhT7rIPC7e+0nnA==} + engines: {node: '>=14', npm: '>=6', yarn: '>=1'} + + '@testing-library/react@16.3.2': + resolution: {integrity: sha512-XU5/SytQM+ykqMnAnvB2umaJNIOsLF3PVv//1Ew4CTcpz0/BRyy/af40qqrt7SjKpDdT1saBMc42CUok5gaw+g==} + engines: {node: '>=18'} + peerDependencies: + '@testing-library/dom': ^10.0.0 + '@types/react': 19.2.17 + '@types/react-dom': 19.2.3 + react: ^18.0.0 || ^19.0.0 + react-dom: ^18.0.0 || ^19.0.0 + peerDependenciesMeta: + '@types/react': + optional: true + '@types/react-dom': + optional: true + + '@types/aria-query@5.0.4': + resolution: {integrity: sha512-rfT93uj5s0PRL7EzccGMs3brplhcrghnDoV26NqKhCAS1hVo+WdNsPvE/yb6ilfr5hi2MEk6d5EWJTKdxg8jVw==} + '@types/babel__core@7.20.5': resolution: {integrity: sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==} @@ -690,6 +739,9 @@ packages: '@types/babel__traverse@7.28.0': resolution: {integrity: sha512-8PvcXf70gTDZBgt9ptxJ8elBeBjcLOAcOtoO/mPJjtji1+CdGbHgm77om1GrsPxsiE+uXIpNSK64UYaIwQXd4Q==} + '@types/chai@5.2.3': + resolution: {integrity: sha512-Mw558oeA9fFbv65/y4mHtXDs9bPnFMZAL/jxdPFUpOHHIXX91mcgEHbS5Lahr+pwZFR8A7GQleRWeI6cGFC2UA==} + '@types/d3-array@3.2.2': resolution: {integrity: sha512-hOLWVbm7uRza0BYXpIIW5pxfrKe0W+D5lrFiAEYR+pb6w3N2SwSMaJbXdUfSEv+dT4MfHBLtn5js0LAWaO6otw==} @@ -786,6 +838,9 @@ packages: '@types/debug@4.1.13': resolution: {integrity: sha512-KSVgmQmzMwPlmtljOomayoR89W4FynCAi3E8PPs7vmDVPe84hT+vGPKkJfThkmXs0x0jAaa9U8uW8bbfyS2fWw==} + '@types/deep-eql@4.0.2': + resolution: {integrity: sha512-c9h9dVVMigMPc4bwTvC5dxqtqJZwQPePsWjPlpSOnojbor6pGqdk541lfA7AqFQr5pB1BRdq0juY9db81BwyFw==} + '@types/estree-jsx@1.0.5': resolution: {integrity: sha512-52CcUVNFyfb1A2ALocQw/Dd1BQFNmSdkuC3BkZ6iqhdMfQz7JWOFRuJFloOzjk+6WijU56m9oKXFAXc7o3Towg==} @@ -837,12 +892,64 @@ packages: peerDependencies: vite: ^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0 || ^8.0.0 + '@vitest/expect@3.2.7': + resolution: {integrity: sha512-E8eBXaKibuvH2pSZErOjdVb5vF4PbKYcrnluBTYxEk1l/VhhwZg1kZQsdtjq+CsF5CFydf2Rdkz7jDHKSisi3w==} + + '@vitest/mocker@3.2.7': + resolution: {integrity: sha512-Trr0hYO9CM3Wj6ksWHRhK9IZpIY6wTMO5u/MqXurMxT57sWBaOPEtP3Oq60ihZuh5JsiagKfz95OcxdEP6dBrA==} + peerDependencies: + msw: ^2.4.9 + vite: ^5.0.0 || ^6.0.0 || ^7.0.0-0 + peerDependenciesMeta: + msw: + optional: true + vite: + optional: true + + '@vitest/pretty-format@3.2.7': + resolution: {integrity: sha512-KUHlwqVu0sRlhCdyPdQ/wBoTfRahjUky1MubOmYw9fWfIZy1gNoHpuaaQBPAaMaVYdQYHJLurzj8ECCj5OwTqA==} + + '@vitest/runner@3.2.7': + resolution: {integrity: sha512-sB9y4ovltoQP+WaUPwmSxO9WIg9Ig694Di5PalVPsYHklAdE027mehpWF2SQSVq+k6sFgaivbTjTJwZLSHbedA==} + + '@vitest/snapshot@3.2.7': + resolution: {integrity: sha512-7C+MwShwtBSI5Buwoyg3s/iY1eHL9PKAf+O1wVh/TdnjXUtkoL/9YQtre90i4MtNXM6edP1wJ2zOBpfCyhIS7g==} + + '@vitest/spy@3.2.7': + resolution: {integrity: sha512-Q2eQGI6d2L/hBtZ0qNuKcAGid68XK6cv1xsoaIma6PaJhHPoqcEJhYpXZ/5myCMqkNgtP6UKuBhbc0nHKnrkuQ==} + + '@vitest/utils@3.2.7': + resolution: {integrity: sha512-x6BDOd7dyo3PFLY3I9/HJ25X/6OurhGXk2/B9gOZNPF7XDVjeBK4k01lQE5uvDpbuheErh91qYuE1E2OEjK3Rw==} + + agent-base@7.1.4: + resolution: {integrity: sha512-MnA+YT8fwfJPgBx3m60MNqakm30XOkyIoH1y6huTQvC0PwZG7ki8NacLBcrPbNoo8vEZy7Jpuk7+jMO+CUovTQ==} + engines: {node: '>= 14'} + ai@6.0.214: resolution: {integrity: sha512-9MlePEXT5pXtQv4fXqmiR0RG3DZU4Dbv+kU9ktEJC2COi2RH2WvI2GiyG9MuCqgPII6f1w+5kB5fNIiArqPzaQ==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 + ansi-regex@5.0.1: + resolution: {integrity: sha512-quJQXlTSUGL2LH9SUXo8VwsY4soanhgo6LNSm84E1LBcE8s3O0wpdiRzyR9z/ZZJMlMWv37qOOb9pdJlMUEKFQ==} + engines: {node: '>=8'} + + ansi-styles@5.2.0: + resolution: {integrity: sha512-Cxwpt2SfTzTtXcfOlzGEee8O+c+MmUgGrNiBcXnuWxuFJHe6a5Hz7qwhwe5OgaSYI0IJvkLqWX1ASG+cJOkEiA==} + engines: {node: '>=10'} + + aria-query@5.3.0: + resolution: {integrity: sha512-b0P0sZPKtyu8HkeRAfCq0IfURZK+SuwMjY1UXGBU27wpAiTwQAIlq56IbIO+ytk/JjS1fMR14ee5WBBfKi5J6A==} + + aria-query@5.3.2: + resolution: {integrity: sha512-COROpnaoap1E2F000S62r6A60uHZnmlvomhfyT2DlTcrY1OrBKn2UhH7qn5wTC9zMvD0AY7csdPSNwKP+7WiQw==} + engines: {node: '>= 0.4'} + + assertion-error@2.0.1: + resolution: {integrity: sha512-Izi8RQcffqCeNVgFigKli1ssklIbpHnCYc6AknXGYoB6grJqyeby7jv12JUQgmTAnIDnbck1uxksT4dzN3PWBA==} + engines: {node: '>=12'} + bail@2.0.2: resolution: {integrity: sha512-0xO6mYd7JB2YesxDKplafRpsiOzPt9V02ddPCLbY1xYGPOX24NTyN50qnUxgCPcSoYMhKpAuBTjQoRZCAkUDRw==} @@ -856,12 +963,20 @@ packages: engines: {node: ^6 || ^7 || ^8 || ^9 || ^10 || ^11 || ^12 || >=13.7} hasBin: true + cac@6.7.14: + resolution: {integrity: sha512-b6Ilus+c3RrdDk+JhLKUAQfzzgLEPy6wcXqS7f/xe1EETvsDP6GORG7SFuOs6cID5YkqchW/LXZbX5bc8j7ZcQ==} + engines: {node: '>=8'} + caniuse-lite@1.0.30001799: resolution: {integrity: sha512-hG1bReV+OUU+MOqK4t/ZWI0tZOyz3rqS9XuhOUz1cIcbwBKjOyJEJuw9ER5JuNyqxNk8u/JUVbGibBOL1yrjFw==} ccount@2.0.1: resolution: {integrity: sha512-eyrF0jiFpY+3drT6383f1qhkbGsLSifNAjA61IUjZjmLCWjItY6LB9ft9YhoDgwfmclB2zhu51Lc7+95b8NRAg==} + chai@5.3.3: + resolution: {integrity: sha512-4zNhdJD/iOjSH0A05ea+Ke6MU5mmpQcbQsSOkgdaUMJ9zTlDTD/GYlwohmIE2u0gaxHYiVHEn1Fw9mZ/ktJWgw==} + engines: {node: '>=18'} + character-entities-html4@2.1.0: resolution: {integrity: sha512-1v7fgQRj6hnSwFpq1Eu0ynr/CDEw0rXo2B61qXrLNdHZmPKgb7fqS1a2JwF0rISo9q77jDI8VMEHoApn8qDoZA==} @@ -874,8 +989,9 @@ packages: character-reference-invalid@2.0.1: resolution: {integrity: sha512-iBZ4F4wRbyORVsu0jPV7gXkOsGYjGHPmAyv+HiHG8gi5PtC9KI2j1+v8/tlibRvjoWX027ypmG/n0HtO5t7unw==} - class-variance-authority@0.7.1: - resolution: {integrity: sha512-Ka+9Trutv7G8M6WT6SeiRWz792K5qEqIGEGzXKhAE6xOWAY6pPH8U+9IY3oCMv6kqTmLsv7Xh/2w2RigkePMsg==} + check-error@2.1.3: + resolution: {integrity: sha512-PAJdDJusoxnwm1VwW07VWwUN1sl7smmC3OKggvndJFadxxDRyFJBX/ggnu/KE4kQAB7a3Dp8f/YXC1FlUprWmA==} + engines: {node: '>= 16'} clsx@2.1.1: resolution: {integrity: sha512-eYm0QWBtUrBWZWG0d386OGAw16Z995PiOVo2B7bjWSbHedGl5e0ZWaq65kOGgUSNesEIDkB9ISbTg/JK9dhCZA==} @@ -901,6 +1017,13 @@ packages: cose-base@2.2.0: resolution: {integrity: sha512-AzlgcsCbUMymkADOJtQm3wO9S3ltPfYOFD5033keQn9NJzIbtnZj+UdBJe7DYml/8TdbtHJW3j58SOnKhWY/5g==} + css.escape@1.5.1: + resolution: {integrity: sha512-YUifsXXuknHlUsmlgyY0PKzgPOr7/FjCePfHNt0jxm83wHZi44VDMQ7/fGNkjY3/jV1MC+1CmZbaHzugyeRtpg==} + + cssstyle@4.6.0: + resolution: {integrity: sha512-2z+rWdzbbSZv6/rhtvzvqeZQHrBaqgogqt85sqFNbabZOuFbCVFb8kPeEtZjiKkbrm395irpNKiYeFeLiQnFPg==} + engines: {node: '>=18'} + csstype@3.2.3: resolution: {integrity: sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==} @@ -1058,7 +1181,11 @@ packages: engines: {node: '>=12'} dagre-d3-es@7.0.14: - resolution: {integrity: sha512-P4rFMVq9ESWqmOgK+dlXvOtLwYg0i7u0HBGJER0LZDJT2VHIPAMZ/riPxqJceWMStH5+E61QxFra9kIS3AqdMg==} + resolution: {tarball: https://registry.npmjs.org/dagre-d3-es/-/dagre-d3-es-7.0.14.tgz} + + data-urls@5.0.0: + resolution: {integrity: sha512-ZYP5VBHshaDAiVZxjbRVcFJpc+4xGgT0bK3vzy1HLN8jTO975HEbuYzZJcHoQEY5K1a0z8YayJkyVETa08eNTg==} + engines: {node: '>=18'} dayjs@1.11.21: resolution: {integrity: sha512-98IT+HOahAisibz/yjKbzuOBwYcjJ7BCLPzARyHiyEBmRz4fatF+KPJszEHXsGYjUG234aH/cOjW1wwTbKUZlA==} @@ -1072,9 +1199,16 @@ packages: supports-color: optional: true + decimal.js@10.6.0: + resolution: {integrity: sha512-YpgQiITW3JXGntzdUmyUR1V812Hn8T1YVXhCu+wO3OpS4eU9l4YdD3qjyiKdV6mvV29zapkMeD390UVEf2lkUg==} + decode-named-character-reference@1.3.0: resolution: {integrity: sha512-GtpQYB283KrPp6nRw50q3U9/VfOutZOe103qlN7BPP6Ad27xYnOIWv4lPzo8HCAL+mMZofJ9KEy30fq6MfaK6Q==} + deep-eql@5.0.2: + resolution: {integrity: sha512-h5k/5U50IJJFpzfL6nO9jaaumfjO/f2NjK/oYB2Djzm4p9L+3T9qWpZqZ2hAbLPuuYq9wrU08WQyBTL5GbPk5Q==} + engines: {node: '>=6'} + delaunator@5.1.0: resolution: {integrity: sha512-AGrQ4QSgssa1NGmWmLPqN5NY2KajF5MqxetNEO+o0n3ZwZZeTmt7bBnvzHWrmkZFxGgr4HdyFgelzgi06otLuQ==} @@ -1089,6 +1223,12 @@ packages: devlop@1.1.0: resolution: {integrity: sha512-RWmIqhcFf1lRYBvNmr7qTNuyCt/7/ns2jbpp1+PalgE/rDQcBT0fioSMUpJ93irlUhC5hrg4cYqe6U+0ImW0rA==} + dom-accessibility-api@0.5.16: + resolution: {integrity: sha512-X7BJ2yElsnOJ30pZF4uIIDfBEVgF4XEBxL9Bxhy6dnrm5hkzqmsWHGTiHqRiITNhMyFLyAiWndIJP7Z1NTteDg==} + + dom-accessibility-api@0.6.3: + resolution: {integrity: sha512-7ZgogeTnjuHbo+ct10G9Ffp0mif17idi0IyWNVA/wcwcm7NPOD/WEHVP3n7n3MhXqxoIYm8d6MuZohYWIZ4T3w==} + dompurify@3.4.11: resolution: {integrity: sha512-zhlUV12GsaRzMsf9q5M254YhA4+VuF0fG+QFqu6aYpoGlKtz+w8//jBcGVYBgQkR5GHjUomejY84AV+/uPbWdw==} @@ -1106,6 +1246,9 @@ packages: resolution: {integrity: sha512-aN97NXWF6AWBTahfVOIrB/NShkzi5H7F9r1s9mD3cDj4Ko5f2qhhVoYMibXF7GlLveb/D2ioWay8lxI97Ven3g==} engines: {node: '>=0.12'} + es-module-lexer@1.7.0: + resolution: {integrity: sha512-jEQoCwk8hyb2AZziIOLhDqpm5+2ww5uIE6lkO/6jcOCusfk6LhMHpXXfBLXTZ7Ydyt0j4VoUQv6uGNYbdW+kBA==} + es-toolkit@1.49.0: resolution: {integrity: sha512-G5iZ6Pc/FNRY/soKZHC+TxGDD83rHUDXxzaWhGCX44vAv/tMs56WMusnm/KMNK+luUPsgA9U28cGr4RDlSzL2g==} @@ -1125,10 +1268,17 @@ packages: estree-util-is-identifier-name@3.0.0: resolution: {integrity: sha512-hFtqIDZTIUZ9BXLb8y4pYGyk6+wekIivNVTcmvk8NoOh+VeRn5y6cEHzbURrWbfp1fIqdVipilzj+lfaadNZmg==} + estree-walker@3.0.3: + resolution: {integrity: sha512-7RUKfXgSMMkzt6ZuXmqapOurLGPPfgj6l9uRZ7lRGolvk0y2yocc35LdcxKC5PQZdn2DMqioAQ2NoWcrTKmm6g==} + eventsource-parser@3.1.0: resolution: {integrity: sha512-kJezFj9YFAMLeORyi7aCLxLbD5/qWMQnoMVlVPyHIll7lgRJCc3JVln9Vgl9nwQi0YkMnhdGTMNn7CkRRAptMg==} engines: {node: '>=18.0.0'} + expect-type@1.4.0: + resolution: {integrity: sha512-KfYbmpRm0VbLjEvVa9yGwCi9GI34xvi7A/HXYWQO65CSD2u3MczUJSuwXKFIxlGsgBQizV9q5J9NHj4VG0n+pA==} + engines: {node: '>=12.0.0'} + extend@3.0.2: resolution: {integrity: sha512-fjquC59cD7CyW6urNXK0FBufkZcoiGG80wTuPujX590cB5Ttln20E2UB4S/WARVqhXffZl2LNgS+gQdPIIim/g==} @@ -1183,12 +1333,24 @@ packages: hastscript@9.0.1: resolution: {integrity: sha512-g7df9rMFX/SPi34tyGCyUBREQoKkapwdY/T04Qn9TDWfHhAYt4/I0gMVirzK5wEzeUqIjEB+LXC/ypb7Aqno5w==} + html-encoding-sniffer@4.0.0: + resolution: {integrity: sha512-Y22oTqIU4uuPgEemfz7NDJz6OeKf12Lsu+QC+s3BVpda64lTiMYCyGwg5ki4vFxkMwQdeZDl2adZoqUgdFuTgQ==} + engines: {node: '>=18'} + html-url-attributes@3.0.1: resolution: {integrity: sha512-ol6UPyBWqsrO6EJySPz2O7ZSr856WDrEzM5zMqp+FJJLGMW35cLYmmZnl0vztAZxRUoNZJFTCohfjuIJ8I4QBQ==} html-void-elements@3.0.0: resolution: {integrity: sha512-bEqo66MRXsUGxWHV5IP0PUiAWwoEjba4VCzg0LjFJBpchPaTfyfCKTG6bc5F8ucKec3q5y6qOdGyYTSBEvhCrg==} + http-proxy-agent@7.0.2: + resolution: {integrity: sha512-T1gkAiYYDWYx3V5Bmyu7HcfcvL7mUrTWiM6yOfa3PIphViJ/gFPbvidQ+veqSOHci/PxBcDabeUNCzpOODJZig==} + engines: {node: '>= 14'} + + https-proxy-agent@7.0.6: + resolution: {integrity: sha512-vK9P5/iUfdl95AI+JVyUuIcVtd4ofvtrOr3HNtM2yxC9bnMbEdp3x01OhQNnjb8IJYi38VlTE3mBXwcfvywuSw==} + engines: {node: '>= 14'} + iconv-lite@0.6.3: resolution: {integrity: sha512-4fCk79wshMdzMp2rH06qWrJE4iolqLhCUH+OiuIgU++RB0+94NlDL81atO7GX55uUKueo0txHNtvEyI6D7WdMw==} engines: {node: '>=0.10.0'} @@ -1196,6 +1358,10 @@ packages: import-meta-resolve@4.2.0: resolution: {integrity: sha512-Iqv2fzaTQN28s/FwZAoFq0ZSs/7hMAHJVX+w8PZl3cY19Pxk6jFFalxQoIfW2826i/fDLXv8IiEZRIT0lDuWcg==} + indent-string@4.0.0: + resolution: {integrity: sha512-EdDDZu4A2OyIK7Lr/2zG+w5jmbuk1DVBnEwREQvBzspBJkCEbRa8GxU1lghYcaGJCnRWibjDXlq779X1/y5xwg==} + engines: {node: '>=8'} + inline-style-parser@0.2.7: resolution: {integrity: sha512-Nb2ctOyNR8DqQoR0OwRG95uNWIC0C1lCgf5Naz5H6Ji72KZ8OcFZLz2P5sNgwlyoJ8Yif11oMuYs5pBQa86csA==} @@ -1222,6 +1388,9 @@ packages: resolution: {integrity: sha512-+Pgi+vMuUNkJyExiMBt5IlFoMyKnr5zhJ4Uspz58WOhBF5QoIZkFyNHIbBAtHwzVAgk5RtndVNsDRN61/mmDqg==} engines: {node: '>=12'} + is-potential-custom-element-name@1.0.1: + resolution: {integrity: sha512-bCYeRA2rVibKZd+s2625gGnGF/t7DSqDs4dP7CrLA1m7jKWz6pps0LpYLJN8Q64HtmPKJ1hrN3nzPNKFEKOUiQ==} + jiti@2.7.0: resolution: {integrity: sha512-AC/7JofJvZGrrneWNaEnJeOLUx+JlGt7tNa0wZiRPT4MY1wmfKjt2+6O2p2uz2+skll8OZZmJMNqeke7kKbNgQ==} hasBin: true @@ -1229,6 +1398,18 @@ packages: js-tokens@4.0.0: resolution: {integrity: sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==} + js-tokens@9.0.1: + resolution: {integrity: sha512-mxa9E9ITFOt0ban3j6L5MpjwegGz6lBQmM1IJkWeBZGcMxto50+eWdjC/52xDbS2vy0k7vIMK0Fe2wfL9OQSpQ==} + + jsdom@26.1.0: + resolution: {integrity: sha512-Cvc9WUhxSMEo4McES3P7oK3QaXldCfNWp7pl2NNeiIFlCoLr3kfq9kb1fxftiwk1FLV7CvpvDfonxtzUDeSOPg==} + engines: {node: '>=18'} + peerDependencies: + canvas: ^3.0.0 + peerDependenciesMeta: + canvas: + optional: true + jsesc@3.1.0: resolution: {integrity: sha512-/sM3dO2FOzXjKQhJuo0Q173wf2KOo8t4I8vHy6lF9poUp7bKT0/NHE8fPX23PwfhnykfqnC2xRxOnVw5XuGIaA==} engines: {node: '>=6'} @@ -1339,9 +1520,19 @@ packages: resolution: {integrity: sha512-lyuxPGr/Wfhrlem2CL/UcnUc1zcqKAImBDzukY7Y5F/yQiNdko6+fRLevlw1HgMySw7f611UIY408EtxRSoK3Q==} hasBin: true + loupe@3.2.1: + resolution: {integrity: sha512-CdzqowRJCeLU72bHvWqwRBBlLcMEtIvGrlvef74kMnV2AolS9Y8xUv1I0U/MNAWMhBlKIoyuEgoJ0t/bbwHbLQ==} + + lru-cache@10.4.3: + resolution: {integrity: sha512-JNAzZcXrCt42VGLuYz0zfAzDfAvJWW6AfYlDBQyDV5DClI2m5sAmK+OIO7s59XfsRsWHp02jAJrRadPRGTt6SQ==} + lru-cache@5.1.1: resolution: {integrity: sha512-KpNARQA3Iwv+jTA0utUVVbrh+Jlrr1Fv0e56GGzAFOXN7dk/FviaDW8LHmK52DlcH4WP2n6gI8vN1aesBFgo9w==} + lz-string@1.5.0: + resolution: {integrity: sha512-h5bgJWpxJNswbU7qCrV0tIKQCaS3blPDrqKWx+QxzuzL1zGUzij9XCWLrSLsJPu5t+eWA/ycetzYAO5IOMcWAQ==} + hasBin: true + magic-string@0.30.21: resolution: {integrity: sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ==} @@ -1495,6 +1686,13 @@ packages: micromark@4.0.2: resolution: {integrity: sha512-zpe98Q6kvavpCr1NPVSCMebCKfD7CA2NqZ+rykeNhONIJBpc1tFKt9hucLGwha3jNTNI8lHpctWJWoimVF4PfA==} + min-indent@1.0.1: + resolution: {integrity: sha512-I9jwMn07Sy/IwOj3zVkVik2JTvgpaykDZEigL6Rx6N9LbMywwUSMtxET+7lVoDLLd3O3IXwJwvuuns8UB/HeAg==} + engines: {node: '>=4'} + + monaco-editor@0.48.0: + resolution: {integrity: sha512-goSDElNqFfw7iDHMg8WDATkfcyeLTNpBHQpO8incK6p5qZt5G/1j41X0xdGzpIkGojGXM+QiRQyLjnfDVvrpwA==} + ms@2.1.3: resolution: {integrity: sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==} @@ -1507,6 +1705,9 @@ packages: resolution: {integrity: sha512-J6l92tKHX6w8Jy5nO1Vuc01NoIiRGi/d6qBKVxh+IQ8Cr3b6HbVNfKiF8ZpFKufTwpwxMmce2W3iQZ861ZRyTg==} engines: {node: '>=18'} + nwsapi@2.2.24: + resolution: {integrity: sha512-7YRhZ3jS45LwmSCT4b2sVFHt/WuovaktDU07QrtOBY2PXskss5a9jfmR9jptyumwXST+rFjrmppMY1KT/yn35A==} + object-assign@4.1.1: resolution: {integrity: sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg==} engines: {node: '>=0.10.0'} @@ -1526,6 +1727,13 @@ packages: path-data-parser@0.1.0: resolution: {integrity: sha512-NOnmBpt5Y2RWbuv0LMzsayp3lVylAHLPUTut412ZA3l+C4uw4ZVkQbjShYCQ8TCpUMdPapr4YjUqLYD6v68j+w==} + pathe@2.0.3: + resolution: {integrity: sha512-WUjGcAqP1gQacoQe+OBJsFA7Ld4DyXuUIjZ5cc75cLHvJ7dtNsTugphxIADwspS+AraAUePCKrSVtPLFj/F88w==} + + pathval@2.0.1: + resolution: {integrity: sha512-//nshmD55c46FuFw26xV/xFAaB5HF9Xdap7HJBBnrKdAd6/GxDBaNA1870O79+9ueg61cZLSVc+OaFlfmObYVQ==} + engines: {node: '>= 14.16'} + picocolors@1.1.1: resolution: {integrity: sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA==} @@ -1543,12 +1751,20 @@ packages: resolution: {integrity: sha512-vuwillviilfKZsg0VGj5R/YwwcHx4SLsIOI/7K6mQkWx+l5cUHTjj5g0AasTBcyXsbfTgrwsUNmVUb5xVwyPwg==} engines: {node: ^10 || ^12 || >=14} + pretty-format@27.5.1: + resolution: {integrity: sha512-Qb1gy5OrP5+zDf2Bvnzdl3jsTf1qXVMazbvCoKhtKqVs4/YK4ozX4gKQJJVyNe+cajNPn0KoC0MC3FUmaHWEmQ==} + engines: {node: ^10.13.0 || ^12.13.0 || ^14.15.0 || >=15.0.0} + prop-types@15.8.1: resolution: {integrity: sha512-oj87CgZICdulUohogVAR7AjlC0327U4el4L6eAvOqCeudMDVU0NThNaV+b9Df4dXgSP1gXMTnPdhfe/2qDH5cg==} property-information@7.2.0: resolution: {integrity: sha512-IAtzIB6sUiWaJYrX9smp3V46pBGbBeLFRGdh25kg1334VcBlD8HzhPeNIWQH9zhGmo2itIe25EHt9dQP7G5hmg==} + punycode@2.3.1: + resolution: {integrity: sha512-vYt7UD1U9Wg6138shLtLOvdAu+8DsC/ilFtEVHcH+wydcSpNE20AfSOduf6MkRFahL5FY7X1oU7nKVZFtfq8Fg==} + engines: {node: '>=6'} + re-resizable@6.11.2: resolution: {integrity: sha512-2xI2P3OHs5qw7K0Ud1aLILK6MQxW50TcO+DetD9eIV58j84TqYeHoZcL9H4GXFXXIh7afhH8mv5iUCXII7OW7A==} peerDependencies: @@ -1569,6 +1785,9 @@ packages: react-is@16.13.1: resolution: {integrity: sha512-24e6ynE2H+OKt4kqsOvNd8kBpV65zoxbA4BVsEOB3ARVWQki/DHzaUoC5KuON/BiccDaCCTZBuOcfZs70kR8bQ==} + react-is@17.0.2: + resolution: {integrity: sha512-w2GsyukL62IJnlaff/nRegPQR94C/XXamvMWmSHRJ4y7Ts/4ocGRmTHvOs8PSE6pB3dWOrD/nueuU5sduBsQ4w==} + react-refresh@0.18.0: resolution: {integrity: sha512-QgT5//D3jfjJb6Gsjxv0Slpj23ip+HtOpnNgnb2S5zU3CB26G/IDPGoy4RJB42wzFE46DRsstbW6tKHoKbhAxw==} engines: {node: '>=0.10.0'} @@ -1583,6 +1802,10 @@ packages: resolution: {integrity: sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==} engines: {node: '>=0.10.0'} + redent@3.0.0: + resolution: {integrity: sha512-6tDA8g98We0zd0GvVeMT9arEOnTw9qM03L9cJXaCjrip1OO764RDBLBfrB4cwzNGDj5OA5ioymC9GkizgWJDUg==} + engines: {node: '>=8'} + regex-recursion@5.1.1: resolution: {integrity: sha512-ae7SBCbzVNrIjgSbh7wMznPcQel1DNlDtzensnFxpiNpXt1U2ju/bHugH422r+4LAVS1FpW1YCwilmnNsjum9w==} @@ -1627,12 +1850,19 @@ packages: roughjs@4.6.6: resolution: {integrity: sha512-ZUz/69+SYpFN/g/lUlo2FXcIjRkSu3nDarreVdGGndHEBJ6cXPdKguS8JGxwj5HA5xIbVKSmLgr5b3AWxtRfvQ==} + rrweb-cssom@0.8.0: + resolution: {integrity: sha512-guoltQEx+9aMf2gDZ0s62EcV8lsXR+0w8915TC3ITdn2YueuNjdAYh/levpU9nFaoChh9RUS5ZdQMrKfVEN9tw==} + rw@1.3.3: resolution: {integrity: sha512-PdhdWy89SiZogBLaw42zdeqtRJ//zFd2PgQavcICDUgJT5oW10QCRKbJ6bg4r0/UY2M6BWd5tkxuGFRvCkgfHQ==} safer-buffer@2.1.2: resolution: {integrity: sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==} + saxes@6.0.0: + resolution: {integrity: sha512-xAg7SOnEhrm5zI3puOOKyy1OMcMlIJZYNJY7xLBwSze0UjhPLnWfj2GF2EpT0jmzaJKIWKHLsaSSajf35bcYnA==} + engines: {node: '>=v12.22.7'} + scheduler@0.23.2: resolution: {integrity: sha512-UOShsPwz7NrMUqhR6t0hWjFduvOzbtv7toDH1/hIrfRNIDBnnBWd0CwJTGvTpngVlmwGCdP9/Zl/tVrDqcuYzQ==} @@ -1643,6 +1873,9 @@ packages: shiki@1.29.2: resolution: {integrity: sha512-njXuliz/cP+67jU2hukkxCNuH1yUi4QfdZZY+sMr5PPrIyXSu5iTb/qYC4BiWWB0vZ+7TbdvYUCeL23zpwCfbg==} + siginfo@2.0.0: + resolution: {integrity: sha512-ybx0WO1/8bSBLEWXZvEd7gMW3Sn3JFlW3TvX1nREbDLRNQNaeNN8WK0meBwPdAaOI7TtRRRJn/Es1zhrrCHu7g==} + source-map-js@1.2.1: resolution: {integrity: sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==} engines: {node: '>=0.10.0'} @@ -1650,6 +1883,12 @@ packages: space-separated-tokens@2.0.2: resolution: {integrity: sha512-PEGlAwrG8yXGXRjW32fGbg66JAlOAwbObuqVoJpv/mRgoWDQfgH1wDPvtzWyUSNAXBGSk8h755YDbbcEy3SH2Q==} + stackback@0.0.2: + resolution: {integrity: sha512-1XMJE5fQo1jGH6Y/7ebnwPOBEkIEnT4QF32d5R1+VXdXveM0IBMJt8zfaxX1P3QhVwrYe+576+jkANtSS2mBbw==} + + std-env@3.10.0: + resolution: {integrity: sha512-5GS12FdOZNliM5mAOxFRg7Ir0pWz8MdpYm6AY6VPkGpbA7ZzmbzNcBJQ0GPvvyWgcY7QAhCgf9Uy89I03faLkg==} + streamdown@2.5.0: resolution: {integrity: sha512-/tTnURfIOxZK/pqJAxsfCvETG/XCJHoWnk3jq9xLcuz6CSpnjjuxSRBTTL4PKGhxiZQf0lqPxGhImdpwcZ2XwA==} peerDependencies: @@ -1659,6 +1898,13 @@ packages: stringify-entities@4.0.4: resolution: {integrity: sha512-IwfBptatlO+QCJUo19AqvrPNqlVMpW9YEL2LIVY+Rpv2qsjCGxaDLNRgeGsQWJhfItebuJhsGSLjaBbNSQ+ieg==} + strip-indent@3.0.0: + resolution: {integrity: sha512-laJTa3Jb+VQpaC6DseHhF7dXVqHTfJPCRDaEbid/drOhgitgYku/letMUqOXFoWV0zIIUbjpdH2t+tYj4bQMRQ==} + engines: {node: '>=8'} + + strip-literal@3.1.0: + resolution: {integrity: sha512-8r3mkIM/2+PpjHoOtiAW8Rg3jJLHaV7xPwG+YRGrv6FP0wwk/toTpATxWYOW0BKdWwl82VT2tFYi5DlROa0Mxg==} + style-to-js@1.1.21: resolution: {integrity: sha512-RjQetxJrrUJLQPHbLku6U/ocGtzyjbJMP9lCNK7Ag0CNh690nSH8woqWH9u16nMjYBAok+i7JO1NP2pOy8IsPQ==} @@ -1673,8 +1919,8 @@ packages: peerDependencies: react: ^16.11.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 - tailwind-merge@2.6.1: - resolution: {integrity: sha512-Oo6tHdpZsGpkKG88HJ8RR1rg/RdnEkQEfMoEk2x1XRI3F1AxeU+ijRXpiVUF4UbLfcxxRGw6TbUINKYdWVsQTQ==} + symbol-tree@3.2.4: + resolution: {integrity: sha512-9QNk5KwDF+Bvz+PyObkmSYjI5ksVUYtjW7AU22r2NKcfLJcXp96hkDWU3+XndOsUb+AQ9QhfzfCT2O+CNWT5Tw==} tailwind-merge@3.6.0: resolution: {integrity: sha512-uxL7qAVQriqRQPAyK3pj66VqskWqoZ37PW94jwOTwNfq/z9oyu1V+eqrZqtR2+fCiXdYOZe/Modt8GtvqNzu+w==} @@ -1690,6 +1936,12 @@ packages: resolution: {integrity: sha512-nt6AMGKW1p/70DF/hGBdJB57B8Tspmbp5gfJ8ilhLnt7kkr2ye7hzD6NVG8GGErk2HWF34igrL2CXmNIkzKqKw==} engines: {node: '>=18'} + tinybench@2.9.0: + resolution: {integrity: sha512-0+DUvqWMValLmha6lr4kD8iAMK1HzV0/aKnCtWb9v9641TnP/MFb7Pc2bxoxQjTXAErryXVgUOfv2YqNllqGeg==} + + tinyexec@0.3.2: + resolution: {integrity: sha512-KQQR9yN7R5+OSwaK0XQoj22pwHoTlgYqmUscPYoknOoWCWfj/5/ABTMRi69FrKU5ffPVh5QcFikpWJI/P1ocHA==} + tinyexec@1.2.4: resolution: {integrity: sha512-SHf/r48b7vOrjve9PxJo3MN5v5yuyjHvdUcrQffT3WXMUfnGmHDVbC4k3sHJaJTgZCwpUplIaAo5ANtMyp3YHg==} engines: {node: '>=18'} @@ -1698,6 +1950,33 @@ packages: resolution: {integrity: sha512-wXR/dYpcqKmfWpEdZjiKJOwCNFndD0DMnrW/cYjVGttEkBfVgcLFHoNrlj47mjOVic9yyNu65alsgF4NQyTa2g==} engines: {node: '>=12.0.0'} + tinypool@1.1.1: + resolution: {integrity: sha512-Zba82s87IFq9A9XmjiX5uZA/ARWDrB03OHlq+Vw1fSdt0I+4/Kutwy8BP4Y/y/aORMo61FQ0vIb5j44vSo5Pkg==} + engines: {node: ^18.0.0 || >=20.0.0} + + tinyrainbow@2.0.0: + resolution: {integrity: sha512-op4nsTR47R6p0vMUUoYl/a+ljLFVtlfaXkLQmqfLR1qHma1h/ysYk4hEXZ880bf2CYgTskvTa/e196Vd5dDQXw==} + engines: {node: '>=14.0.0'} + + tinyspy@4.0.4: + resolution: {integrity: sha512-azl+t0z7pw/z958Gy9svOTuzqIk6xq+NSheJzn5MMWtWTFywIacg2wUlzKFGtt3cthx0r2SxMK0yzJOR0IES7Q==} + engines: {node: '>=14.0.0'} + + tldts-core@6.1.86: + resolution: {integrity: sha512-Je6p7pkk+KMzMv2XXKmAE3McmolOQFdxkKw0R8EYNr7sELW46JqnNeTX8ybPiQgvg1ymCoF8LXs5fzFaZvJPTA==} + + tldts@6.1.86: + resolution: {integrity: sha512-WMi/OQ2axVTf/ykqCQgXiIct+mSQDFdH2fkwhPwgEwvJ1kSzZRiinb0zF2Xb8u4+OqPChmyI6MEu4EezNJz+FQ==} + hasBin: true + + tough-cookie@5.1.2: + resolution: {integrity: sha512-FVDYdxtnj0G6Qm/DhNPSb8Ju59ULcup3tuJxkFb5K8Bv2pUXILbf0xZWU8PX8Ov19OXljbUyveOFwRMwkXzO+A==} + engines: {node: '>=16'} + + tr46@5.1.1: + resolution: {integrity: sha512-hdF5ZgjTqgAntKkklYw0R03MG2x/bSzTtkxmIRw/sTNV8YXsCJ1tfLAX23lhxhHJlEf3CRCOCGGWw3vI3GaSPw==} + engines: {node: '>=18'} + trim-lines@3.0.1: resolution: {integrity: sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg==} @@ -1758,6 +2037,11 @@ packages: vfile@6.0.3: resolution: {integrity: sha512-KzIbH/9tXat2u30jf+smMwFCsno4wHVdNmzFyL+T/L3UGqqk6JKfVqOFOZEpZSHADH1k40ab6NUIXZq422ov3Q==} + vite-node@3.2.4: + resolution: {integrity: sha512-EbKSKh+bh1E1IFxeO0pg1n4dvoOTt0UDiXMd/qn++r98+jPO1xtJilvXldeuQ8giIB5IkpjCgMleHMNEsGH6pg==} + engines: {node: ^18.0.0 || ^20.0.0 || >=22.0.0} + hasBin: true + vite@7.3.6: resolution: {integrity: sha512-4XP60spRGjSZFf1qYH+dJIkK2znL3zQfl9KkOV9MkkRR/3Dls0dxaBsQPTloEc5BLXWPL9vsOxopxyKoMmDueg==} engines: {node: ^20.19.0 || >=22.12.0} @@ -1798,12 +2082,90 @@ packages: yaml: optional: true + vitest@3.2.7: + resolution: {integrity: sha512-KrxIJ62Fd89gfysR4WotlgZABiz2dqFPgqGzX7s+CwsqLFomRH7777ZcrOD6+WVAh7khPQP41A+BKbpcJFrdEg==} + engines: {node: ^18.0.0 || ^20.0.0 || >=22.0.0} + hasBin: true + peerDependencies: + '@edge-runtime/vm': '*' + '@types/debug': ^4.1.12 + '@types/node': ^18.0.0 || ^20.0.0 || >=22.0.0 + '@vitest/browser': 3.2.7 + '@vitest/ui': 3.2.7 + happy-dom: '*' + jsdom: '*' + peerDependenciesMeta: + '@edge-runtime/vm': + optional: true + '@types/debug': + optional: true + '@types/node': + optional: true + '@vitest/browser': + optional: true + '@vitest/ui': + optional: true + happy-dom: + optional: true + jsdom: + optional: true + + w3c-xmlserializer@5.0.0: + resolution: {integrity: sha512-o8qghlI8NZHU1lLPrpi2+Uq7abh4GGPpYANlalzWxyWteJOCsr/P+oPBA49TOLu5FTZO4d3F9MnWJfiMo4BkmA==} + engines: {node: '>=18'} + web-namespaces@2.0.1: resolution: {integrity: sha512-bKr1DkiNa2krS7qxNtdrtHAmzuYGFQLiQ13TsorsdT6ULTkPLKuu5+GsFpDlg6JFjUTwX2DyhMPG2be8uPrqsQ==} + webidl-conversions@7.0.0: + resolution: {integrity: sha512-VwddBukDzu71offAQR975unBIGqfKZpM+8ZX6ySk8nYhVoo5CYaZyzt3YBvYtRtO+aoGlqxPg/B87NGVZ/fu6g==} + engines: {node: '>=12'} + + whatwg-encoding@3.1.1: + resolution: {integrity: sha512-6qN4hJdMwfYBtE3YBTTHhoeuUrDBPZmbQaxWAqSALV/MeEnR5z1xd8UKud2RAkFoPkmB+hli1TZSnyi84xz1vQ==} + engines: {node: '>=18'} + deprecated: Use @exodus/bytes instead for a more spec-conformant and faster implementation + + whatwg-mimetype@4.0.0: + resolution: {integrity: sha512-QaKxh0eNIi2mE9p2vEdzfagOKHCcj1pJ56EEHGQOVxp8r9/iszLUUV7v89x9O1p/T+NlTM5W7jW6+cz4Fq1YVg==} + engines: {node: '>=18'} + + whatwg-url@14.2.0: + resolution: {integrity: sha512-De72GdQZzNTUBBChsXueQUnPKDkg/5A5zp7pFDuQAj5UFoENpiACU0wlCvzpAGnTkj++ihpKwKyYewn/XNUbKw==} + engines: {node: '>=18'} + + why-is-node-running@2.3.0: + resolution: {integrity: sha512-hUrmaWBdVDcxvYqnyh09zunKzROWjbZTiNy8dBEjkS7ehEDQibXJ7XvlmtbwuTclUiIyN+CyXQD4Vmko8fNm8w==} + engines: {node: '>=8'} + hasBin: true + + ws@8.21.0: + resolution: {integrity: sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==} + engines: {node: '>=10.0.0'} + peerDependencies: + bufferutil: ^4.0.1 + utf-8-validate: '>=5.0.2' + peerDependenciesMeta: + bufferutil: + optional: true + utf-8-validate: + optional: true + + xml-name-validator@5.0.0: + resolution: {integrity: sha512-EvGK8EJ3DhaHfbRlETOWAS5pO9MZITeauHKJyb8wyajUfQUenkIg2MvLDTZ4T/TgIcm3HU0TFBgWWboAZ30UHg==} + engines: {node: '>=18'} + + xmlchars@2.2.0: + resolution: {integrity: sha512-JZnDKK8B0RCDw84FNdDAIpZK+JuJw+s7Lz8nksI7SIuU3UXJJslUthsi+uWBUYOwPFwW7W7PRLRfUKpxjtjFCw==} + yallist@3.1.1: resolution: {integrity: sha512-a4UGQaWPH59mOXUYnAG2ewncQS4i4F43Tv3JoAM+s2VDAmS9NsK8GpDMLrCHPksFT7h3K6TOoUNn2pb7RoXx4g==} + yaml@2.9.0: + resolution: {integrity: sha512-2AvhNX3mb8zd6Zy7INTtSpl1F15HW6Wnqj0srWlkKLcpYl/gMIMJiyuGq2KeI2YFxUPjdlB+3Lc10seMLtL4cA==} + engines: {node: '>= 14.6'} + hasBin: true + zod@4.4.3: resolution: {integrity: sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ==} @@ -1812,6 +2174,8 @@ packages: snapshots: + '@adobe/css-tools@4.5.0': {} + '@ai-sdk/gateway@3.0.139(zod@4.4.3)': dependencies: '@ai-sdk/provider': 3.0.12 @@ -1845,6 +2209,14 @@ snapshots: package-manager-detector: 1.6.0 tinyexec: 1.2.4 + '@asamuzakjp/css-color@3.2.0': + dependencies: + '@csstools/css-calc': 2.1.4(@csstools/css-parser-algorithms@3.0.5(@csstools/css-tokenizer@3.0.4))(@csstools/css-tokenizer@3.0.4) + '@csstools/css-color-parser': 3.1.0(@csstools/css-parser-algorithms@3.0.5(@csstools/css-tokenizer@3.0.4))(@csstools/css-tokenizer@3.0.4) + '@csstools/css-parser-algorithms': 3.0.5(@csstools/css-tokenizer@3.0.4) + '@csstools/css-tokenizer': 3.0.4 + lru-cache: 10.4.3 + '@babel/code-frame@7.29.7': dependencies: '@babel/helper-validator-identifier': 7.29.7 @@ -1934,6 +2306,8 @@ snapshots: '@babel/core': 7.29.7 '@babel/helper-plugin-utils': 7.29.7 + '@babel/runtime@7.29.7': {} + '@babel/template@7.29.7': dependencies: '@babel/code-frame': 7.29.7 @@ -1961,6 +2335,26 @@ snapshots: '@chevrotain/types@11.1.2': {} + '@csstools/color-helpers@5.1.0': {} + + '@csstools/css-calc@2.1.4(@csstools/css-parser-algorithms@3.0.5(@csstools/css-tokenizer@3.0.4))(@csstools/css-tokenizer@3.0.4)': + dependencies: + '@csstools/css-parser-algorithms': 3.0.5(@csstools/css-tokenizer@3.0.4) + '@csstools/css-tokenizer': 3.0.4 + + '@csstools/css-color-parser@3.1.0(@csstools/css-parser-algorithms@3.0.5(@csstools/css-tokenizer@3.0.4))(@csstools/css-tokenizer@3.0.4)': + dependencies: + '@csstools/color-helpers': 5.1.0 + '@csstools/css-calc': 2.1.4(@csstools/css-parser-algorithms@3.0.5(@csstools/css-tokenizer@3.0.4))(@csstools/css-tokenizer@3.0.4) + '@csstools/css-parser-algorithms': 3.0.5(@csstools/css-tokenizer@3.0.4) + '@csstools/css-tokenizer': 3.0.4 + + '@csstools/css-parser-algorithms@3.0.5(@csstools/css-tokenizer@3.0.4)': + dependencies: + '@csstools/css-tokenizer': 3.0.4 + + '@csstools/css-tokenizer@3.0.4': {} + '@esbuild/aix-ppc64@0.28.1': optional: true @@ -2072,19 +2466,6 @@ snapshots: '@opentelemetry/api@1.9.1': {} - '@radix-ui/react-compose-refs@1.1.3(@types/react@19.2.17)(react@18.3.1)': - dependencies: - react: 18.3.1 - optionalDependencies: - '@types/react': 19.2.17 - - '@radix-ui/react-slot@1.3.0(@types/react@19.2.17)(react@18.3.1)': - dependencies: - '@radix-ui/react-compose-refs': 1.1.3(@types/react@19.2.17)(react@18.3.1) - react: 18.3.1 - optionalDependencies: - '@types/react': 19.2.17 - '@rolldown/pluginutils@1.0.0-rc.3': {} '@rollup/rollup-android-arm-eabi@4.62.2': @@ -2265,12 +2646,12 @@ snapshots: '@tailwindcss/oxide-win32-arm64-msvc': 4.3.1 '@tailwindcss/oxide-win32-x64-msvc': 4.3.1 - '@tailwindcss/vite@4.3.1(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0))': + '@tailwindcss/vite@4.3.1(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0))': dependencies: '@tailwindcss/node': 4.3.1 '@tailwindcss/oxide': 4.3.1 tailwindcss: 4.3.1 - vite: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0) + vite: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0) '@tanstack/query-core@5.101.2': {} @@ -2279,6 +2660,38 @@ snapshots: '@tanstack/query-core': 5.101.2 react: 18.3.1 + '@testing-library/dom@10.4.1': + dependencies: + '@babel/code-frame': 7.29.7 + '@babel/runtime': 7.29.7 + '@types/aria-query': 5.0.4 + aria-query: 5.3.0 + dom-accessibility-api: 0.5.16 + lz-string: 1.5.0 + picocolors: 1.1.1 + pretty-format: 27.5.1 + + '@testing-library/jest-dom@6.9.1': + dependencies: + '@adobe/css-tools': 4.5.0 + aria-query: 5.3.2 + css.escape: 1.5.1 + dom-accessibility-api: 0.6.3 + picocolors: 1.1.1 + redent: 3.0.0 + + '@testing-library/react@16.3.2(@testing-library/dom@10.4.1)(@types/react-dom@19.2.3(@types/react@19.2.17))(@types/react@19.2.17)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': + dependencies: + '@babel/runtime': 7.29.7 + '@testing-library/dom': 10.4.1 + react: 18.3.1 + react-dom: 18.3.1(react@18.3.1) + optionalDependencies: + '@types/react': 19.2.17 + '@types/react-dom': 19.2.3(@types/react@19.2.17) + + '@types/aria-query@5.0.4': {} + '@types/babel__core@7.20.5': dependencies: '@babel/parser': 7.29.7 @@ -2300,6 +2713,11 @@ snapshots: dependencies: '@babel/types': 7.29.7 + '@types/chai@5.2.3': + dependencies: + '@types/deep-eql': 4.0.2 + assertion-error: 2.0.1 + '@types/d3-array@3.2.2': {} '@types/d3-axis@3.0.6': @@ -2421,6 +2839,8 @@ snapshots: dependencies: '@types/ms': 2.1.0 + '@types/deep-eql@4.0.2': {} + '@types/estree-jsx@1.0.5': dependencies: '@types/estree': 1.0.9 @@ -2463,7 +2883,7 @@ snapshots: '@vercel/oidc@3.2.0': {} - '@vitejs/plugin-react@5.2.0(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0))': + '@vitejs/plugin-react@5.2.0(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0))': dependencies: '@babel/core': 7.29.7 '@babel/plugin-transform-react-jsx-self': 7.29.7(@babel/core@7.29.7) @@ -2471,10 +2891,54 @@ snapshots: '@rolldown/pluginutils': 1.0.0-rc.3 '@types/babel__core': 7.20.5 react-refresh: 0.18.0 - vite: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0) + vite: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0) transitivePeerDependencies: - supports-color + '@vitest/expect@3.2.7': + dependencies: + '@types/chai': 5.2.3 + '@vitest/spy': 3.2.7 + '@vitest/utils': 3.2.7 + chai: 5.3.3 + tinyrainbow: 2.0.0 + + '@vitest/mocker@3.2.7(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0))': + dependencies: + '@vitest/spy': 3.2.7 + estree-walker: 3.0.3 + magic-string: 0.30.21 + optionalDependencies: + vite: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0) + + '@vitest/pretty-format@3.2.7': + dependencies: + tinyrainbow: 2.0.0 + + '@vitest/runner@3.2.7': + dependencies: + '@vitest/utils': 3.2.7 + pathe: 2.0.3 + strip-literal: 3.1.0 + + '@vitest/snapshot@3.2.7': + dependencies: + '@vitest/pretty-format': 3.2.7 + magic-string: 0.30.21 + pathe: 2.0.3 + + '@vitest/spy@3.2.7': + dependencies: + tinyspy: 4.0.4 + + '@vitest/utils@3.2.7': + dependencies: + '@vitest/pretty-format': 3.2.7 + loupe: 3.2.1 + tinyrainbow: 2.0.0 + + agent-base@7.1.4: {} + ai@6.0.214(zod@4.4.3): dependencies: '@ai-sdk/gateway': 3.0.139(zod@4.4.3) @@ -2483,6 +2947,18 @@ snapshots: '@opentelemetry/api': 1.9.1 zod: 4.4.3 + ansi-regex@5.0.1: {} + + ansi-styles@5.2.0: {} + + aria-query@5.3.0: + dependencies: + dequal: 2.0.3 + + aria-query@5.3.2: {} + + assertion-error@2.0.1: {} + bail@2.0.2: {} baseline-browser-mapping@2.10.40: {} @@ -2495,10 +2971,20 @@ snapshots: node-releases: 2.0.50 update-browserslist-db: 1.2.3(browserslist@4.28.4) + cac@6.7.14: {} + caniuse-lite@1.0.30001799: {} ccount@2.0.1: {} + chai@5.3.3: + dependencies: + assertion-error: 2.0.1 + check-error: 2.1.3 + deep-eql: 5.0.2 + loupe: 3.2.1 + pathval: 2.0.1 + character-entities-html4@2.1.0: {} character-entities-legacy@3.0.0: {} @@ -2507,9 +2993,7 @@ snapshots: character-reference-invalid@2.0.1: {} - class-variance-authority@0.7.1: - dependencies: - clsx: 2.1.1 + check-error@2.1.3: {} clsx@2.1.1: {} @@ -2529,6 +3013,13 @@ snapshots: dependencies: layout-base: 2.0.1 + css.escape@1.5.1: {} + + cssstyle@4.6.0: + dependencies: + '@asamuzakjp/css-color': 3.2.0 + rrweb-cssom: 0.8.0 + csstype@3.2.3: {} cytoscape-cose-bilkent@4.1.0(cytoscape@3.34.0): @@ -2715,16 +3206,25 @@ snapshots: d3: 7.9.0 lodash-es: 4.18.1 + data-urls@5.0.0: + dependencies: + whatwg-mimetype: 4.0.0 + whatwg-url: 14.2.0 + dayjs@1.11.21: {} debug@4.4.3: dependencies: ms: 2.1.3 + decimal.js@10.6.0: {} + decode-named-character-reference@1.3.0: dependencies: character-entities: 2.0.2 + deep-eql@5.0.2: {} + delaunator@5.1.0: dependencies: robust-predicates: 3.0.3 @@ -2737,6 +3237,10 @@ snapshots: dependencies: dequal: 2.0.3 + dom-accessibility-api@0.5.16: {} + + dom-accessibility-api@0.6.3: {} + dompurify@3.4.11: optionalDependencies: '@types/trusted-types': 2.0.7 @@ -2752,6 +3256,8 @@ snapshots: entities@6.0.1: {} + es-module-lexer@1.7.0: {} + es-toolkit@1.49.0: {} esbuild@0.28.1: @@ -2789,8 +3295,14 @@ snapshots: estree-util-is-identifier-name@3.0.0: {} + estree-walker@3.0.3: + dependencies: + '@types/estree': 1.0.9 + eventsource-parser@3.1.0: {} + expect-type@1.4.0: {} + extend@3.0.2: {} fdir@6.5.0(picomatch@4.0.4): @@ -2899,16 +3411,36 @@ snapshots: property-information: 7.2.0 space-separated-tokens: 2.0.2 + html-encoding-sniffer@4.0.0: + dependencies: + whatwg-encoding: 3.1.1 + html-url-attributes@3.0.1: {} html-void-elements@3.0.0: {} + http-proxy-agent@7.0.2: + dependencies: + agent-base: 7.1.4 + debug: 4.4.3 + transitivePeerDependencies: + - supports-color + + https-proxy-agent@7.0.6: + dependencies: + agent-base: 7.1.4 + debug: 4.4.3 + transitivePeerDependencies: + - supports-color + iconv-lite@0.6.3: dependencies: safer-buffer: 2.1.2 import-meta-resolve@4.2.0: {} + indent-string@4.0.0: {} + inline-style-parser@0.2.7: {} internmap@1.0.1: {} @@ -2928,10 +3460,41 @@ snapshots: is-plain-obj@4.1.0: {} + is-potential-custom-element-name@1.0.1: {} + jiti@2.7.0: {} js-tokens@4.0.0: {} + js-tokens@9.0.1: {} + + jsdom@26.1.0: + dependencies: + cssstyle: 4.6.0 + data-urls: 5.0.0 + decimal.js: 10.6.0 + html-encoding-sniffer: 4.0.0 + http-proxy-agent: 7.0.2 + https-proxy-agent: 7.0.6 + is-potential-custom-element-name: 1.0.1 + nwsapi: 2.2.24 + parse5: 7.3.0 + rrweb-cssom: 0.8.0 + saxes: 6.0.0 + symbol-tree: 3.2.4 + tough-cookie: 5.1.2 + w3c-xmlserializer: 5.0.0 + webidl-conversions: 7.0.0 + whatwg-encoding: 3.1.1 + whatwg-mimetype: 4.0.0 + whatwg-url: 14.2.0 + ws: 8.21.0 + xml-name-validator: 5.0.0 + transitivePeerDependencies: + - bufferutil + - supports-color + - utf-8-validate + jsesc@3.1.0: {} json-schema@0.4.0: {} @@ -3005,10 +3568,16 @@ snapshots: dependencies: js-tokens: 4.0.0 + loupe@3.2.1: {} + + lru-cache@10.4.3: {} + lru-cache@5.1.1: dependencies: yallist: 3.1.1 + lz-string@1.5.0: {} + magic-string@0.30.21: dependencies: '@jridgewell/sourcemap-codec': 1.5.5 @@ -3389,12 +3958,18 @@ snapshots: transitivePeerDependencies: - supports-color + min-indent@1.0.1: {} + + monaco-editor@0.48.0: {} + ms@2.1.3: {} nanoid@3.3.15: {} node-releases@2.0.50: {} + nwsapi@2.2.24: {} + object-assign@4.1.1: {} oniguruma-to-es@2.3.0: @@ -3421,6 +3996,10 @@ snapshots: path-data-parser@0.1.0: {} + pathe@2.0.3: {} + + pathval@2.0.1: {} + picocolors@1.1.1: {} picomatch@4.0.4: {} @@ -3438,6 +4017,12 @@ snapshots: picocolors: 1.1.1 source-map-js: 1.2.1 + pretty-format@27.5.1: + dependencies: + ansi-regex: 5.0.1 + ansi-styles: 5.2.0 + react-is: 17.0.2 + prop-types@15.8.1: dependencies: loose-envify: 1.4.0 @@ -3446,6 +4031,8 @@ snapshots: property-information@7.2.0: {} + punycode@2.3.1: {} + re-resizable@6.11.2(react-dom@18.3.1(react@18.3.1))(react@18.3.1): dependencies: react: 18.3.1 @@ -3466,6 +4053,8 @@ snapshots: react-is@16.13.1: {} + react-is@17.0.2: {} + react-refresh@0.18.0: {} react-rnd@10.5.3(react-dom@18.3.1(react@18.3.1))(react@18.3.1): @@ -3480,6 +4069,11 @@ snapshots: dependencies: loose-envify: 1.4.0 + redent@3.0.0: + dependencies: + indent-string: 4.0.0 + strip-indent: 3.0.0 + regex-recursion@5.1.1: dependencies: regex: 5.1.1 @@ -3582,10 +4176,16 @@ snapshots: points-on-curve: 0.2.0 points-on-path: 0.2.1 + rrweb-cssom@0.8.0: {} + rw@1.3.3: {} safer-buffer@2.1.2: {} + saxes@6.0.0: + dependencies: + xmlchars: 2.2.0 + scheduler@0.23.2: dependencies: loose-envify: 1.4.0 @@ -3603,10 +4203,16 @@ snapshots: '@shikijs/vscode-textmate': 10.0.2 '@types/hast': 3.0.4 + siginfo@2.0.0: {} + source-map-js@1.2.1: {} space-separated-tokens@2.0.2: {} + stackback@0.0.2: {} + + std-env@3.10.0: {} + streamdown@2.5.0(react-dom@18.3.1(react@18.3.1))(react@18.3.1): dependencies: clsx: 2.1.1 @@ -3635,6 +4241,14 @@ snapshots: character-entities-html4: 2.1.0 character-entities-legacy: 3.0.0 + strip-indent@3.0.0: + dependencies: + min-indent: 1.0.1 + + strip-literal@3.1.0: + dependencies: + js-tokens: 9.0.1 + style-to-js@1.1.21: dependencies: style-to-object: 1.0.14 @@ -3651,7 +4265,7 @@ snapshots: react: 18.3.1 use-sync-external-store: 1.6.0(react@18.3.1) - tailwind-merge@2.6.1: {} + symbol-tree@3.2.4: {} tailwind-merge@3.6.0: {} @@ -3661,6 +4275,10 @@ snapshots: throttleit@2.1.0: {} + tinybench@2.9.0: {} + + tinyexec@0.3.2: {} + tinyexec@1.2.4: {} tinyglobby@0.2.17: @@ -3668,6 +4286,26 @@ snapshots: fdir: 6.5.0(picomatch@4.0.4) picomatch: 4.0.4 + tinypool@1.1.1: {} + + tinyrainbow@2.0.0: {} + + tinyspy@4.0.4: {} + + tldts-core@6.1.86: {} + + tldts@6.1.86: + dependencies: + tldts-core: 6.1.86 + + tough-cookie@5.1.2: + dependencies: + tldts: 6.1.86 + + tr46@5.1.1: + dependencies: + punycode: 2.3.1 + trim-lines@3.0.1: {} trough@2.2.0: {} @@ -3738,7 +4376,28 @@ snapshots: '@types/unist': 3.0.3 vfile-message: 4.0.3 - vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0): + vite-node@3.2.4(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0): + dependencies: + cac: 6.7.14 + debug: 4.4.3 + es-module-lexer: 1.7.0 + pathe: 2.0.3 + vite: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0) + transitivePeerDependencies: + - '@types/node' + - jiti + - less + - lightningcss + - sass + - sass-embedded + - stylus + - sugarss + - supports-color + - terser + - tsx + - yaml + + vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0): dependencies: esbuild: 0.28.1 fdir: 6.5.0(picomatch@4.0.4) @@ -3750,11 +4409,84 @@ snapshots: fsevents: 2.3.3 jiti: 2.7.0 lightningcss: 1.32.0 + yaml: 2.9.0 + + vitest@3.2.7(@types/debug@4.1.13)(jiti@2.7.0)(jsdom@26.1.0)(lightningcss@1.32.0)(yaml@2.9.0): + dependencies: + '@types/chai': 5.2.3 + '@vitest/expect': 3.2.7 + '@vitest/mocker': 3.2.7(vite@7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0)) + '@vitest/pretty-format': 3.2.7 + '@vitest/runner': 3.2.7 + '@vitest/snapshot': 3.2.7 + '@vitest/spy': 3.2.7 + '@vitest/utils': 3.2.7 + chai: 5.3.3 + debug: 4.4.3 + expect-type: 1.4.0 + magic-string: 0.30.21 + pathe: 2.0.3 + picomatch: 4.0.4 + std-env: 3.10.0 + tinybench: 2.9.0 + tinyexec: 0.3.2 + tinyglobby: 0.2.17 + tinypool: 1.1.1 + tinyrainbow: 2.0.0 + vite: 7.3.6(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0) + vite-node: 3.2.4(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0) + why-is-node-running: 2.3.0 + optionalDependencies: + '@types/debug': 4.1.13 + jsdom: 26.1.0 + transitivePeerDependencies: + - jiti + - less + - lightningcss + - msw + - sass + - sass-embedded + - stylus + - sugarss + - supports-color + - terser + - tsx + - yaml + + w3c-xmlserializer@5.0.0: + dependencies: + xml-name-validator: 5.0.0 web-namespaces@2.0.1: {} + webidl-conversions@7.0.0: {} + + whatwg-encoding@3.1.1: + dependencies: + iconv-lite: 0.6.3 + + whatwg-mimetype@4.0.0: {} + + whatwg-url@14.2.0: + dependencies: + tr46: 5.1.1 + webidl-conversions: 7.0.0 + + why-is-node-running@2.3.0: + dependencies: + siginfo: 2.0.0 + stackback: 0.0.2 + + ws@8.21.0: {} + + xml-name-validator@5.0.0: {} + + xmlchars@2.2.0: {} + yallist@3.1.1: {} + yaml@2.9.0: {} + zod@4.4.3: {} zwitch@2.0.4: {} diff --git a/pkg/cli/webapp/pnpm-workspace.yaml b/pkg/cli/webapp/pnpm-workspace.yaml new file mode 100644 index 0000000..2245205 --- /dev/null +++ b/pkg/cli/webapp/pnpm-workspace.yaml @@ -0,0 +1,5 @@ +packages: + - . + +minimumReleaseAge: 10080 +trustPolicy: no-downgrade diff --git a/pkg/cli/webapp/src/App.tsx b/pkg/cli/webapp/src/App.tsx index 442fa9c..337cd3e 100644 --- a/pkg/cli/webapp/src/App.tsx +++ b/pkg/cli/webapp/src/App.tsx @@ -1,6 +1,6 @@ -import { useMemo, type ReactNode } from "react"; +import { useCallback, useMemo, useState, useSyncExternalStore } from "react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { AppShell } from "@flanksource/clicky-ui/components"; +import { AppShell, type AppShellProps } from "@flanksource/clicky-ui/components"; import { RouterProvider, useBrowserRouter, @@ -10,16 +10,27 @@ import { EntityExplorerApp } from "@flanksource/clicky-ui/rpc"; import { apiClient } from "./api"; import { AgentLauncher } from "./AgentLauncher"; import { ChatLayer } from "./ChatLayer"; +import { + CommandPalette, + SearchTrigger, + useCommandPaletteShortcut, +} from "./CommandPalette"; import { ChatRoute } from "./ChatRoute"; import { HomeDashboard } from "./HomeDashboard"; import { PromptWorkbench } from "./PromptWorkbench"; import { SessionBrowser } from "./SessionBrowser"; +import { ShellActions } from "./shell"; +import { WhoamiPage } from "./WhoamiPage"; import { - CAPTAIN_SIDEBAR_COLLAPSE_KEY, - ShellActions, captainNavSections, + CAPTAIN_SIDEBAR_COLLAPSE_KEY, + getProjectScopeSnapshot, + setProjectScopeInLocation, + subscribeProjectScope, + withProjectScope, type PrimaryRoute, -} from "./shell"; +} from "./shellHelpers"; +import type { ProjectScope } from "./sessionData"; export function App() { const queryClient = useMemo( @@ -31,6 +42,24 @@ export function App() { ); const router = useBrowserRouter(); const route = parseRoute(router.pathname, window.location.search); + const projectScope = useSyncExternalStore( + subscribeProjectScope, + getProjectScopeSnapshot, + getProjectScopeSnapshot, + ); + + const setProjectScope = (scope: ProjectScope) => { + setProjectScopeInLocation(scope, router.navigate, router.pathname, window.location.search); + }; + const shellActions = ( + + ); + + const [paletteOpen, setPaletteOpen] = useState(false); + useCommandPaletteShortcut( + useCallback(() => setPaletteOpen((prev) => !prev), []), + ); + const shellSearch = setPaletteOpen(true)} />; return ( @@ -40,20 +69,31 @@ export function App() { } + navSections={captainNavSections("sessions", projectScope)} + actions={shellActions} + search={shellSearch} + projectScope={projectScope} /> ) : route.kind === "prompts" ? ( } + navSections={captainNavSections("prompts", projectScope)} + actions={shellActions} + search={shellSearch} /> ) : ( - + {route.kind === "dashboard" ? ( - + + ) : route.kind === "whoami" ? ( + ) : route.kind === "operations" ? ( )} + setPaletteOpen(false)} + onNavigate={router.navigate} + /> @@ -83,19 +128,26 @@ export function App() { function CaptainShell({ active, onNavigate, + projectScope, + actions, + search, children, }: { active: PrimaryRoute; onNavigate: (to: string, opts?: { replace?: boolean }) => void; - children: ReactNode; + projectScope: ProjectScope; + actions: AppShellProps["actions"]; + search: AppShellProps["search"]; + children: AppShellProps["children"]; }) { return ( } - navSections={captainNavSections(active)} + brand={} + navSections={captainNavSections(active, projectScope)} collapsedStorageKey={CAPTAIN_SIDEBAR_COLLAPSE_KEY} - actions={} + actions={actions} + search={search} contentClassName="p-0 overflow-hidden" > {children} @@ -105,14 +157,16 @@ function CaptainShell({ function ShellBrand({ onNavigate, + projectScope, }: { onNavigate: (to: string, opts?: { replace?: boolean }) => void; + projectScope: ProjectScope; }) { return ( @@ -124,12 +178,14 @@ type Route = | { kind: "agent" } | { kind: "sessions"; sessionId?: string } | { kind: "prompts"; promptId?: string } + | { kind: "whoami" } | { kind: "operations" } | { kind: "chat"; threadId: string; model?: string }; function primaryRoute(route: Route): PrimaryRoute { if (route.kind === "dashboard") return "dashboard"; if (route.kind === "operations") return "operations"; + if (route.kind === "whoami") return "whoami"; if (route.kind === "prompts") return "prompts"; if (route.kind === "sessions") return "sessions"; return "agent"; @@ -137,6 +193,7 @@ function primaryRoute(route: Route): PrimaryRoute { function parseRoute(pathname: string, search: string): Route { if (pathname.startsWith("/operations")) return { kind: "operations" }; + if (pathname.startsWith("/whoami")) return { kind: "whoami" }; if (pathname.startsWith("/prompts")) { const raw = pathname.slice("/prompts".length).replace(/^\/+/, ""); const promptId = raw ? decodeURIComponent(raw.split("/")[0] ?? "") : undefined; diff --git a/pkg/cli/webapp/src/ChatLayer.tsx b/pkg/cli/webapp/src/ChatLayer.tsx index 0bf49d5..0a51ff5 100644 --- a/pkg/cli/webapp/src/ChatLayer.tsx +++ b/pkg/cli/webapp/src/ChatLayer.tsx @@ -22,9 +22,8 @@ export function ChatLayer() { chat={{ api: "/api/chat", modelsApi: "/api/chat/models", - defaultModel: "claude-agent-sonnet", + defaultModel: "claude-sonnet-5", enableAttachments: true, - toolApproval: "manual", suggestions: [ "Summarize this run", "Show recent changed files", diff --git a/pkg/cli/webapp/src/CommandPalette.test.tsx b/pkg/cli/webapp/src/CommandPalette.test.tsx new file mode 100644 index 0000000..706ac3d --- /dev/null +++ b/pkg/cli/webapp/src/CommandPalette.test.tsx @@ -0,0 +1,260 @@ +import { useCallback, useState } from "react"; +import { + cleanup, + fireEvent, + render, + screen, + waitFor, +} from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + CommandPalette, + useCommandPaletteShortcut, +} from "./CommandPalette"; +import { directSessionId } from "./commandPaletteHelpers"; +import type { SessionRecord } from "./sessionData"; +import type { PromptSummary } from "./promptData"; + +const fetchSessionSearch = vi.hoisted(() => vi.fn()); +const fetchPromptList = vi.hoisted(() => vi.fn()); + +vi.mock("./sessionData", async (importOriginal) => ({ + ...(await importOriginal()), + fetchSessionSearch, +})); + +vi.mock("./promptData", async (importOriginal) => ({ + ...(await importOriginal()), + fetchPromptList, + resolvePromptOps: () => ({ list: { path: "/api/v1/prompt", method: "GET" } }), +})); + +vi.mock("@flanksource/clicky-ui/rpc", async (importOriginal) => ({ + ...(await importOriginal()), + useOperations: () => ({ operations: [], isLoading: false, error: null }), +})); + +// Fixtures use values distinct from anything the component derives, so an +// assertion can only pass if the component actually plumbed them through. +const SESSION_KEY = "055781c7-360a-4eb2-80be-452b3937fcfe"; +const SESSION_ID = "019f7c25-9adf-7901-add9-8c46693472fb"; +const PROMPT_ID = "local:review-diff"; + +function session(overrides: Partial = {}): SessionRecord { + return { + key: SESSION_KEY, + id: SESSION_ID, + source: "claude", + title: "Refactor the billing importer", + project: "/home/dev/acme/billing", + toolCalls: 0, + messages: 0, + ...overrides, + } as SessionRecord; +} + +function prompt(overrides: Partial = {}): PromptSummary { + return { + id: PROMPT_ID, + name: "review-diff", + sourceKind: "local", + sourceId: "local", + source: "local", + path: "/home/dev/prompts/review-diff.prompt", + relPath: "prompts/review-diff.prompt", + writable: true, + ...overrides, + } as PromptSummary; +} + +function renderPalette(onNavigate = vi.fn(), onClose = vi.fn()) { + const client = new QueryClient({ + defaultOptions: { queries: { retry: false, gcTime: 0 } }, + }); + render( + + + , + ); + return { onNavigate, onClose }; +} + +async function type(value: string) { + const input = screen.getByLabelText("Search sessions and prompts"); + fireEvent.change(input, { target: { value } }); + return input; +} + +beforeEach(() => { + fetchSessionSearch.mockReset(); + fetchPromptList.mockReset(); + fetchSessionSearch.mockResolvedValue({ sessions: [], total: 0 }); + fetchPromptList.mockResolvedValue([]); +}); + +afterEach(cleanup); + +describe("directSessionId", () => { + it("accepts a token long enough to be an unambiguous id prefix", () => { + // Identity resolution takes Captain UUIDs and provider-id prefixes, so + // this must not be gated on a UUID shape. + expect(directSessionId(SESSION_ID)).toBe(SESSION_ID); + expect(directSessionId("a1b2c3d4")).toBe("a1b2c3d4"); + expect(directSessionId(` ${SESSION_KEY} `)).toBe(SESSION_KEY); + }); + + it("rejects prose and tokens too short to disambiguate", () => { + expect(directSessionId("billing importer")).toBeNull(); + expect(directSessionId("a1b2c3")).toBeNull(); + expect(directSessionId("")).toBeNull(); + }); +}); + +describe("useCommandPaletteShortcut", () => { + function Harness() { + const [count, setCount] = useState(0); + const toggle = useCallback(() => setCount((c) => c + 1), []); + useCommandPaletteShortcut(toggle); + return {count}; + } + + it("fires on Cmd+K and Ctrl+K but not on a bare k or Alt+Cmd+K", () => { + render(); + const count = () => screen.getByTestId("count").textContent; + + fireEvent.keyDown(window, { key: "k", metaKey: true }); + expect(count()).toBe("1"); + + fireEvent.keyDown(window, { key: "k", ctrlKey: true }); + expect(count()).toBe("2"); + + fireEvent.keyDown(window, { key: "k" }); + fireEvent.keyDown(window, { key: "k", metaKey: true, altKey: true }); + expect(count()).toBe("2"); + }); +}); + +describe("CommandPalette", () => { + it("prompts for input before any query is typed", () => { + renderPalette(); + expect( + screen.getByText(/Type to search sessions and prompts/), + ).toBeInTheDocument(); + expect(fetchSessionSearch).not.toHaveBeenCalled(); + }); + + it("searches all projects rather than the active project scope", async () => { + renderPalette(); + await type("billing"); + await waitFor(() => + expect(fetchSessionSearch).toHaveBeenCalledWith({ + query: "billing", + limit: 20, + }), + ); + }); + + it("opens the highlighted session on Enter using its Captain UUID key", async () => { + fetchSessionSearch.mockResolvedValue({ sessions: [session()], total: 1 }); + const { onNavigate, onClose } = renderPalette(); + const input = await type("billing"); + + await screen.findByText("Refactor the billing importer"); + fireEvent.keyDown(input, { key: "Enter" }); + + expect(onNavigate).toHaveBeenCalledWith( + `/sessions/${encodeURIComponent(SESSION_KEY)}`, + ); + expect(onClose).toHaveBeenCalled(); + }); + + it("opens a prompt result at its prompt route", async () => { + fetchPromptList.mockResolvedValue([prompt()]); + const { onNavigate } = renderPalette(); + const input = await type("review"); + + await screen.findByText("review-diff"); + fireEvent.keyDown(input, { key: "Enter" }); + + expect(onNavigate).toHaveBeenCalledWith( + `/prompts/${encodeURIComponent(PROMPT_ID)}`, + ); + }); + + it("wraps arrow navigation across the session and prompt groups", async () => { + fetchSessionSearch.mockResolvedValue({ sessions: [session()], total: 1 }); + fetchPromptList.mockResolvedValue([prompt()]); + const { onNavigate } = renderPalette(); + const input = await type("review"); + + await screen.findByText("review-diff"); + + // Two rows total, starting on the session. ArrowDown crosses into the + // prompt group... + fireEvent.keyDown(input, { key: "ArrowDown" }); + fireEvent.keyDown(input, { key: "Enter" }); + expect(onNavigate).toHaveBeenCalledWith( + `/prompts/${encodeURIComponent(PROMPT_ID)}`, + ); + + // ...and a further ArrowDown wraps past the end back to the session. + onNavigate.mockClear(); + fireEvent.keyDown(input, { key: "ArrowDown" }); + fireEvent.keyDown(input, { key: "Enter" }); + expect(onNavigate).toHaveBeenCalledWith( + `/sessions/${encodeURIComponent(SESSION_KEY)}`, + ); + + // ArrowUp wraps in the opposite direction, off the first row to the last. + onNavigate.mockClear(); + fireEvent.keyDown(input, { key: "ArrowUp" }); + fireEvent.keyDown(input, { key: "Enter" }); + expect(onNavigate).toHaveBeenCalledWith( + `/prompts/${encodeURIComponent(PROMPT_ID)}`, + ); + }); + + it("offers a direct-open row that defers resolution to the session route", async () => { + const { onNavigate } = renderPalette(); + const input = await type(SESSION_ID); + + await screen.findByText("Open session by id"); + fireEvent.keyDown(input, { key: "Enter" }); + + expect(onNavigate).toHaveBeenCalledWith( + `/sessions/${encodeURIComponent(SESSION_ID)}`, + ); + }); + + it("does not offer a direct-open row for a multi-word query", async () => { + renderPalette(); + await type("billing importer"); + await waitFor(() => expect(fetchSessionSearch).toHaveBeenCalled()); + expect(screen.queryByText("Open session by id")).not.toBeInTheDocument(); + }); + + it("caps each group and reports the remainder", async () => { + const many = Array.from({ length: 11 }, (_, i) => + session({ key: `key-${i}`, id: `id-${i}`, title: `Session ${i}` }), + ); + fetchSessionSearch.mockResolvedValue({ sessions: many, total: many.length }); + renderPalette(); + await type("session"); + + await screen.findByText("Session 0"); + expect(screen.getByText("Session 7")).toBeInTheDocument(); + expect(screen.queryByText("Session 8")).not.toBeInTheDocument(); + expect(screen.getByText("+3 more")).toBeInTheDocument(); + }); + + it("reports when nothing matches", async () => { + renderPalette(); + // Multi-word so no direct-open row is offered; otherwise the palette always + // has at least one row and the empty state is unreachable. + await type("no such thing"); + expect( + await screen.findByText(/No sessions or prompts match/), + ).toBeInTheDocument(); + }); +}); diff --git a/pkg/cli/webapp/src/CommandPalette.tsx b/pkg/cli/webapp/src/CommandPalette.tsx new file mode 100644 index 0000000..c0d3121 --- /dev/null +++ b/pkg/cli/webapp/src/CommandPalette.tsx @@ -0,0 +1,406 @@ +import { + useEffect, + useMemo, + useRef, + useState, + type KeyboardEvent, + type ReactNode, +} from "react"; +import { useQuery } from "@tanstack/react-query"; +import { Modal } from "@flanksource/clicky-ui/components"; +import { + UiCode2, + UiSearch, + UiTerminal, + type IconComponent, +} from "@flanksource/clicky-ui/icons"; +import { useOperations } from "@flanksource/clicky-ui/rpc"; +import { apiClient } from "./api"; +import { + fetchSessionSearch, + projectLabel, + sessionTitle, + shortID, + type SessionRecord, +} from "./sessionData"; +import { + fetchPromptList, + resolvePromptOps, + type PromptSummary, +} from "./promptData"; +import { + directSessionId, + paletteShortcutLabel, +} from "./commandPaletteHelpers"; + +// The ⌘K palette is captain's single global search: it spans sessions and +// prompts regardless of the active route and jumps to the chosen item rather +// than filtering a list in place. The per-page SearchInput on /sessions still +// owns scoped, in-place filtering; this owns "find that one thing". +// +// Session results deliberately ignore the app-bar project scope so a session is +// findable from anywhere, and the direct-open row hands the raw text to +// /sessions/{id}, whose RunSessionGet identity resolution is +// the same code path as `captain sessions get`. + +// GROUP_CAP keeps each section short so the list stays scannable; overflow is +// surfaced as a "+N more" hint rather than silently dropped. +const GROUP_CAP = 8; +const DEBOUNCE_MS = 200; + +// Session identity resolution accepts a full Captain UUID or a +// provider-session-id prefix, so the direct-open row cannot gate on a UUID +// shape. Anything token-shaped and long enough to be unambiguous counts. +/** + * Binds ⌘K / Ctrl+K to `toggle` for the lifetime of the caller. Works even while + * a field is focused — it isn't a text-editing shortcut. `SearchInput`'s own + * `onShortcut` cannot serve here: that listener lives in `InputField` and only + * runs while the field is mounted, which a closed palette is not. + */ +export function useCommandPaletteShortcut(toggle: () => void) { + useEffect(() => { + const onKey = (event: globalThis.KeyboardEvent) => { + if ( + (event.metaKey || event.ctrlKey) && + !event.altKey && + event.key.toLowerCase() === "k" + ) { + event.preventDefault(); + toggle(); + } + }; + window.addEventListener("keydown", onKey); + return () => window.removeEventListener("keydown", onKey); + }, [toggle]); +} + +interface Row { + key: string; + icon: IconComponent; + title: string; + subtitle: string; + meta: string; + onSelect: () => void; +} + +/** + * Top-bar affordance that stands in for an inline search box: a click (or the + * ⌘K/Ctrl+K shortcut) opens the palette. Rendered on every route so global + * search is always one key away. + */ +export function SearchTrigger({ onOpen }: { onOpen: () => void }) { + return ( + + ); +} + +function CommandPaletteContent({ + open, + onClose, + onNavigate, +}: { + open: boolean; + onClose: () => void; + onNavigate: (to: string, opts?: { replace?: boolean }) => void; +}) { + const [query, setQuery] = useState(""); + const [debounced, setDebounced] = useState(""); + const [active, setActive] = useState(0); + const inputRef = useRef(null); + const listRef = useRef(null); + + // Move focus to the input when this freshly keyed palette mounts. The Modal focuses its own panel in a + // passive effect whose flush time varies (opening cascades state updates), so a + // single deferred focus keeps losing the race. Instead retry briefly until the + // input holds focus, then stop — there is no focus trap, so once it lands it + // stays. This self-corrects regardless of when the Modal grabs focus. + useEffect(() => { + if (!open) return; + let tries = 0; + const timer = setInterval(() => { + const el = inputRef.current; + if (el && document.activeElement !== el) el.focus(); + if (++tries >= 10 || (el && document.activeElement === el)) + clearInterval(timer); + }, 30); + return () => clearInterval(timer); + }, [open]); + + // Results are server-side, so hold off a beat rather than firing a request per + // keystroke. + useEffect(() => { + const timer = setTimeout(() => setDebounced(query.trim()), DEBOUNCE_MS); + return () => clearTimeout(timer); + }, [query]); + + const enabled = open && debounced.length > 0; + + const sessionsQuery = useQuery({ + queryKey: ["palette-sessions", debounced], + queryFn: () => fetchSessionSearch({ query: debounced, limit: 20 }), + enabled, + }); + + const { operations } = useOperations(apiClient); + const promptOps = useMemo(() => resolvePromptOps(operations), [operations]); + const promptsQuery = useQuery({ + queryKey: ["palette-prompts", promptOps.list?.path, debounced], + queryFn: () => + fetchPromptList(promptOps.list!, { source: "all", query: debounced }), + enabled: enabled && Boolean(promptOps.list), + }); + + const sessions = sessionsQuery.data?.sessions ?? EMPTY_SESSIONS; + const prompts = promptsQuery.data ?? EMPTY_PROMPTS; + const directId = directSessionId(query); + + // One flat, ordered list (direct id, sessions, then prompts) backs keyboard + // navigation; the rendered groups index into it so the highlight and Enter + // stay in sync. + const rows = useMemo(() => { + const out: Row[] = []; + if (directId) { + out.push({ + key: `id:${directId}`, + icon: UiSearch, + title: "Open session by id", + subtitle: directId, + meta: "id", + onSelect: () => { + onClose(); + onNavigate(`/sessions/${encodeURIComponent(directId)}`); + }, + }); + } + for (const session of sessions.slice(0, GROUP_CAP)) { + out.push({ + key: `session:${session.key}`, + icon: UiTerminal, + title: sessionTitle(session), + subtitle: shortID(session.id) || session.key, + meta: session.project ? projectLabel(session.project) : session.source, + onSelect: () => { + onClose(); + onNavigate(`/sessions/${encodeURIComponent(session.key)}`); + }, + }); + } + for (const prompt of prompts.slice(0, GROUP_CAP)) { + out.push({ + key: `prompt:${prompt.id}`, + icon: UiCode2, + title: prompt.name, + subtitle: prompt.relPath || prompt.path, + meta: prompt.sourceKind, + onSelect: () => { + onClose(); + onNavigate(`/prompts/${encodeURIComponent(prompt.id)}`); + }, + }); + } + return out; + }, [directId, sessions, prompts, onClose, onNavigate]); + + // Keep the active index in range as results change while typing. + useEffect(() => { + setActive((a) => (rows.length === 0 ? 0 : Math.min(a, rows.length - 1))); + }, [rows.length]); + + // Scroll the highlighted row into view as the selection moves. + useEffect(() => { + listRef.current + ?.querySelector(`[data-row="${active}"]`) + ?.scrollIntoView({ block: "nearest" }); + }, [active]); + + function onKeyDown(event: KeyboardEvent) { + if (rows.length === 0) return; + if (event.key === "ArrowDown") { + event.preventDefault(); + setActive((a) => (a + 1) % rows.length); + } else if (event.key === "ArrowUp") { + event.preventDefault(); + setActive((a) => (a - 1 + rows.length) % rows.length); + } else if (event.key === "Enter") { + event.preventDefault(); + rows[active]?.onSelect(); + } + } + + const directBase = directId ? 1 : 0; + const promptBase = directBase + Math.min(sessions.length, GROUP_CAP); + const loading = sessionsQuery.isFetching || promptsQuery.isFetching; + + return ( + +
+ + { + setQuery(event.target.value); + setActive(0); + }} + onKeyDown={onKeyDown} + placeholder="Search sessions, prompts, or paste a session id…" + aria-label="Search sessions and prompts" + className="min-w-0 flex-1 bg-transparent text-sm outline-none placeholder:text-muted-foreground" + /> + + esc + +
+ +
+ {!query.trim() ? ( +
+ Type to search sessions and prompts, or paste a session id. +
+ ) : rows.length === 0 && !loading ? ( +
+ No sessions or prompts match “{query.trim()}”. +
+ ) : ( + <> + {directId && ( + + + + )} + {sessions.length > 0 && ( + + {sessions.slice(0, GROUP_CAP).map((_, i) => ( + + ))} + + )} + {prompts.length > 0 && ( + + {prompts.slice(0, GROUP_CAP).map((_, i) => ( + + ))} + + )} + {loading && ( +
+ Searching… +
+ )} + + )} +
+
+ ); +} + +export function CommandPalette(props: { + open: boolean; + onClose: () => void; + onNavigate: (to: string, opts?: { replace?: boolean }) => void; +}) { + return ; +} + +function Group({ + label, + overflow, + children, +}: { + label: string; + overflow: number; + children: ReactNode; +}) { + return ( +
+
+ {label} + {overflow > 0 && ( + +{overflow} more + )} +
+ {children} +
+ ); +} + +function PaletteRow({ + row, + index, + active, + onHover, +}: { + row: Row; + index: number; + active: number; + onHover: (index: number) => void; +}) { + const isActive = index === active; + const Icon = row.icon; + return ( + + ); +} + +const EMPTY_SESSIONS: SessionRecord[] = []; +const EMPTY_PROMPTS: PromptSummary[] = []; diff --git a/pkg/cli/webapp/src/HomeDashboard.tsx b/pkg/cli/webapp/src/HomeDashboard.tsx index cd530d3..6c990ce 100644 --- a/pkg/cli/webapp/src/HomeDashboard.tsx +++ b/pkg/cli/webapp/src/HomeDashboard.tsx @@ -1,34 +1,31 @@ -import { useMemo, useState, type ComponentProps, type KeyboardEvent, type ReactNode } from "react"; +import { useMemo, useState, type ReactNode } from "react"; import { useQuery } from "@tanstack/react-query"; import { Button, SearchInput, SegmentedControl, - Switch, } from "@flanksource/clicky-ui/components"; import { Icon, - ProgressBars, + TimeseriesPanel, UiActivity, - UiArrowDown, - UiArrowUp, - UiBrain, UiChartBar, UiChip, UiClock, - UiCopy, UiHistory, UiMemoryStick, UiRefresh, UiRobotAi, - UiSparkles, UiTerminal, + type TimeseriesResponse, + type TimeseriesSeries, } from "@flanksource/clicky-ui/data"; import { SOURCE_OPTIONS, commandLabel, errorMessage, fetchLiveSessions, + fetchSessionThroughput, formatCompactNumber, formatCost, formatTime, @@ -38,29 +35,42 @@ import { sessionSortTime, sessionTitle, type SessionDashboard, - type SessionLive, type SessionRecord, + type SessionThroughputGroup, + type SessionThroughputResult, + type ProjectScope, type SourceFilter, } from "./sessionData"; +import { + ContextCell, + ProjectGroupHeader, + SessionActions, + SessionIdentity, + SessionTable, + UsageBarsCell, + type DashboardSort, + type ProjectSessionGroup, + type SessionIcon, + type SortDirection, +} from "./SessionTable"; +import { + compareSessions, + defaultSortDirection, + effortLabel, + formatRelativeTime, + groupSessionsByProject, + hasLiveProcess, + modelIcon, +} from "./sessionTableHelpers"; +import { withProjectScope } from "./shellHelpers"; type DashboardView = "list" | "cards"; -type DashboardSort = "model" | "health" | "context" | "cpu" | "memory" | "tokens" | "recent"; -type SortDirection = "asc" | "desc"; type Navigate = (to: string, opts?: { replace?: boolean }) => void; -type DashboardIcon = NonNullable["icon"]>; -type LiveSessionRecord = SessionRecord & { live: SessionLive }; -type ProjectSessionGroup = { - key: string; - label: string; - detail?: string; - sessions: LiveSessionRecord[]; -}; - -const PERCENT_UNIT = { - perBar: 25, - label: "%", - barLabel: "25%", - format: (units: number) => `${Math.round(units * 25)}`, +type DashboardIcon = SessionIcon; +type ThroughputMetric = "outputTokensPerSecond" | "contextTokensPerSecond"; +type ThroughputChartModel = { + series: TimeseriesSeries[]; + responses: Record; }; const VIEW_OPTIONS = [ @@ -78,36 +88,35 @@ const SORT_OPTIONS = [ { id: "recent", label: "Recent" }, ] satisfies Array<{ id: DashboardSort; label: string }>; -const SESSION_GRID_CLASS = - "grid grid-cols-[minmax(12rem,1.4fr)_5.25rem_6.25rem] sm:grid-cols-[minmax(13rem,1.5fr)_5.25rem_6.25rem_6.25rem] lg:grid-cols-[minmax(15rem,1.6fr)_5.5rem_7rem_7rem_7rem_6rem_7rem_5.5rem]"; - -const SESSION_COLUMNS = [ - { label: "Model", sort: "model" }, - { label: "Status", sort: "health" }, - { label: "CPU", sort: "cpu" }, - { label: "Memory", sort: "memory" }, - { label: "Context", sort: "context" }, - { label: "Tokens", sort: "tokens" }, - { label: "Updated", sort: "recent" }, - { label: "Actions" }, -] satisfies Array<{ label: string; sort?: DashboardSort }>; +const EMPTY_SESSIONS: SessionRecord[] = []; +const EMPTY_THROUGHPUT_GROUPS: SessionThroughputGroup[] = []; -export function HomeDashboard({ onNavigate }: { onNavigate: Navigate }) { +export function HomeDashboard({ + onNavigate, + projectScope, +}: { + onNavigate: Navigate; + projectScope: ProjectScope; +}) { const [source, setSource] = useState("all"); - const [allProjects, setAllProjects] = useState(false); const [query, setQuery] = useState(""); const [view, setView] = useState("list"); const [sort, setSort] = useState("health"); const [sortDirection, setSortDirection] = useState("desc"); const sessionsQuery = useQuery({ - queryKey: ["sessions-dashboard", source, allProjects, query], - queryFn: () => fetchLiveSessions({ source, allProjects, query, limit: 200 }), + queryKey: ["sessions-dashboard", source, projectScope, query], + queryFn: () => fetchLiveSessions({ source, project: projectScope, query, limit: 200 }), refetchInterval: 5000, refetchIntervalInBackground: false, }); + const throughputQuery = useQuery({ + queryKey: ["sessions-throughput", source, projectScope, query], + queryFn: () => fetchSessionThroughput({ source, project: projectScope, query, limit: 500 }), + refetchInterval: false, + }); - const sessions = sessionsQuery.data?.sessions ?? []; + const sessions = sessionsQuery.data?.sessions ?? EMPTY_SESSIONS; const liveSessions = useMemo( () => sessions @@ -137,7 +146,7 @@ export function HomeDashboard({ onNavigate }: { onNavigate: Navigate }) { const openSession = (session: SessionRecord) => { if (session.detailAvailable === false) return; - onNavigate(`/sessions/${encodeURIComponent(session.key)}`); + onNavigate(withProjectScope(`/sessions/${encodeURIComponent(session.key)}`, projectScope)); }; const selectSort = (nextSort: DashboardSort) => { @@ -158,16 +167,17 @@ export function HomeDashboard({ onNavigate }: { onNavigate: Navigate }) { void sessionsQuery.refetch()} + loading={sessionsQuery.isFetching || throughputQuery.isFetching} + onRefresh={() => { + void sessionsQuery.refetch(); + void throughputQuery.refetch(); + }} onNewAgent={() => onNavigate("/agent")} /> @@ -183,6 +193,11 @@ export function HomeDashboard({ onNavigate }: { onNavigate: Navigate }) { liveCount={liveSessions.length} loading={sessionsQuery.isLoading} /> +
@@ -223,8 +238,6 @@ export function HomeDashboard({ onNavigate }: { onNavigate: Navigate }) { function DashboardToolbar({ source, onSourceChange, - allProjects, - onAllProjectsChange, query, onQueryChange, view, @@ -237,8 +250,6 @@ function DashboardToolbar({ }: { source: SourceFilter; onSourceChange: (source: SourceFilter) => void; - allProjects: boolean; - onAllProjectsChange: (enabled: boolean) => void; query: string; onQueryChange: (query: string) => void; view: DashboardView; @@ -270,7 +281,7 @@ function DashboardToolbar({
-
+
- void; - onOpen: (session: SessionRecord) => void; + result?: SessionThroughputResult; + loading: boolean; + error: unknown; }) { + const groups = result?.groups ?? EMPTY_THROUGHPUT_GROUPS; + const outputChart = useMemo( + () => buildThroughputChart(groups, "outputTokensPerSecond"), + [groups], + ); + const contextChart = useMemo( + () => buildThroughputChart(groups, "contextTokensPerSecond"), + [groups], + ); + return ( -
-
- {SESSION_COLUMNS.map((column) => ( - - ))} -
-
- {groups.map((group) => ( -
- -
- {group.sessions.map((session) => ( -
activateSession(session, onOpen)} - onKeyDown={(event) => handleSessionKeyDown(event, session, onOpen)} - className={`${SESSION_GRID_CLASS} cursor-pointer items-center gap-density-2 px-density-3 py-density-2 text-sm transition-colors hover:bg-muted/40 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring data-[disabled=true]:cursor-default data-[disabled=true]:hover:bg-transparent`} - data-disabled={session.detailAvailable === false ? "true" : undefined} - > - - - - - - - - -
- ))} -
-
- ))} +
+
+
+
Token Throughput
+
+ {loading ? "Loading completed sessions..." : `${result?.total ?? 0} completed sessions sampled`} +
+
+ {result?.skipped ? ( +
{result.skipped} skipped
+ ) : null}
-
+ + {error ? ( +
+ {errorMessage(error)} +
+ ) : loading ? ( +
+ Loading throughput... +
+ ) : groups.length === 0 ? ( +
+ No completed sessions with duration and token usage. +
+ ) : ( + <> +
+ + +
+ + + )} + ); } -function SessionHeaderCell({ - column, - sort, - sortDirection, - onSortChange, +function StaticThroughputChart({ + title, + icon, + model, + empty, }: { - column: (typeof SESSION_COLUMNS)[number]; - sort: DashboardSort; - sortDirection: SortDirection; - onSortChange: (sort: DashboardSort) => void; + title: string; + icon: DashboardIcon; + model: ThroughputChartModel; + empty: string; }) { - if (!column.sort) { - return
{column.label}
; + const fetcher = useMemo( + () => staticThroughputFetcher(model.responses), + [model.responses], + ); + if (model.series.length === 0) { + return ( +
+ {empty} +
+ ); } + return ( + + ); +} - const active = sort === column.sort; +function ThroughputTable({ groups }: { groups: SessionThroughputGroup[] }) { return ( - +
+
+ + + + + + + + + + + + + {groups.map((group) => ( + + + + + + + + + ))} + +
ModelSessionsOutputTotalContextContext Used
+
+ + + + + {group.model} + + {group.source} / {effortLabel(group.reasoningEffort) ?? "default"} + + +
+
{group.sessions}{formatRate(group.outputTokensPerSecond)}{formatRate(group.totalTokensPerSecond)}{formatOptionalRate(group.contextTokensPerSecond)}{formatOptionalPercent(group.avgContextUsedPercent)}
+
+
); } @@ -521,8 +573,8 @@ function SessionCardGrid({ key={session.key} role={session.detailAvailable === false ? undefined : "button"} tabIndex={session.detailAvailable === false ? undefined : 0} - onClick={() => activateSession(session, onOpen)} - onKeyDown={(event) => handleSessionKeyDown(event, session, onOpen)} + onClick={() => session.detailAvailable !== false && onOpen(session)} + onKeyDown={(event) => handleCardKeyDown(event, session, onOpen)} className="min-w-0 cursor-pointer rounded border border-border bg-card p-density-3 transition-colors hover:border-muted-foreground/40 hover:bg-muted/30 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring data-[disabled=true]:cursor-default data-[disabled=true]:hover:border-border data-[disabled=true]:hover:bg-card" data-disabled={session.detailAvailable === false ? "true" : undefined} > @@ -531,7 +583,7 @@ function SessionCardGrid({
- + @@ -552,7 +604,7 @@ function SessionCardGrid({ value={ @@ -564,9 +616,11 @@ function SessionCardGrid({
-
- {commandLabel(session.live.command)} -
+ {session.live?.command ? ( +
+ {commandLabel(session.live.command)} +
+ ) : null}
))}
@@ -576,168 +630,15 @@ function SessionCardGrid({ ); } -function ProjectGroupHeader({ group }: { group: ProjectSessionGroup }) { - return ( -
-
-
{group.label}
- {group.detail ? ( -
{group.detail}
- ) : null} -
-
- {group.sessions.length} -
-
- ); -} - -function UsageBarsCell({ - title, - value, - icon, - thresholds, -}: { - title: string; - value: number | undefined; - icon: DashboardIcon; - thresholds: [warning: number, danger: number]; -}) { - return ( -
- - - {formatPercent(value)} - -
- ); -} - -function SessionIdentity({ session }: { session: SessionRecord }) { - const model = modelLabel(session); - const effort = effortLabel(session.reasoningEffort); - return ( -
-
- - - - {model} - - {session.source} - -
-
- {effort ? ( - - - {effort} - - ) : null} - {sessionTitle(session)} -
-
- {session.live?.pid ? `pid ${session.live.pid} - ` : ""} - {commandLabel(session.live?.command)} -
-
- ); -} - -function StatusCell({ session }: { session: SessionRecord & { live: SessionLive } }) { - const signal = session.health?.[0]; - return ( -
-
- - {session.live.status ?? "active"} -
-
- ); -} - -function ContextCell({ - session, - expanded = false, -}: { - session: SessionRecord; - expanded?: boolean; -}) { - const percent = session.context?.freePercent; - if (percent === undefined) { - return ; - } - return ( -
-
- {percent}% free - {expanded && session.context?.windowTokens ? ( - - {formatCompactNumber(session.context.windowTokens)} - - ) : null} -
-
-
-
-
- ); -} - -function SessionActions({ - session, - onOpen, - compact = false, -}: { - session: SessionRecord; - onOpen: (session: SessionRecord) => void; - compact?: boolean; -}) { - return ( -
- - -
- ); +function handleCardKeyDown( + event: React.KeyboardEvent, + session: SessionRecord, + onOpen: (session: SessionRecord) => void, +) { + if (event.key !== "Enter" && event.key !== " ") return; + if (session.detailAvailable === false) return; + event.preventDefault(); + onOpen(session); } function HealthPanel({ @@ -870,173 +771,82 @@ function MetricBox({ label, value }: { label: string; value: ReactNode }) { ); } -function MetricText({ value }: { value: ReactNode }) { - return
{value}
; -} - -function hasLiveProcess(session: SessionRecord): session is LiveSessionRecord { - return Boolean(session.live); -} - -function compareSessions( - left: LiveSessionRecord, - right: LiveSessionRecord, - sort: DashboardSort, - direction: SortDirection, -) { - if (sort === "model") { - return directionalCompare( - modelLabel(left).localeCompare(modelLabel(right)) || - sessionSortTime(right) - sessionSortTime(left), - direction, - "asc", - ); - } - if (sort === "context") { - return directionalCompare( - (left.context?.freePercent ?? 101) - (right.context?.freePercent ?? 101) || - sessionSortTime(right) - sessionSortTime(left), - direction, - "asc", - ); - } - if (sort === "cpu") { - return directionalCompare( - (right.live.cpuPercent ?? -1) - (left.live.cpuPercent ?? -1) || - sessionSortTime(right) - sessionSortTime(left), - direction, - "desc", - ); - } - if (sort === "memory") { - return directionalCompare( - (right.live.memoryPercent ?? -1) - (left.live.memoryPercent ?? -1) || - sessionSortTime(right) - sessionSortTime(left), - direction, - "desc", - ); - } - if (sort === "tokens") { - return directionalCompare( - tokenTotal(right) - tokenTotal(left) || - sessionSortTime(right) - sessionSortTime(left), - direction, - "desc", - ); - } - if (sort === "recent") { - return directionalCompare( - sessionSortTime(right) - sessionSortTime(left), - direction, - "desc", - ); - } - return directionalCompare( - healthRank(right) - healthRank(left) || - (left.context?.freePercent ?? 101) - (right.context?.freePercent ?? 101) || - (right.live.cpuPercent ?? -1) - (left.live.cpuPercent ?? -1) || - sessionSortTime(right) - sessionSortTime(left), - direction, - "desc", - ); -} - -function directionalCompare(value: number, direction: SortDirection, natural: SortDirection) { - return direction === natural ? value : -value; -} - -function defaultSortDirection(sort: DashboardSort): SortDirection { - return sort === "model" || sort === "context" ? "asc" : "desc"; -} - -function tokenTotal(session: SessionRecord) { - return session.tokens?.totalTokens ?? 0; -} - -function groupSessionsByProject(sessions: LiveSessionRecord[]): ProjectSessionGroup[] { - const groups = new Map(); - - for (const session of sessions) { - const cwd = session.live.cwd ?? session.cwd; - const key = cwd || "unknown"; - const label = projectLabel(cwd); - const existing = groups.get(key); - - if (existing) { - existing.sessions.push(session); - continue; +function buildThroughputChart( + groups: SessionThroughputGroup[], + metric: ThroughputMetric, +): ThroughputChartModel { + const responses: Record = {}; + const series: TimeseriesSeries[] = []; + const ranked = [...groups] + .filter((group) => (group.points?.length ?? 0) >= 2 && (group[metric] ?? 0) > 0) + .sort((left, right) => (right[metric] ?? 0) - (left[metric] ?? 0)) + .slice(0, 5); + + for (const [index, group] of ranked.entries()) { + const id = `${metric}-${index}-${hashThroughputGroup(group)}`; + const points = []; + for (const point of group.points ?? []) { + const value = point[metric] ?? 0; + if (!Number.isFinite(value)) continue; + points.push({ at: point.at, value }); } - - groups.set(key, { - key, - label, - detail: cwd && cwd !== label ? cwd : undefined, - sessions: [session], + if (points.length < 2) continue; + series.push({ + id, + label: throughputGroupLabel(group), + unit: " tok/s", }); + responses[id] = { id, points }; } - return [...groups.values()]; -} - -function activateSession(session: SessionRecord, onOpen: (session: SessionRecord) => void) { - if (session.detailAvailable === false) return; - onOpen(session); -} - -function handleSessionKeyDown( - event: KeyboardEvent, - session: SessionRecord, - onOpen: (session: SessionRecord) => void, -) { - if (event.key !== "Enter" && event.key !== " ") return; - event.preventDefault(); - activateSession(session, onOpen); -} - -function modelLabel(session: SessionRecord) { - return session.model || session.provider || session.source; + return { series, responses }; } -function modelIcon(session: SessionRecord): DashboardIcon { - const value = `${session.provider ?? ""} ${session.model ?? ""} ${session.source}`.toLowerCase(); - if (value.includes("claude") || value.includes("anthropic")) return UiSparkles; - if (value.includes("codex") || value.includes("openai") || value.includes("gpt")) return UiRobotAi; - return UiTerminal; -} - -function effortLabel(value: string | undefined) { - return value ? value.replace(/_/g, " ") : undefined; -} - -function formatPercent(value: number | undefined) { - return value === undefined ? "--" : `${value.toFixed(1)}%`; +function staticThroughputFetcher(responses: Record) { + return async (url: string): Promise => { + const path = url.split("?")[0] ?? url; + const id = decodeURIComponent(path.split("/").filter(Boolean).pop() ?? path); + return responses[id] ?? { id, points: [] }; + }; } -function formatRelativeTime(value: string | undefined) { - if (!value) return "--"; - const time = new Date(value).getTime(); - if (Number.isNaN(time)) return value; - const delta = Date.now() - time; - if (delta < 60_000) return "now"; - if (delta < 3_600_000) return `${Math.round(delta / 60_000)}m ago`; - if (delta < 86_400_000) return `${Math.round(delta / 3_600_000)}h ago`; - return `${Math.round(delta / 86_400_000)}d ago`; +function throughputGroupLabel(group: SessionThroughputGroup) { + const effort = effortLabel(group.reasoningEffort); + return effort && effort !== "default" ? `${group.model} / ${effort}` : group.model; +} + +function hashThroughputGroup(group: SessionThroughputGroup) { + const lastPoint = group.points?.[group.points.length - 1]; + const input = [ + group.key, + group.sessions, + group.durationSeconds, + group.outputTokens, + group.totalTokens, + group.contextTokens, + lastPoint?.at, + lastPoint?.outputTokensPerSecond, + lastPoint?.contextTokensPerSecond, + ].join("|"); + let hash = 0; + for (let i = 0; i < input.length; i++) { + hash = (hash * 31 + input.charCodeAt(i)) >>> 0; + } + return hash.toString(36); } -function contextTone(percent: number) { - if (percent <= 10) return "font-medium text-destructive"; - if (percent <= 25) return "font-medium text-amber-700"; - return "font-medium text-emerald-700"; +function formatRate(value: number) { + if (!Number.isFinite(value) || value <= 0) return "0/s"; + if (value >= 1000) return `${formatCompactNumber(Math.round(value))}/s`; + if (value >= 10) return `${value.toFixed(0)}/s`; + return `${value.toFixed(1)}/s`; } -function contextBarTone(percent: number) { - if (percent <= 10) return "bg-destructive"; - if (percent <= 25) return "bg-amber-500"; - return "bg-emerald-500"; +function formatOptionalRate(value: number | undefined) { + return value !== undefined && value > 0 ? formatRate(value) : "--"; } -function copySessionRef(session: SessionRecord) { - if (!navigator.clipboard) return; - const value = session.live?.pid ? `${session.source}:${session.live.pid}` : session.key; - void navigator.clipboard.writeText(value).catch(() => undefined); +function formatOptionalPercent(value: number | undefined) { + if (value === undefined || !Number.isFinite(value) || value <= 0) return "--"; + return `${value.toFixed(value >= 10 ? 0 : 1)}%`; } diff --git a/pkg/cli/webapp/src/PromptBatchInspector.tsx b/pkg/cli/webapp/src/PromptBatchInspector.tsx new file mode 100644 index 0000000..5ea34d8 --- /dev/null +++ b/pkg/cli/webapp/src/PromptBatchInspector.tsx @@ -0,0 +1,213 @@ +import { memo, useCallback, useEffect, useMemo, useState } from "react"; +import { useQuery } from "@tanstack/react-query"; +import { Button, Select } from "@flanksource/clicky-ui/components"; +import { + SessionChatComposer, + SessionInspector, + type UnifiedSessionInput, +} from "@flanksource/clicky-ui/ai"; +import { + usePromptRunStream, + type PromptBatchHandle, + type PromptBatchRunHandle, + type PromptRunStreamState, +} from "./hooks/usePromptRunStream"; +import { useSessionChat, mergeSessionMessages } from "./hooks/useSessionChat"; +import { fetchSession } from "./sessionData"; +import { + batchChatTargetState, + batchChatTargets, + batchSessionCollection, +} from "./sessionCollection"; + +export function PromptBatchInspector({ + handle, + onEdit, +}: { + handle: PromptBatchHandle; + onEdit: () => void; +}) { + const query = useQuery({ + queryKey: ["prompt-batch", handle.batchId], + queryFn: () => fetchSession(handle.batchId), + refetchInterval: 2_000, + }); + const [streams, setStreams] = useState>( + {}, + ); + const updateStream = useCallback( + (runID: string, state: PromptRunStreamState) => + setStreams((current) => + current[runID] === state ? current : { ...current, [runID]: state }, + ), + [], + ); + const liveSessions = useMemo(() => { + const sessions = new Map(); + const savedSessions = new Map( + query.data?.sessions.map((item) => [item.captainId, item.detail]) ?? [], + ); + for (const run of handle.runs) { + const stream = streams[run.runId]; + if (!stream) continue; + const saved = savedSessions.get(run.sessionId); + sessions.set(run.sessionId, { + ...saved, + id: run.sessionId, + source: run.backend || saved?.source || "captain", + title: run.selector || run.model || saved?.title, + backend: run.backend || saved?.backend, + model: run.model || saved?.model, + reasoningEffort: run.effort || saved?.reasoningEffort, + messages: mergeSessionMessages(saved?.messages ?? [], stream.messages), + }); + } + return sessions; + }, [handle.runs, query.data?.sessions, streams]); + const collection = useMemo( + () => batchSessionCollection(handle, query.data, liveSessions), + [handle, liveSessions, query.data], + ); + const targets = useMemo(() => batchChatTargets(handle), [handle]); + const [targetID, setTargetID] = useState(targets[0]?.sessionId); + const target = + targets.find((candidate) => candidate.sessionId === targetID) ?? targets[0]; + const targetState = batchChatTargetState(query.data, target); + const chat = useSessionChat({ + initialRunID: target?.runId, + sessionID: target?.sessionId, + initialCapabilities: targetState?.capabilities ?? target?.capabilities, + initialState: targetState, + onTerminal: async () => { + await query.refetch(); + }, + }); + + useEffect(() => { + if (target && target.sessionId !== targetID) setTargetID(target.sessionId); + }, [target, targetID]); + + return ( +
+ {handle.runs.map((run) => ( + + ))} +
+
+
Multi-model run
+
+ {handle.batchId} +
+
+ +
+ {query.error && ( +
+ {errorMessage(query.error)} +
+ )} + { + const run = handle.runs.find( + (candidate) => candidate.sessionId === item.id, + ); + return run ? : undefined; + }} + {...(target + ? { + composer: ( + +