diff --git a/.gavel.yaml b/.gavel.yaml new file mode 100644 index 00000000..772a0fba --- /dev/null +++ b/.gavel.yaml @@ -0,0 +1,28 @@ +ai: {} +checks: {} +commit: + allow: + - pkg/cli/webapp/dist/.gitkeep + grouping: {} + lint: {} + message: {} + precommit: {} + summary: {} + tidy: {} +fixtures: {} +lint: + fix: {} +pr: + content: {} + fix: {} +procfile: {} +secrets: {} +ssh: {} +status: + summary: {} +test: + outlineSummary: {} +todos: + plan: {} + run: {} + verify: {} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 698dc2e1..61fbf9a7 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -23,10 +23,9 @@ jobs: with: go-version: "1.26.x" - name: Lint with Gavel - uses: flanksource/gavel@25b5e48fb166ba33a47544606aaa3e16644fa4bb # v0.0.53 + uses: flanksource/gavel@43a9189b99e71ed928f418e01cd34aa797c2c0e0 # v0.0.54 with: args: lint golangci-lint - version: v0.0.53 artifact-name: gavel-lint-results comment-header: captain-gavel-lint comment: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name == github.repository }} @@ -55,7 +54,7 @@ jobs: # pkg/cli/webapp is ignored: its vitest suite needs a sibling clicky-ui # checkout (a link: dependency), which does not exist on CI runners. - name: Test with Gavel - uses: flanksource/gavel@25b5e48fb166ba33a47544606aaa3e16644fa4bb # v0.0.53 + uses: flanksource/gavel@43a9189b99e71ed928f418e01cd34aa797c2c0e0 # v0.0.54 with: args: test --ignore tests/e2e --ignore pkg/cli/webapp version: v0.0.53 diff --git a/Procfile b/Procfile index 6e55617d..e1c54767 100644 --- a/Procfile +++ b/Procfile @@ -1 +1 @@ -dev: go run ./cmd/captain serve --dev +dev: killport 9020; go run ./cmd/captain serve --dev --port 9020 diff --git a/README.md b/README.md index bf99d835..5d358fb4 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,7 @@ Supported backends are inferred from code and dependencies, including: ### 8. Utility commands +- `cmux info ` — resolves all processes in a cmux surface, pane, or workspace and reports runtime, CPU, memory, listening TCP ports, and optional Go stacks - `cmux screenshot` — captures a screenshot of the active browser surface in cmux and copies the path to the clipboard - `port kill ` — finds and kills the process listening on a TCP port @@ -184,7 +185,7 @@ captain/ ├── pkg/claude/ # Claude history, sessions, parsing, formatting ├── pkg/cli/ # Cobra/clicky command implementations ├── pkg/cli/webapp/ # Embedded React web UI (served by captain serve) -├── pkg/cmux/ # Terminal multiplexer integration (screenshot) +├── pkg/cmux/ # Terminal multiplexer integration (processes and screenshots) ├── pkg/collections/ # Generic collection utilities ├── pkg/container/ # Sandbox discovery, generation, build/run logic ├── pkg/dod/ # Definition of Done persistence and execution @@ -491,10 +492,15 @@ Important generate/build flags: captain whoami captain whoami --backend anthropic captain whoami --models=false +captain whoami --no-cache ``` Lists every AI adapter (API providers and CLI agents: `anthropic`, `openai`, `gemini`, `claude-cli`, `claude-agent`, `codex-cli`, `gemini-cli`), their authentication method, binary availability, and a live model listing. +Provider model listings are resolved through the persisted cache in `~/.config/captain/models.json` (24h TTL, invalidated when the set of configured API keys changes), and priced from the OpenRouter snapshot in `~/Library/Caches/flanksource/openrouter-pricing.json` (24h TTL). `--no-cache` skips both, re-queries every provider's model endpoint plus OpenRouter pricing, and rewrites each cache with the fresh result. Add `-v` to see an access line per request, `-vv` for headers and query params, `-vvv` for request bodies, `-vvvv` for response bodies; credentials are redacted at every rung. Failed requests (status >= 400 or a transport error) are logged at the default verbosity. `-Plog.level.http=` raises HTTP logging alone, and `-Phttp.log.base-level=` (or `HTTP_LOG_BASE_LEVEL`) shifts the whole ladder. + +To keep the traffic instead of watching it scroll past, `-Phttp.har=` writes every request/response pair — including redirect hops and retries — to a HAR 1.2 archive you can open in browser DevTools. It applies to any command, not just `whoami`, and the file is written even when the command fails. `-Phttp.har.level=metadata` records headers, query strings and timings without bodies; `-Phttp.har.maxBodySize=` changes the 64 KB per-body cap (`0` for none). Credentials are masked the same way as in the wire log, which means the archive is safe to attach to a bug report but cannot be replayed — `-Phttp.har.sensitive=true` keeps them verbatim and writes the file `0600`. Use `-Phttp.captain.har=` to capture captain's own traffic when a shared `http.har` is already set. + ### Configure ```bash @@ -523,7 +529,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`, binds both the Go API and Vite UI to random free ports, and proxies `/api` to the API. Pass `--port` or `--ui-port` to use a specific development port. 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` keeps the Go API on the configured `--port` (`9020` by default), starts Vite on a random free port, and proxies `/api` to the API. Pass `--ui-port` to use a specific Vite port. 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 @@ -536,6 +542,17 @@ Exposes captain commands as MCP tools. Auto-exposes all commands except `sandbox ### Utility commands ```bash +# Inspect every process attributed to a copied cmux surface +captain cmux info \ + surface_ref=surface:21 \ + surface_id=65BB4725-B785-48DE-B3FD-31167ECB8300 + +# Pipe the Copy IDs block directly from the clipboard +pbpaste | captain cmux info + +# Inspect a PID and request a Go goroutine stack through gops +captain cmux info 33745 --stack + # Screenshot active browser surface in cmux captain cmux screenshot diff --git a/cmd/captain/help.go b/cmd/captain/help.go new file mode 100644 index 00000000..7f4d4fb5 --- /dev/null +++ b/cmd/captain/help.go @@ -0,0 +1,33 @@ +package main + +import ( + "fmt" + + "github.com/flanksource/clicky" + "github.com/flanksource/commons/help" + "github.com/spf13/cobra" +) + +// installRootHelp appends commons' documentation of the runtime knobs captain +// inherits — log verbosity, HTTP wire logging, HAR capture and output +// formatting — to `captain --help`. +// +// Cobra resolves a command's help function by walking up to the root, so this +// one runs for every command; the block is emitted only when the target is the +// root itself, leaving subcommand help as cobra renders it and leaving the +// sandbox/container SetHelpFunc overrides untouched. +func installRootHelp(root *cobra.Command) { + cobraHelp := root.HelpFunc() + root.SetHelpFunc(func(cmd *cobra.Command, args []string) { + cobraHelp(cmd, args) + if cmd != root { + return + } + text := help.Help() + out := text.ANSI() + if clicky.Flags.NoColor { + out = text.String() + } + fmt.Fprintf(cmd.OutOrStdout(), "\n%s\n", out) + }) +} diff --git a/cmd/captain/help_test.go b/cmd/captain/help_test.go new file mode 100644 index 00000000..ec21d08b --- /dev/null +++ b/cmd/captain/help_test.go @@ -0,0 +1,59 @@ +package main + +import ( + "bytes" + + "github.com/flanksource/clicky" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/spf13/cobra" +) + +var _ = Describe("root help", func() { + // commonsMarker is a property only the commons help block documents, so its + // presence distinguishes the appended block from cobra's own output. + const commonsMarker = "http.har.maxBodySize" + + newRoot := func() (*cobra.Command, *cobra.Command, *bytes.Buffer) { + out := &bytes.Buffer{} + root := &cobra.Command{Use: "captain", Short: "test root"} + child := &cobra.Command{Use: "child", Short: "test child", Run: func(*cobra.Command, []string) {}} + root.AddCommand(child) + root.SetOut(out) + root.SetErr(out) + installRootHelp(root) + return root, child, out + } + + It("appends the commons runtime knobs to the root help", func() { + root, _, out := newRoot() + root.SetArgs([]string{"--help"}) + + Expect(root.Execute()).To(Succeed()) + Expect(out.String()).To(ContainSubstring("Usage:"), "cobra's own help must still render") + Expect(out.String()).To(ContainSubstring("Available Commands:")) + Expect(out.String()).To(ContainSubstring(commonsMarker)) + Expect(out.String()).To(ContainSubstring("HTTP wire logging")) + }) + + It("leaves subcommand help untouched", func() { + root, _, out := newRoot() + root.SetArgs([]string{"child", "--help"}) + + Expect(root.Execute()).To(Succeed()) + Expect(out.String()).To(ContainSubstring("Usage:")) + Expect(out.String()).ToNot(ContainSubstring(commonsMarker)) + }) + + It("drops ANSI escapes when --no-color is set", func() { + defer func(previous bool) { clicky.Flags.NoColor = previous }(clicky.Flags.NoColor) + clicky.Flags.NoColor = true + + root, _, out := newRoot() + root.SetArgs([]string{"--help"}) + + Expect(root.Execute()).To(Succeed()) + Expect(out.String()).To(ContainSubstring(commonsMarker)) + Expect(out.String()).ToNot(ContainSubstring("\x1b[")) + }) +}) diff --git a/cmd/captain/main.go b/cmd/captain/main.go index da48e814..171734c7 100644 --- a/cmd/captain/main.go +++ b/cmd/captain/main.go @@ -25,25 +25,43 @@ func main() { // respect it (cobra walks up the tree), so a runtime error from any // command prints just the error. SilenceUsage: true, - PersistentPreRun: func(cmd *cobra.Command, args []string) { + PersistentPreRunE: func(cmd *cobra.Command, args []string) error { clicky.Flags.UseFlags() - cli.EnableHTTPWireLogging() + // A malformed -Phttp.har.level is a hard stop rather than a run + // that silently captures nothing. + if err := cli.EnableHTTPWireLogging(); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } // Every model/effort resolution path reads a process-wide opt-out set; // a malformed config is a hard stop rather than a silent full catalog. if err := cli.InstallDisabledSelections(); err != nil { fmt.Fprintln(os.Stderr, err) os.Exit(1) } + // Bind the database context every command reads from, so an unknown + // --context fails before the command runs rather than at first query. + name, err := cli.ResolveDatabaseContextName(cmd.Context()) + if err != nil { + return err + } + cmd.SetContext(cli.ContextWithDatabaseContext(cmd.Context(), name)) + return nil }, } configureVersion(rootCmd, currentBuildInfo()) clicky.BindAllFlags(rootCmd.PersistentFlags(), "format") - cli.BindDatabaseURLFlag(rootCmd.PersistentFlags()) + cli.BindDatabaseFlags(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. + // wire logging can be toggled, e.g. -Plog.level.http=trace or + // -Phttp.log.base-level=info. properties.BindFlags(rootCmd.PersistentFlags()) + // Document those properties where they are discoverable: appended to the + // root --help, after the flags they complement. + installRootHelp(rootCmd) + // Bind HistoryOptions directly on rootCmd so 'captain' IS 'captain history'. // All history flags (--tool, --category, --since, --limit, -f, ...) work // at the root. RunHistory auto-detects piped stdin when File is empty. @@ -67,7 +85,7 @@ func main() { 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 := clicky.AddNamedCommandWithContext("cost", rootCmd, cli.CostOptions{}, cli.RunCost) costCmd.Short = "Show token usage and estimated costs" costCmd.Long = "Display token consumption (input, output, cache read/write) and estimated costs across Claude Code sessions." @@ -79,6 +97,10 @@ func main() { planCmd.Short = "Show the exit-plan-mode plan for a session" planCmd.Long = "Determine the plan file path and content for a Claude Code or Codex session. Pass a session ID (exact or prefix) to target a specific session; otherwise the most recent session with a plan in the current directory is used. Claude plans resolve to a ~/.claude/plans/.md file; Codex plans are inline update_plan checklists. Use --path to print only the plan file path." + contextsCmd := clicky.AddNamedCommandWithContext("contexts", rootCmd, cli.ContextsOptions{}, cli.RunContexts) + contextsCmd.Short = "List the databases captain can read" + contextsCmd.Long = "List the configured database contexts. Only the default context is monitored and written to; every other context is read-only and is selected per command with --context or per browser session from the UI's project picker. Pass --check to connect to each one and report whether it is reachable." + sessionsCmd := &cobra.Command{Use: "sessions", Short: "Browse Claude and Codex sessions"} rootCmd.AddCommand(sessionsCmd) clicky.AddNamedCommandWithContext("list", sessionsCmd, cli.SessionListOptions{}, cli.RunSessionList).Short = "List discovered sessions" @@ -106,9 +128,11 @@ func main() { Short: "AI provider commands", Long: "AI provider commands.\n\n" + "Logging: increase application verbosity with -v/-vv or --log-level=debug. " + - "To log HTTP requests/responses to the provider APIs (with sensitive headers " + - "redacted), set -Plog.level.http=trace3 for headers and timing, or trace4 to " + - "also include request/response bodies.", + "HTTP calls to the provider APIs are logged on the same ladder (with credentials " + + "redacted): failed requests are logged by default, -v adds an access line per " + + "request, -vv adds headers and query params, -vvv request bodies, -vvvv response " + + "bodies. Use -Plog.level.http= to raise only HTTP logging, or " + + "-Phttp.har= to write the exchanges to a HAR archive instead.", } rootCmd.AddCommand(aiCmd) aiCmd.AddCommand(cli.NewCommandAlias(cli.CommandAliasOptions{ @@ -125,7 +149,7 @@ 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 (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." + 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. Disabled models are hidden by default; pass --disabled=true to include them. Pass --models=false to skip the network probes, --backend to inspect a single adapter, or --no-cache to bypass the persisted model and pricing caches and re-query both live." configureCmd := clicky.AddNamedCommandWithContext("configure", rootCmd, cli.ConfigureOptions{}, cli.RunConfigure) configureCmd.Use = "configure [provider]" @@ -248,6 +272,10 @@ func main() { cmuxCmd := &cobra.Command{Use: "cmux", Short: "Cmux terminal multiplexer commands"} rootCmd.AddCommand(cmuxCmd) + cmuxInfoCmd := clicky.AddNamedCommandWithContext("info", cmuxCmd, cli.CmuxInfoOptions{}, cli.RunCmuxInfo) + cmuxInfoCmd.Use = "info [copy-id-lines|pid]..." + cmuxInfoCmd.Short = "Inspect processes running in a cmux target" + cmuxInfoCmd.Long = "Resolve copied cmux workspace, pane, or surface IDs (or a PID) and report each process, runtime, CPU, memory, and listening TCP ports. Use --stack to request Go goroutine stacks through gops." screenshotCmd := clicky.AddNamedCommand("screenshot", cmuxCmd, cli.CmuxScreenshotOptions{}, cli.RunCmuxScreenshot) screenshotCmd.Short = "Take a screenshot of the active browser surface" screenshotCmd.Long = "Capture a screenshot of the currently focused browser panel in cmux and copy the file path to the clipboard." @@ -258,7 +286,11 @@ func main() { portKillCmd.Short = "Kill the process listening on a TCP port" portKillCmd.Long = "Find the process bound to the specified TCP port using lsof and kill it with SIGKILL. Reports the process name and PID before killing." - if err := rootCmd.Execute(); err != nil { + err := rootCmd.Execute() + // Flush before exiting: PersistentPostRun does not run when a command + // fails, and a failed run is the one whose HAR you want. + cli.FlushHAR() + if err != nil { os.Exit(1) } } diff --git a/go.mod b/go.mod index 01817de1..50d3bfb0 100644 --- a/go.mod +++ b/go.mod @@ -11,7 +11,7 @@ require ( github.com/firebase/genkit/go v1.11.0 github.com/flanksource/clicky v1.21.52 github.com/flanksource/clicky/aichat v1.21.48 - github.com/flanksource/commons v1.54.0 + github.com/flanksource/commons v1.55.0 github.com/flanksource/sandbox-runtime v1.0.2 github.com/fsnotify/fsnotify v1.9.0 github.com/google/dotprompt/go v0.0.0-20260502013637-5cd4a8405ca3 @@ -23,6 +23,7 @@ require ( github.com/samber/lo v1.53.0 github.com/segmentio/encoding v0.5.4 github.com/sergi/go-diff v1.4.0 + github.com/shirou/gopsutil/v3 v3.24.5 github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.11.1 github.com/t14raptor/go-fast v0.1.0 @@ -60,7 +61,7 @@ require ( github.com/nukilabs/ftoa v1.0.0 // indirect github.com/nukilabs/unicodeid v0.1.0 // indirect github.com/pgplex/pgparser v0.2.0 // indirect - github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect + github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 github.com/xi2/xz v0.0.0-20171230120015-48954b6210f8 // indirect github.com/zclconf/go-cty v1.14.4 // indirect github.com/zclconf/go-cty-yaml v1.1.0 // indirect @@ -136,7 +137,7 @@ require ( github.com/charmbracelet/bubbletea v1.3.10 // indirect github.com/charmbracelet/colorprofile v0.4.3 // indirect github.com/charmbracelet/lipgloss v1.1.0 // indirect - github.com/charmbracelet/x/ansi v0.11.6 // indirect + github.com/charmbracelet/x/ansi v0.11.6 github.com/charmbracelet/x/cellbuf v0.0.15 // indirect github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0 // indirect github.com/charmbracelet/x/term v0.2.2 // indirect @@ -292,7 +293,6 @@ require ( github.com/rogpeppe/go-internal v1.14.1 // indirect github.com/samber/oops v1.21.0 // indirect github.com/segmentio/asm v1.2.1 // indirect - github.com/shirou/gopsutil/v3 v3.24.5 // indirect github.com/shoenig/go-m1cpu v0.1.7 // indirect github.com/shopspring/decimal v1.4.0 // indirect github.com/sirupsen/logrus v1.9.4 // indirect diff --git a/go.sum b/go.sum index 9c81f698..74e16fbe 100644 --- a/go.sum +++ b/go.sum @@ -278,8 +278,8 @@ github.com/flanksource/clicky v1.21.52 h1:JtcBD05mIbE0cLuhu218Z9uO6O2mFARmgrtEIB github.com/flanksource/clicky v1.21.52/go.mod h1:eonv42hF6W1IPjQXOL6roBm/nHCYNomqc3Kps2E3RZE= github.com/flanksource/clicky/aichat v1.21.48 h1:f8Kvl96Lfp1qcqPuve1zsjaYN8ZcK1/FYqnkGA7VN30= github.com/flanksource/clicky/aichat v1.21.48/go.mod h1:PGN/lVAgxpchRctciUCpR4YIuqWoDwRxDh339A6wi3w= -github.com/flanksource/commons v1.54.0 h1:tXTJ9rzko94HdiZe9vcG857i5+02TyZ3qkl+GVElWOs= -github.com/flanksource/commons v1.54.0/go.mod h1:gupTCRqGpgD8dd2ooE7bMDJxkfcVKvkPVuHg3cWkW+Q= +github.com/flanksource/commons v1.55.0 h1:gj9zBY3V1qgAAnEiLaeGbkqCmNK0p1tJVQCDurdTZ2k= +github.com/flanksource/commons v1.55.0/go.mod h1:gupTCRqGpgD8dd2ooE7bMDJxkfcVKvkPVuHg3cWkW+Q= github.com/flanksource/commons-db v0.1.26 h1:NXAP0WvMs4ufyDfl1L2ryRxBv5qxV67GiI1nINd4YIw= github.com/flanksource/commons-db v0.1.26/go.mod h1:i378WIxy8g9xOeLBvhh1y3FO99oCumUqNmfhuDF79kM= github.com/flanksource/gomplate/v3 v3.24.84 h1:UOE0yCJsczTIKRaHUvhD6tjCYrbNvOugAizuy0FVlhE= diff --git a/migrations/02_merge_duplicate_sessions.sql b/migrations/02_merge_duplicate_sessions.sql index 577c7366..5a92a19e 100644 --- a/migrations/02_merge_duplicate_sessions.sql +++ b/migrations/02_merge_duplicate_sessions.sql @@ -74,6 +74,9 @@ BEGIN UPDATE captain_prompt_runs r SET execution_session_id = m.winner FROM captain_duplicate_session_map m WHERE r.execution_session_id = m.loser; + UPDATE captain_session_processes p SET session_id = m.winner + FROM captain_duplicate_session_map m WHERE p.session_id = m.loser; + -- captain_sessions.parent_session_id and root_session_id are self-referential -- and ON DELETE CASCADE, so a ghost that any subagent row named as its parent -- takes that whole subtree -- sessions, messages, turns, artifacts -- down with diff --git a/migrations/20_prompt_runs_and_plans.pg.hcl b/migrations/20_prompt_runs_and_plans.pg.hcl index dcaa587c..5238c5b3 100644 --- a/migrations/20_prompt_runs_and_plans.pg.hcl +++ b/migrations/20_prompt_runs_and_plans.pg.hcl @@ -10,6 +10,10 @@ table "captain_prompt_runs" { null = false type = uuid } + column "turn_id" { + null = true + type = uuid + } column "root_session_id" { null = false type = uuid @@ -87,6 +91,22 @@ table "captain_prompt_runs" { null = true type = jsonb } + column "approval_state" { + null = true + type = jsonb + } + column "provider_checkpoint_codec" { + null = true + type = text + } + column "provider_checkpoint_version" { + null = true + type = integer + } + column "provider_checkpoint" { + null = true + type = bytea + } column "error" { null = true type = text @@ -130,6 +150,12 @@ table "captain_prompt_runs" { on_update = NO_ACTION on_delete = CASCADE } + foreign_key "captain_prompt_runs_turn_id_fkey" { + columns = [column.turn_id] + ref_columns = [table.captain_turns.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } foreign_key "captain_prompt_runs_root_session_id_fkey" { columns = [column.root_session_id] ref_columns = [table.captain_sessions.column.id] @@ -192,6 +218,9 @@ table "captain_prompt_runs" { index "captain_prompt_runs_state_idx" { columns = [column.state, column.phase, column.updated_at] } + index "captain_prompt_runs_turn_id_idx" { + columns = [column.turn_id] + } check "captain_prompt_runs_iteration_nonnegative" { expr = "current_iteration >= 0" @@ -208,6 +237,9 @@ table "captain_prompt_runs" { check "captain_prompt_runs_time_order" { expr = "(started_at IS NULL OR started_at >= queued_at) AND (finished_at IS NULL OR (started_at IS NOT NULL AND finished_at >= started_at))" } + check "captain_prompt_runs_provider_checkpoint" { + expr = "(provider_checkpoint IS NULL AND provider_checkpoint_codec IS NULL AND provider_checkpoint_version IS NULL) OR (provider_checkpoint IS NOT NULL AND length(btrim(provider_checkpoint_codec)) > 0 AND provider_checkpoint_version > 0)" + } } table "captain_prompt_run_iterations" { diff --git a/migrations/30_execution.pg.hcl b/migrations/30_execution.pg.hcl index 18e672d2..1de42264 100644 --- a/migrations/30_execution.pg.hcl +++ b/migrations/30_execution.pg.hcl @@ -90,7 +90,6 @@ table "captain_turns" { expr = "ended_at IS NULL OR (started_at IS NOT NULL AND ended_at >= started_at)" } } - table "captain_model_calls" { schema = schema.public @@ -200,6 +199,11 @@ table "captain_model_calls" { type = numeric(20, 8) default = 0 } + column "provider_cost_usd" { + null = false + type = numeric(20, 8) + default = 0 + } column "currency" { null = false type = text @@ -302,7 +306,7 @@ table "captain_model_calls" { expr = "input_tokens >= 0 AND output_tokens >= 0 AND reasoning_tokens >= 0 AND cache_read_tokens >= 0 AND cache_write_tokens >= 0 AND context_tokens >= 0 AND context_window_tokens >= 0" } check "captain_model_calls_costs_nonnegative" { - expr = "input_cost >= 0 AND output_cost >= 0 AND reasoning_cost >= 0 AND cache_read_cost >= 0 AND cache_write_cost >= 0" + expr = "input_cost >= 0 AND output_cost >= 0 AND reasoning_cost >= 0 AND cache_read_cost >= 0 AND cache_write_cost >= 0 AND provider_cost_usd >= 0" } check "captain_model_calls_currency" { expr = "length(currency) = 3" @@ -314,432 +318,3 @@ table "captain_model_calls" { expr = "iteration_id IS NULL OR prompt_run_id IS NOT NULL" } } - -table "captain_messages" { - 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 "provider_message_id" { - null = true - type = text - } - column "sequence" { - null = false - type = bigint - } - column "role" { - null = false - type = text - } - column "parts" { - null = false - type = jsonb - default = sql("'[]'::jsonb") - } - column "raw" { - null = true - type = jsonb - } - column "source_line" { - null = true - type = bigint - } - 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_messages_session_id_fkey" { - columns = [column.session_id] - ref_columns = [table.captain_sessions.column.id] - on_update = NO_ACTION - on_delete = CASCADE - } - foreign_key "captain_messages_turn_id_fkey" { - columns = [column.turn_id] - ref_columns = [table.captain_turns.column.id] - on_update = NO_ACTION - on_delete = CASCADE - } - foreign_key "captain_messages_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_messages_session_sequence_key" { - unique = true - columns = [column.session_id, column.sequence] - } - index "captain_messages_provider_message_id_key" { - unique = true - columns = [column.session_id, column.provider_message_id] - where = "provider_message_id IS NOT NULL" - } - index "captain_messages_turn_id_idx" { - columns = [column.turn_id] - } - # Partial on purpose. captain_messages_model_call_id_fkey is ON DELETE SET - # NULL, so deleting a model call must find its referencing messages -- without - # an index that is a sequential scan of the largest table in the schema. But - # the ingest path never sets model_call_id, so a full index on the column was - # 8.4 MB of nothing but NULLs, maintained on every message insert. Excluding - # the NULLs costs the FK check nothing (model_call_id = 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/31_execution_events.pg.hcl b/migrations/31_execution_events.pg.hcl new file mode 100644 index 00000000..f3eb81eb --- /dev/null +++ b/migrations/31_execution_events.pg.hcl @@ -0,0 +1,291 @@ +table "captain_messages" { + 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 "provider_message_id" { + null = true + type = text + } + column "sequence" { + null = false + type = bigint + } + column "role" { + null = false + type = text + } + column "parts" { + null = false + type = jsonb + default = sql("'[]'::jsonb") + } + column "raw" { + null = true + type = jsonb + } + column "source_line" { + null = true + type = bigint + } + 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_messages_session_id_fkey" { + columns = [column.session_id] + ref_columns = [table.captain_sessions.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + foreign_key "captain_messages_turn_id_fkey" { + columns = [column.turn_id] + ref_columns = [table.captain_turns.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + foreign_key "captain_messages_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_messages_session_sequence_key" { + unique = true + columns = [column.session_id, column.sequence] + } + index "captain_messages_provider_message_id_key" { + unique = true + columns = [column.session_id, column.provider_message_id] + where = "provider_message_id IS NOT NULL" + } + index "captain_messages_turn_id_idx" { + columns = [column.turn_id] + } + # Partial on purpose. captain_messages_model_call_id_fkey is ON DELETE SET + # NULL, so deleting a model call must find its referencing messages -- without + # an index that is a sequential scan of the largest table in the schema. But + # the ingest path never sets model_call_id, so a full index on the column was + # 8.4 MB of nothing but NULLs, maintained on every message insert. Excluding + # the NULLs costs the FK check nothing (model_call_id = 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" + } + + # `sequence >= 0` used to be checked here. The sequence space now has two + # halves and the constraint no longer describes it. Non-negative sequences are + # the provider's: the ingester keys every message on its line number in the + # transcript JSONL and converges on (session_id, sequence), so that half is + # reserved for lines the file already has or is about to grow into. Negative + # sequences are the harness's own insertions -- lifecycle notices recorded by + # hooks acting between turns, which have no line in any transcript. Writing + # those at MAX(sequence)+1 would put them in the provider's half, where the + # next ingest pass overwrites them with the real line landing on that number. + # + # Ordering across the two halves is carried by occurred_at, not by sequence. + + 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" + } +} + diff --git a/migrations/32_execution_approvals.pg.hcl b/migrations/32_execution_approvals.pg.hcl new file mode 100644 index 00000000..46ad4ac9 --- /dev/null +++ b/migrations/32_execution_approvals.pg.hcl @@ -0,0 +1,248 @@ +table "captain_session_mcp_credentials" { + schema = schema.public + + column "id" { + null = false + type = uuid + default = sql("gen_random_uuid()") + } + column "session_id" { + null = false + type = uuid + } + column "prompt_run_id" { + null = false + type = uuid + } + column "backend" { + null = false + type = text + } + column "secret_hash" { + null = false + type = bytea + } + column "policy" { + null = false + type = jsonb + default = sql("'{}'::jsonb") + } + column "expires_at" { + null = true + type = timestamptz + } + column "revoked_at" { + null = true + type = timestamptz + } + column "revocation_reason" { + null = true + type = text + } + column "created_at" { + null = false + type = timestamptz + default = sql("now()") + } + + primary_key { + columns = [column.id] + } + + foreign_key "captain_session_mcp_credentials_session_id_fkey" { + columns = [column.session_id] + ref_columns = [table.captain_sessions.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + foreign_key "captain_session_mcp_credentials_prompt_run_id_fkey" { + columns = [column.prompt_run_id] + ref_columns = [table.captain_prompt_runs.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + + index "captain_session_mcp_credentials_secret_hash_key" { + unique = true + columns = [column.secret_hash] + } + index "captain_session_mcp_credentials_active_session_idx" { + columns = [column.session_id, column.created_at] + where = "revoked_at IS NULL" + } + + check "captain_session_mcp_credentials_hash_length" { + expr = "octet_length(secret_hash) = 32" + } + check "captain_session_mcp_credentials_expiry" { + expr = "expires_at IS NULL OR expires_at > created_at" + } + check "captain_session_mcp_credentials_revocation" { + expr = "(revoked_at IS NULL AND revocation_reason IS NULL) OR revoked_at 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 "credential_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 + } + foreign_key "captain_turn_requests_credential_id_fkey" { + columns = [column.credential_id] + ref_columns = [table.captain_session_mcp_credentials.column.id] + on_update = NO_ACTION + on_delete = CASCADE + } + + 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] + } + index "captain_turn_requests_credential_call_key" { + unique = true + columns = [column.credential_id, column.tool_call_id] + where = "credential_id IS NOT NULL AND tool_call_id IS NOT NULL" + } + + 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" + } + check "captain_turn_requests_tool_approval_identity" { + expr = "kind <> 'tool_approval' OR (prompt_run_id IS NOT NULL AND turn_id IS NOT NULL AND model_call_id IS NOT NULL AND tool_call_id IS NOT NULL)" + } +} diff --git a/migrations/60_view_session_overview.sql b/migrations/60_view_session_overview.sql index e2db594f..223ae7dd 100644 --- a/migrations/60_view_session_overview.sql +++ b/migrations/60_view_session_overview.sql @@ -64,6 +64,9 @@ SELECT 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, + -- Always 0: nothing writes captain_artifacts. Kept because CREATE OR REPLACE + -- VIEW cannot drop a column, and commons-db only drops a view when a *table* + -- diff would break it, so removing these here fails every existing database. 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, @@ -100,7 +103,19 @@ SELECT process.cpu_percent, process.memory_percent, process.memory_rss_bytes, - process.sampled_at AS process_sampled_at + process.sampled_at AS process_sampled_at, + latest_run.execution_mode, + -- The portion of cost_usd the providers themselves reported. cost_usd falls + -- back to list-priced buckets per call, so it alone cannot tell a billed + -- figure from a reconstruction; this can, and callers must not present an + -- estimate as billed cost. The bucket sums come with it so a caller holding + -- an api.Cost can fall back the same way cost_usd does. + COALESCE(call_stats.provider_cost_usd, 0::numeric) AS provider_cost_usd, + COALESCE(call_stats.input_cost, 0::numeric) AS input_cost, + COALESCE(call_stats.output_cost, 0::numeric) AS output_cost, + COALESCE(call_stats.reasoning_cost, 0::numeric) AS reasoning_cost, + COALESCE(call_stats.cache_read_cost, 0::numeric) AS cache_read_cost, + COALESCE(call_stats.cache_write_cost, 0::numeric) AS cache_write_cost FROM public.captain_sessions s LEFT JOIN LATERAL ( SELECT p.* @@ -152,12 +167,22 @@ LEFT JOIN LATERAL ( 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 + CASE + WHEN c.provider_cost_usd > 0 THEN c.provider_cost_usd + ELSE c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + END + ) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cost_usd, + COALESCE(sum(c.provider_cost_usd) + FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS provider_cost_usd, + COALESCE(sum(c.input_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS input_cost, + COALESCE(sum(c.output_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS output_cost, + COALESCE(sum(c.reasoning_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS reasoning_cost, + COALESCE(sum(c.cache_read_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cache_read_cost, + COALESCE(sum(c.cache_write_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cache_write_cost FROM public.captain_turns t LEFT JOIN public.captain_model_calls c ON c.turn_id = t.id WHERE t.session_id = s.id @@ -205,7 +230,14 @@ LEFT JOIN LATERAL ( FROM public.captain_artifacts a WHERE a.session_id = s.id AND a.kind LIKE 'file.%' -) file_stats ON true; +) file_stats ON true +LEFT JOIN LATERAL ( + SELECT NULLIF(r.runtime ->> 'mode', '') AS execution_mode + FROM public.captain_prompt_runs r + WHERE r.session_id = s.id + ORDER BY COALESCE(r.finished_at, r.started_at, r.created_at) DESC, r.id DESC + LIMIT 1 +) latest_run 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/62_view_session_turns.sql b/migrations/62_view_session_turns.sql index c67a47a9..d52f3447 100644 --- a/migrations/62_view_session_turns.sql +++ b/migrations/62_view_session_turns.sql @@ -52,7 +52,19 @@ SELECT 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 + COALESCE(event_stats.event_ids, ARRAY[]::uuid[]) AS event_ids, + -- Appended after initial release: CREATE OR REPLACE VIEW only allows adding + -- columns at the end, so later additions must stay below this line. + -- + -- The portion of cost_usd the providers reported, plus the list-price buckets + -- it otherwise falls back to. cost_usd alone cannot tell a billed figure from + -- a reconstruction, and callers must not render an estimate as billed cost. + COALESCE(call_stats.provider_cost_usd, 0::numeric) AS provider_cost_usd, + COALESCE(call_stats.input_cost, 0::numeric) AS input_cost, + COALESCE(call_stats.output_cost, 0::numeric) AS output_cost, + COALESCE(call_stats.reasoning_cost, 0::numeric) AS reasoning_cost, + COALESCE(call_stats.cache_read_cost, 0::numeric) AS cache_read_cost, + COALESCE(call_stats.cache_write_cost, 0::numeric) AS cache_write_cost FROM public.captain_turns t LEFT JOIN LATERAL ( SELECT @@ -63,12 +75,22 @@ LEFT JOIN LATERAL ( 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 + CASE + WHEN c.provider_cost_usd > 0 THEN c.provider_cost_usd + ELSE c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + END + ) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cost_usd, + COALESCE(sum(c.provider_cost_usd) + FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS provider_cost_usd, + COALESCE(sum(c.input_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS input_cost, + COALESCE(sum(c.output_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS output_cost, + COALESCE(sum(c.reasoning_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS reasoning_cost, + COALESCE(sum(c.cache_read_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cache_read_cost, + COALESCE(sum(c.cache_write_cost) FILTER (WHERE upper(c.currency) = 'USD'), 0::numeric) AS cache_write_cost FROM public.captain_model_calls c WHERE c.turn_id = t.id ) call_stats ON true diff --git a/migrations/63_view_session_agents.sql b/migrations/63_view_session_agents.sql index a8981321..92558163 100644 --- a/migrations/63_view_session_agents.sql +++ b/migrations/63_view_session_agents.sql @@ -45,11 +45,14 @@ LEFT JOIN LATERAL ( 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 + CASE + WHEN c.provider_cost_usd > 0 THEN c.provider_cost_usd + ELSE c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + END ) 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 diff --git a/migrations/64_view_session_files.sql b/migrations/64_view_session_files.sql deleted file mode 100644 index 55d88da1..00000000 --- a/migrations/64_view_session_files.sql +++ /dev/null @@ -1,25 +0,0 @@ --- 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/67_view_session_costs.sql b/migrations/67_view_session_costs.sql index 40c8e199..b752ce8b 100644 --- a/migrations/67_view_session_costs.sql +++ b/migrations/67_view_session_costs.sql @@ -34,15 +34,23 @@ SELECT 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, + -- Mirrors api.Cost.Total(): the provider's reported figure is authoritative + -- per call, with the list-price bucket sum as the fallback. Decided per row, + -- not per group, so a mix of priced and provider-reported calls still totals. sum( - c.input_cost - + c.output_cost - + c.reasoning_cost - + c.cache_read_cost - + c.cache_write_cost + CASE + WHEN c.provider_cost_usd > 0 THEN c.provider_cost_usd + ELSE c.input_cost + + c.output_cost + + c.reasoning_cost + + c.cache_read_cost + + c.cache_write_cost + END ) AS total_cost, min(c.started_at) AS first_call_at, - max(c.ended_at) AS last_call_at + max(c.ended_at) AS last_call_at, + -- Appended last: CREATE OR REPLACE VIEW can only add columns at the end. + sum(c.provider_cost_usd) AS provider_cost_usd FROM public.captain_turns t JOIN public.captain_model_calls c ON c.turn_id = t.id GROUP BY diff --git a/migrations/74_turn_request_approval_identity.sql b/migrations/74_turn_request_approval_identity.sql new file mode 100644 index 00000000..e453e82f --- /dev/null +++ b/migrations/74_turn_request_approval_identity.sql @@ -0,0 +1,78 @@ +-- phase: post + +WITH candidates AS ( + SELECT + request.id AS request_id, + model_call.id AS model_call_id, + model_call.turn_id, + count(*) OVER (PARTITION BY request.id) AS candidate_count + FROM public.captain_turn_requests request + JOIN public.captain_model_calls model_call + ON model_call.prompt_run_id = request.prompt_run_id + AND (request.model_call_id IS NULL OR request.model_call_id = model_call.id) + AND (request.turn_id IS NULL OR request.turn_id = model_call.turn_id) + JOIN public.captain_turns turn + ON turn.id = model_call.turn_id + AND turn.session_id = request.session_id + WHERE request.kind = 'tool_approval' + AND (request.turn_id IS NULL OR request.model_call_id IS NULL) +), unique_candidates AS ( + SELECT request_id, model_call_id, turn_id + FROM candidates + WHERE candidate_count = 1 +) +UPDATE public.captain_turn_requests request +SET + turn_id = COALESCE(request.turn_id, candidate.turn_id), + model_call_id = COALESCE(request.model_call_id, candidate.model_call_id) +FROM unique_candidates candidate +WHERE request.id = candidate.request_id; + +DO $$ +DECLARE + invalid_ids text; +BEGIN + SELECT string_agg(request.id::text, ', ' ORDER BY request.id) + INTO invalid_ids + FROM public.captain_turn_requests request + WHERE request.kind = 'tool_approval' + AND ( + request.prompt_run_id IS NULL + OR request.turn_id IS NULL + OR request.model_call_id IS NULL + OR request.tool_call_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM public.captain_model_calls model_call + JOIN public.captain_turns turn ON turn.id = model_call.turn_id + WHERE model_call.id = request.model_call_id + AND model_call.prompt_run_id = request.prompt_run_id + AND model_call.turn_id = request.turn_id + AND turn.session_id = request.session_id + ) + ); + + IF invalid_ids IS NOT NULL THEN + RAISE EXCEPTION 'ambiguous legacy tool approval identity for request(s): %', invalid_ids + USING ERRCODE = 'check_violation'; + END IF; +END +$$; + +ALTER TABLE public.captain_turn_requests + DROP CONSTRAINT IF EXISTS captain_turn_requests_tool_approval_identity; + +ALTER TABLE public.captain_turn_requests + ADD CONSTRAINT captain_turn_requests_tool_approval_identity + CHECK ( + kind <> 'tool_approval' + OR ( + prompt_run_id IS NOT NULL + AND turn_id IS NOT NULL + AND model_call_id IS NOT NULL + AND tool_call_id IS NOT NULL + ) + ) NOT VALID; + +ALTER TABLE public.captain_turn_requests + VALIDATE CONSTRAINT captain_turn_requests_tool_approval_identity; diff --git a/migrations/75_model_call_provider_cost.sql b/migrations/75_model_call_provider_cost.sql new file mode 100644 index 00000000..d6367f89 --- /dev/null +++ b/migrations/75_model_call_provider_cost.sql @@ -0,0 +1,27 @@ +-- phase: post + +-- captain_model_calls.provider_cost_usd carries the model provider's own billed +-- total for a call, alongside the five list-price bucket columns. The rollup +-- views prefer it per row, mirroring api.Cost.Total(). +-- +-- The non-negative check lives here rather than only in 30_execution.pg.hcl +-- because Atlas does not diff check-constraint expressions: editing the HCL +-- alone applies to freshly created databases but silently leaves existing ones +-- on the old five-column predicate. + +ALTER TABLE public.captain_model_calls + DROP CONSTRAINT IF EXISTS captain_model_calls_costs_nonnegative; + +ALTER TABLE public.captain_model_calls + ADD CONSTRAINT captain_model_calls_costs_nonnegative + CHECK ( + input_cost >= 0 + AND output_cost >= 0 + AND reasoning_cost >= 0 + AND cache_read_cost >= 0 + AND cache_write_cost >= 0 + AND provider_cost_usd >= 0 + ) NOT VALID; + +ALTER TABLE public.captain_model_calls + VALIDATE CONSTRAINT captain_model_calls_costs_nonnegative; diff --git a/migrations/76_model_call_cost_backfill.sql b/migrations/76_model_call_cost_backfill.sql new file mode 100644 index 00000000..c48d8dbc --- /dev/null +++ b/migrations/76_model_call_cost_backfill.sql @@ -0,0 +1,32 @@ +-- phase: post + +-- Repair rows written before finishChatModelCall split its cost buckets. +-- +-- The old writer put the whole turn's provider-reported cost into output_cost +-- and left the other four buckets at zero. That made every rollup read as a +-- list-price estimate (provider_cost_usd = 0) while actually reporting billed +-- money, and it overstated output_cost by the entire input and cache spend. +-- +-- Old rows are identifiable without ambiguity: a priced call always prices its +-- input, so output_cost > 0 with every other bucket at zero cannot be produced +-- by the current writer. Only the total is recoverable — the per-bucket split +-- is not, and is deliberately left at zero rather than guessed. The rollup +-- views prefer provider_cost_usd per row, so the totals stay exact. + +UPDATE public.captain_model_calls +SET provider_cost_usd = output_cost, + output_cost = 0 +WHERE provider_cost_usd = 0 + AND output_cost > 0 + AND input_cost = 0 + AND reasoning_cost = 0 + AND cache_read_cost = 0 + AND cache_write_cost = 0; + +-- The same writer set context_tokens to the input count alone. Context is what +-- the model actually read, so cache hits and writes belong in it; without them +-- a cache-heavy turn reports a context far below the window it really used. +UPDATE public.captain_model_calls +SET context_tokens = input_tokens + cache_read_tokens + cache_write_tokens +WHERE context_tokens = input_tokens + AND cache_read_tokens + cache_write_tokens > 0; diff --git a/migrations/77_drop_session_files_view.sql b/migrations/77_drop_session_files_view.sql new file mode 100644 index 00000000..b2645c9c --- /dev/null +++ b/migrations/77_drop_session_files_view.sql @@ -0,0 +1,17 @@ +-- phase: post + +-- Drop captain_session_files, which was never wired to anything. +-- +-- The view was added as "normalized file.* artifact rows for the SessionInspector +-- files tab", but no non-test code has ever inserted a captain_artifacts row, so +-- it selected zero rows on every database while the files tab it was built for +-- rendered empty. The changed-file set a session actually reports comes from the +-- monitor's projection in captain_sessions.metadata->'files'. +-- +-- Nothing depends on the view, so a plain DROP suffices. The matching +-- file_read_count/file_written_count columns on captain_session_overview are +-- left in place: CREATE OR REPLACE VIEW cannot drop a column, and commons-db +-- only drops a view when a table diff would break it, so removing them would +-- fail the migration on every existing database. + +DROP VIEW IF EXISTS public.captain_session_files; diff --git a/migrations/approval_identity_upgrade_integration_test.go b/migrations/approval_identity_upgrade_integration_test.go new file mode 100644 index 00000000..2514bcdd --- /dev/null +++ b/migrations/approval_identity_upgrade_integration_test.go @@ -0,0 +1,179 @@ +package migrations + +import ( + "context" + "database/sql" + "fmt" + "strings" + "time" + + "github.com/flanksource/commons-db/dbtest" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +const approvalIdentityMigration = "74_turn_request_approval_identity.sql" + +var _ = Describe("Tool approval identity migration", func() { + It("backfills an unambiguous legacy approval and replaces the credential constraint", func(ctx SpecContext) { + handle := dbtest.ForGinkgo(dbtest.Options{Name: "captain_approval_identity_upgrade"}) + dsn, db := handle.DSN(), handle.SQL() + Expect(Apply(ctx, dsn)).To(Succeed()) + + ids := seedLegacyToolApproval(ctx, db, 1) + Expect(resetApprovalIdentityMigration(ctx, db)).To(Succeed()) + Expect(Apply(ctx, dsn)).To(Succeed()) + + var turnID, modelCallID uuid.UUID + Expect(db.QueryRowContext(ctx, ` + SELECT turn_id, model_call_id + FROM captain_turn_requests + WHERE id = $1 + `, ids.request).Scan(&turnID, &modelCallID)).To(Succeed()) + Expect(turnID).To(Equal(ids.turn)) + Expect(modelCallID).To(Equal(ids.modelCalls[0])) + + var definition string + Expect(db.QueryRowContext(ctx, ` + SELECT pg_get_constraintdef(oid) + FROM pg_constraint + WHERE conname = 'captain_turn_requests_tool_approval_identity' + `).Scan(&definition)).To(Succeed()) + Expect(definition).To(And( + ContainSubstring("turn_id IS NOT NULL"), + ContainSubstring("model_call_id IS NOT NULL"), + Not(ContainSubstring("credential_id IS NOT NULL")), + )) + + _, err := db.ExecContext(ctx, ` + INSERT INTO captain_turn_requests ( + id, session_id, turn_id, prompt_run_id, model_call_id, tool_call_id, + kind, request, idempotency_key, requested_by, expires_at + ) VALUES ($1, $2, $3, $4, $5, 'provider-call', 'tool_approval', + '{"tool":"accounts_edit","input":{}}', $6, 'provider', $7) + `, uuid.New(), ids.session, ids.turn, ids.promptRun, ids.modelCalls[0], + "provider:"+ids.promptRun.String()+":provider-call", time.Now().Add(time.Hour)) + Expect(err).NotTo(HaveOccurred()) + Expect(Apply(ctx, dsn)).To(Succeed()) + }) + + It("fails when a legacy approval cannot be correlated to one model call", func(ctx SpecContext) { + handle := dbtest.ForGinkgo(dbtest.Options{Name: "captain_approval_identity_ambiguous"}) + dsn, db := handle.DSN(), handle.SQL() + Expect(Apply(ctx, dsn)).To(Succeed()) + + ids := seedLegacyToolApproval(ctx, db, 2) + Expect(resetApprovalIdentityMigration(ctx, db)).To(Succeed()) + err := Apply(ctx, dsn) + Expect(err).To(MatchError(And( + ContainSubstring("ambiguous legacy tool approval identity"), + ContainSubstring(ids.request.String()), + ))) + }) + + It("rejects a legacy constraint after its upgrade script is already recorded", func(ctx SpecContext) { + handle := dbtest.ForGinkgo(dbtest.Options{Name: "captain_approval_identity_drift"}) + dsn, db := handle.DSN(), handle.SQL() + Expect(Apply(ctx, dsn)).To(Succeed()) + Expect(installLegacyApprovalConstraint(ctx, db)).To(Succeed()) + + err := Apply(ctx, dsn) + Expect(err).To(MatchError(And( + ContainSubstring("verify Captain database"), + ContainSubstring("credential_id IS NOT NULL"), + ))) + }) +}) + +type legacyApprovalIDs struct { + session uuid.UUID + turn uuid.UUID + promptRun uuid.UUID + request uuid.UUID + modelCalls []uuid.UUID +} + +func seedLegacyToolApproval(ctx context.Context, db *sql.DB, modelCallCount int) legacyApprovalIDs { + ids := legacyApprovalIDs{ + session: uuid.New(), turn: uuid.New(), promptRun: uuid.New(), request: uuid.New(), + } + Expect(installLegacyApprovalConstraint(ctx, db)).To(Succeed()) + + _, err := db.ExecContext(ctx, `INSERT INTO captain_sessions (id, source) VALUES ($1, 'aichat')`, ids.session) + Expect(err).NotTo(HaveOccurred()) + _, err = db.ExecContext(ctx, ` + INSERT INTO captain_turns (id, session_id, provider_turn_id, turn_index, status, started_at) + VALUES ($1, $2, 'legacy-turn', 0, 'open', now()) + `, ids.turn, ids.session) + Expect(err).NotTo(HaveOccurred()) + _, err = db.ExecContext(ctx, ` + INSERT INTO captain_prompt_runs (id, session_id, turn_id, root_session_id, admission_key) + VALUES ($1, $2, $3, $2, 'legacy-approval-run') + `, ids.promptRun, ids.session, ids.turn) + Expect(err).NotTo(HaveOccurred()) + + for i := range modelCallCount { + modelCallID := uuid.New() + ids.modelCalls = append(ids.modelCalls, modelCallID) + _, err = db.ExecContext(ctx, ` + INSERT INTO captain_model_calls ( + id, turn_id, prompt_run_id, call_index, model, backend, status, started_at + ) VALUES ($1, $2, $3, $4, 'gemini', 'google', 'running', now()) + `, modelCallID, ids.turn, ids.promptRun, i) + Expect(err).NotTo(HaveOccurred()) + } + + credentialID := uuid.New() + _, err = db.ExecContext(ctx, ` + INSERT INTO captain_session_mcp_credentials ( + id, session_id, prompt_run_id, backend, secret_hash, policy, expires_at + ) VALUES ($1, $2, $3, 'google', $4, '{"accounts_edit":"ask"}', $5) + `, credentialID, ids.session, ids.promptRun, []byte(strings.Repeat("a", 32)), time.Now().Add(time.Hour)) + Expect(err).NotTo(HaveOccurred()) + _, err = db.ExecContext(ctx, ` + INSERT INTO captain_turn_requests ( + id, session_id, prompt_run_id, credential_id, tool_call_id, kind, + request, idempotency_key, requested_by, expires_at + ) VALUES ($1, $2, $3, $4, 'legacy-call', 'tool_approval', + '{"tool":"accounts_edit","input":{}}', $5, 'caller_tool', $6) + `, ids.request, ids.session, ids.promptRun, credentialID, + "mcp:"+credentialID.String()+":legacy-call", time.Now().Add(time.Hour)) + Expect(err).NotTo(HaveOccurred()) + return ids +} + +func installLegacyApprovalConstraint(ctx context.Context, db *sql.DB) error { + if _, err := db.ExecContext(ctx, ` + ALTER TABLE captain_turn_requests + DROP CONSTRAINT captain_turn_requests_tool_approval_identity + `); err != nil { + return err + } + _, err := db.ExecContext(ctx, ` + ALTER TABLE captain_turn_requests + ADD CONSTRAINT captain_turn_requests_tool_approval_identity + CHECK (kind <> 'tool_approval' OR ( + credential_id IS NOT NULL AND prompt_run_id IS NOT NULL AND tool_call_id IS NOT NULL + )) + `) + return err +} + +func resetApprovalIdentityMigration(ctx context.Context, db *sql.DB) error { + result, err := db.ExecContext(ctx, ` + DELETE FROM schema_migration_scripts + WHERE scope = $1 AND path = $2 + `, Scope, approvalIdentityMigration) + if err != nil { + return err + } + affected, err := result.RowsAffected() + if err != nil { + return fmt.Errorf("read reset migration result: %w", err) + } + if affected != 1 { + return fmt.Errorf("reset migration ledger: deleted %d rows, want 1", affected) + } + return nil +} diff --git a/migrations/merge_duplicate_sessions_integration_test.go b/migrations/merge_duplicate_sessions_integration_test.go index aa87f31c..3ccbed5c 100644 --- a/migrations/merge_duplicate_sessions_integration_test.go +++ b/migrations/merge_duplicate_sessions_integration_test.go @@ -9,10 +9,9 @@ import ( // 02_merge_duplicate_sessions.sql collapses the session rows the old wide // identity key allowed, and it does that with a DELETE. Everything reachable from // a deleted row by an ON DELETE CASCADE is therefore at risk, which is why the -// migration re-points references first. The reference that matters most is -// captain_sessions.parent_session_id: it is self-referential and CASCADEs, so a -// ghost that a subagent row named as its parent takes that subagent -- and its -// transcript -- with it. This pins that the collapse moves those links instead. +// migration re-points references first. The references that matter most are the +// self-referential hierarchy and the process row: deleting a ghost otherwise +// takes both the subagent transcript and the live process identity with it. var _ = Describe("Captain duplicate session collapse", func() { It("re-points a subagent at the surviving row instead of cascading it away", func(ctx SpecContext) { handle := dbtest.ForGinkgo(dbtest.Options{Name: "captain_merge_duplicates"}) @@ -36,6 +35,7 @@ var _ = Describe("Captain duplicate session collapse", func() { winnerID = "11111111-1111-1111-1111-111111111111" ghostID = "22222222-2222-2222-2222-222222222222" subagentID = "33333333-3333-3333-3333-333333333333" + processID = "44444444-4444-4444-4444-444444444444" ) // One rollout, two rows: the monitor's (ingested, provider '') and the @@ -58,6 +58,12 @@ var _ = Describe("Captain duplicate session collapse", func() { VALUES ($1, 'codex', 'host-a', 'rollout-1-sub', $2, $2)`, subagentID, ghostID) Expect(err).NotTo(HaveOccurred()) + _, err = db.ExecContext(ctx, ` + INSERT INTO captain_session_processes + (id, session_id, host_id, boot_id, pid, process_started_at, status) + VALUES ($1, $2, 'host-a', 'boot-a', 42, now(), 'running')`, processID, ghostID) + Expect(err).NotTo(HaveOccurred()) + Expect(Apply(ctx, dsn)).To(Succeed()) var surviving int @@ -75,6 +81,12 @@ var _ = Describe("Captain duplicate session collapse", func() { Expect(parent).To(Equal(winnerID)) Expect(root).To(Equal(winnerID)) + var processSessionID string + Expect(db.QueryRowContext(ctx, + `SELECT session_id FROM captain_session_processes WHERE id = $1`, processID, + ).Scan(&processSessionID)).To(Succeed(), "the process row was cascaded away with the ghost") + Expect(processSessionID).To(Equal(winnerID)) + // The winner keeps its transcript and absorbs the label that existed only on // the ghost. var messages int diff --git a/migrations/migrations.go b/migrations/migrations.go index 98d55590..cf01b0f9 100644 --- a/migrations/migrations.go +++ b/migrations/migrations.go @@ -44,6 +44,7 @@ type migrationLockHandle interface { type applyDependencies struct { acquireLock func(context.Context, string) (migrationLockHandle, error) migrate func(context.Context, string) error + verify func(context.Context, string) error } var defaultApplyDependencies = applyDependencies{ @@ -54,6 +55,7 @@ var defaultApplyDependencies = applyDependencies{ commonsmigrate.WithExclude("todo_*"), ) }, + verify: verifyToolApprovalIdentity, } // Apply reconciles the checked-in HCL schema, then applies the colocated SQL @@ -83,6 +85,53 @@ func apply(ctx context.Context, connection string, deps applyDependencies) (resu if err := deps.migrate(ctx, connection); err != nil { return fmt.Errorf("migrate Captain database: %w", err) } + if err := deps.verify(ctx, connection); err != nil { + return fmt.Errorf("verify Captain database: %w", err) + } + return nil +} + +func verifyToolApprovalIdentity(ctx context.Context, connection string) (resultErr error) { + db, err := commonsdb.NewDB(connection) + if err != nil { + return fmt.Errorf("open schema verification database: %w", err) + } + defer func() { + if err := db.Close(); err != nil { + resultErr = errors.Join(resultErr, fmt.Errorf("close schema verification database: %w", err)) + } + }() + + var validated bool + var definition string + err = db.QueryRowContext(ctx, ` + SELECT c.convalidated, pg_get_constraintdef(c.oid) + FROM pg_constraint c + JOIN pg_class relation ON relation.oid = c.conrelid + JOIN pg_namespace namespace ON namespace.oid = relation.relnamespace + WHERE namespace.nspname = 'public' + AND relation.relname = 'captain_turn_requests' + AND c.conname = 'captain_turn_requests_tool_approval_identity' + AND c.contype = 'c' + `).Scan(&validated, &definition) + if errors.Is(err, sql.ErrNoRows) { + return errors.New("captain_turn_requests_tool_approval_identity constraint is missing") + } + if err != nil { + return fmt.Errorf("read captain_turn_requests_tool_approval_identity: %w", err) + } + normalized := strings.ToLower(strings.Join(strings.Fields(definition), " ")) + if !validated || strings.Contains(normalized, "credential_id is not null") { + return fmt.Errorf("captain_turn_requests_tool_approval_identity is invalid: %s", definition) + } + for _, required := range []string{ + "prompt_run_id is not null", "turn_id is not null", + "model_call_id is not null", "tool_call_id is not null", + } { + if !strings.Contains(normalized, required) { + return fmt.Errorf("captain_turn_requests_tool_approval_identity omits %q: %s", required, definition) + } + } return nil } diff --git a/migrations/migrations_test.go b/migrations/migrations_test.go index cee76efd..e241d7e7 100644 --- a/migrations/migrations_test.go +++ b/migrations/migrations_test.go @@ -19,6 +19,8 @@ func TestSchemaBundleContainsGavelIntegrationContract(t *testing.T) { "20_prompt_runs_and_plans.pg.hcl", "21_plans.pg.hcl", "30_execution.pg.hcl", + "31_execution_events.pg.hcl", + "32_execution_approvals.pg.hcl", "40_artifacts.pg.hcl", "50_constraints.sql", "51_state_triggers.sql", @@ -27,7 +29,6 @@ func TestSchemaBundleContainsGavelIntegrationContract(t *testing.T) { "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", @@ -36,6 +37,8 @@ func TestSchemaBundleContainsGavelIntegrationContract(t *testing.T) { "70_prompt_run_runtime.sql", "71_session_storage_params.sql", "72_ingest_storage_params.sql", + "73_normalize_session_cwd.sql", + "74_turn_request_approval_identity.sql", } for _, name := range expectedFiles { if _, err := fs.Stat(schemaFS, name); err != nil { @@ -82,8 +85,17 @@ func TestSchemaBundleContainsGavelIntegrationContract(t *testing.T) { `column "turn_id"`, `column "prompt_run_id"`, `column "iteration_id"`, + ) + assertContainsAll(t, "31_execution_events.pg.hcl", + `table "captain_messages"`, `table "captain_events"`, + ) + assertContainsAll(t, "32_execution_approvals.pg.hcl", + `table "captain_session_mcp_credentials"`, + `column "secret_hash"`, + `column "policy"`, `table "captain_turn_requests"`, + `column "credential_id"`, `column "state"`, `column "version"`, ) @@ -142,6 +154,14 @@ func TestSchemaBundleContainsGavelIntegrationContract(t *testing.T) { "ALTER TABLE public.captain_messages SET (\n fillfactor", ) assertContainsNone(t, "30_execution.pg.hcl", "fillfactor", "autovacuum_") + assertContainsNone(t, "31_execution_events.pg.hcl", "fillfactor", "autovacuum_") + assertContainsNone(t, "32_execution_approvals.pg.hcl", "fillfactor", "autovacuum_") + assertContainsAll(t, "74_turn_request_approval_identity.sql", + "-- phase: post", + "ambiguous legacy tool approval identity", + "DROP CONSTRAINT IF EXISTS captain_turn_requests_tool_approval_identity", + "VALIDATE CONSTRAINT captain_turn_requests_tool_approval_identity", + ) for _, name := range expectedFiles { assertContainsNone(t, name, `table "captain_outbox"`, @@ -156,7 +176,6 @@ func TestSchemaBundleContainsGavelIntegrationContract(t *testing.T) { "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", @@ -221,11 +240,40 @@ func TestApplyHoldsMigrationLockAcrossMigration(t *testing.T) { events = append(events, "migrate") return nil }, + verify: func(context.Context, string) error { + events = append(events, "verify") + return nil + }, }) if err != nil { t.Fatalf("apply: %v", err) } - assertEventsEqual(t, events, []string{"lock", "migrate", "unlock"}) + assertEventsEqual(t, events, []string{"lock", "migrate", "verify", "unlock"}) +} + +func TestApplyReleasesMigrationLockOnVerificationFailure(t *testing.T) { + t.Parallel() + + var events []string + verificationErr := errors.New("constraint drifted") + 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 + }, + verify: func(context.Context, string) error { + events = append(events, "verify") + return verificationErr + }, + }) + if !errors.Is(err, verificationErr) { + t.Fatalf("apply error = %v, want errors.Is(_, %v)", err, verificationErr) + } + assertEventsEqual(t, events, []string{"lock", "migrate", "verify", "unlock"}) } func TestApplyReleasesMigrationLockOnMigrationFailure(t *testing.T) { diff --git a/pkg/ai/adapter_models.go b/pkg/ai/adapter_models.go index 06c0e3b0..6e662c21 100644 --- a/pkg/ai/adapter_models.go +++ b/pkg/ai/adapter_models.go @@ -23,8 +23,9 @@ var resolveModelRows = ResolveModels // 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 { +// cache instead of hitting providers every time; refresh bypasses that cache and +// re-queries every provider listing. +func fetchAPIModels(backends []Backend, probe AuthProbe, refresh bool) map[Backend]modelFetch { apis := map[Backend]bool{} for _, b := range backends { if b.Kind() != "api" { @@ -46,7 +47,7 @@ func fetchAPIModels(backends []Backend, probe AuthProbe) map[Backend]modelFetch wg.Add(1) go func(backend Backend) { defer wg.Done() - rows, err := resolveModelRows(context.Background(), ResolveOptions{Backend: backend, UseTokens: true}) + rows, err := resolveModelRows(context.Background(), ResolveOptions{Backend: backend, UseTokens: true, Refresh: refresh}) m := liveModelDefs(rows, backend) mu.Lock() out[backend] = modelFetch{models: m, err: err} diff --git a/pkg/ai/adapters.go b/pkg/ai/adapters.go index b14251e4..d2c66384 100644 --- a/pkg/ai/adapters.go +++ b/pkg/ai/adapters.go @@ -17,9 +17,11 @@ import ( // 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"` + 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"` + IncludeDisabled bool `flag:"disabled" help:"Include disabled models" default:"false"` + NoCache bool `flag:"no-cache" help:"Bypass the persisted model and OpenRouter pricing caches and re-query both live" default:"false"` } // AdapterStatus is the resolved auth/availability of a single agent adapter @@ -31,16 +33,19 @@ type AdapterStatus struct { // Provider and Mode are the two axes Backend is a pair of. They are carried // on the wire so the whoami page can group and filter cards from registry // truth instead of re-deriving the mapping in TypeScript. - Provider string `json:"provider"` - Mode string `json:"mode"` - 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"` + Provider string `json:"provider"` + Mode string `json:"mode"` + Authenticated bool `json:"authenticated"` + AuthMethod string `json:"authMethod,omitempty"` + AuthDetail string `json:"authDetail,omitempty"` + Binary string `json:"binary,omitempty"` + BinaryMissing string `json:"binaryMissing,omitempty"` + DependencyMissing string `json:"dependencyMissing,omitempty"` + Provisioner string `json:"provisioner,omitempty"` + RuntimeError string `json:"runtimeError,omitempty"` + ModelCount int `json:"modelCount"` + Models []string `json:"models,omitempty"` + ModelError string `json:"modelError,omitempty"` ModelDetails []ModelDef `json:"modelDetails,omitempty"` @@ -60,7 +65,7 @@ func (a AdapterStatus) Ready() bool { return false } if a.Type == "cli" { - return a.Binary != "" + return (a.Binary != "" || a.Provisioner != "") && a.DependencyMissing == "" && a.RuntimeError == "" } return true } @@ -212,7 +217,13 @@ func resolveAdapter(backend Backend, p AuthProbe) AdapterStatus { } if cli, ok := cliAdapters()[backend]; ok { - if path, err := p.LookPath(cli.binary); err == nil { + if runtime, custom := probeRuntime(backend); custom { + st.Binary = runtime.Binary + st.BinaryMissing = runtime.BinaryMissing + st.DependencyMissing = runtime.DependencyMissing + st.Provisioner = runtime.Provisioner + st.RuntimeError = runtime.Error + } else if path, err := p.LookPath(cli.binary); err == nil { st.Binary = path } else { st.BinaryMissing = cli.binary @@ -273,7 +284,7 @@ func ProbeAdapters(opts WhoamiOptions, probe AuthProbe) ([]AdapterStatus, error) var models map[Backend]modelFetch var codexModels modelFetch if opts.Models { - models = fetchAPIModels(backends, probe) + models = fetchAPIModels(backends, probe, opts.NoCache) codexModels = fetchCodexModels(backends, probe) } diff --git a/pkg/ai/adapters_test.go b/pkg/ai/adapters_test.go index 5e8c512b..f0dfb826 100644 --- a/pkg/ai/adapters_test.go +++ b/pkg/ai/adapters_test.go @@ -263,13 +263,53 @@ func TestProbeAdaptersFiltersNoisyOpenAIModelsForDirectAPI(t *testing.T) { 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"} { + for _, hidden := range []string{"gpt-realtime-2.1", "gpt-image-2", "gpt-audio-1.5", "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 TestProbeAdaptersNoCacheBypassesPersistedModelCache(t *testing.T) { + for _, tc := range []struct { + name string + noCache bool + wantRefresh bool + }{ + {name: "default reuses the persisted resolve", noCache: false, wantRefresh: false}, + {name: "no-cache re-resolves live", noCache: true, wantRefresh: true}, + } { + t.Run(tc.name, func(t *testing.T) { + prev := resolveModelRows + resolved := false + resolveModelRows = func(_ context.Context, opts ResolveOptions) ([]ResolvedModel, error) { + resolved = true + if opts.Refresh != tc.wantRefresh { + t.Fatalf("ResolveOptions.Refresh = %v, want %v", opts.Refresh, tc.wantRefresh) + } + return []ResolvedModel{ + {Model: Model{ID: "anthropic/claude-opus-5", Backend: BackendAnthropic, Label: "Opus 5", ReleaseDate: "2026-05-01"}, Live: true}, + }, nil + } + t.Cleanup(func() { resolveModelRows = prev }) + + adapters, err := ProbeAdapters( + WhoamiOptions{Backend: string(BackendAnthropic), Models: true, NoCache: tc.noCache}, + fakeProbe(map[string]string{"ANTHROPIC_API_KEY": "sk-ant-test"}, nil, nil, "/home/u"), + ) + if err != nil { + t.Fatalf("ProbeAdapters: %v", err) + } + if !resolved { + t.Fatal("expected the model resolver to run for an authenticated API backend") + } + if len(adapters) != 1 || !stringSliceContains(adapters[0].Models, "claude-opus-5") { + t.Fatalf("adapters = %+v, want the resolved model", adapters) + } + }) + } +} + 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 diff --git a/pkg/ai/agent.go b/pkg/ai/agent.go index 7179ef77..68a393b2 100644 --- a/pkg/ai/agent.go +++ b/pkg/ai/agent.go @@ -207,15 +207,22 @@ func (a *Agent) accrue(resp *Response) Cost { // 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 { + return PriceUsage(backend, model, resp.Usage, resp.CostUSD) +} + +// PriceUsage is PriceResponse for callers that hold a bare Usage plus the +// provider's reported total (stream events, persisted model calls) rather than +// a full Response. +func PriceUsage(backend Backend, model string, usage Usage, providerCostUSD float64) Cost { cost := Cost{ Model: model, - InputTokens: resp.Usage.InputTokens, - OutputTokens: resp.Usage.OutputTokens, - ReasoningTokens: resp.Usage.ReasoningTokens, - CacheReadTokens: resp.Usage.CacheReadTokens, - CacheWriteTokens: resp.Usage.CacheWriteTokens, - TotalTokens: resp.Usage.TotalTokens(), - ProviderCostUSD: resp.CostUSD, + InputTokens: usage.InputTokens, + OutputTokens: usage.OutputTokens, + ReasoningTokens: usage.ReasoningTokens, + CacheReadTokens: usage.CacheReadTokens, + CacheWriteTokens: usage.CacheWriteTokens, + TotalTokens: usage.TotalTokens(), + ProviderCostUSD: providerCostUSD, } // The pricing registry is keyed on OpenRouter-style ids (provider/model); // try the backend-prefixed id first, then the bare model. @@ -232,6 +239,18 @@ func PriceResponse(backend Backend, model string, resp *Response) Cost { return cost } +// ContextWindowFor returns a model's context window from the pricing registry, +// or 0 when unknown, so persisted model calls can record context occupancy +// server-side instead of relying on the UI's catalog for the denominator. +func ContextWindowFor(backend Backend, model string) int { + for _, id := range PricingIDs(backend, model) { + if info, ok := pricing.GetModelInfo(id); ok && info.ContextWindow > 0 { + return info.ContextWindow + } + } + return 0 +} + // 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 diff --git a/pkg/ai/agent/commit/commit.go b/pkg/ai/agent/commit/commit.go index c0e40460..065fc9b4 100644 --- a/pkg/ai/agent/commit/commit.go +++ b/pkg/ai/agent/commit/commit.go @@ -78,6 +78,7 @@ type Hook struct { anchor string // SHA the chain fixes up onto, set by the run's first commit subject string // resolved once, shared by every commit in the run fixups int // fixup commits cut so far; a chain of 0 needs no squash + failed error // first commit failure, so the agent-phase sweep does not repeat it } // New builds a Hook for a policy. @@ -119,6 +120,7 @@ func commitPhase(p api.CommitPhase) agent.Phase { func (h *Hook) Post(hc *agent.HookContext, phase agent.Phase) error { if h.commitsAt(phase) { if err := h.commit(hc, phase); err != nil { + h.failed = err return err } } @@ -130,11 +132,18 @@ func (h *Hook) Post(hc *agent.HookContext, phase agent.Phase) error { // commitsAt reports whether a commit is cut at this phase: the declared one, or // the agent-phase sweep that closes out a per-turn policy. +// +// The sweep is skipped once a commit has already failed. It exists to catch a +// turn that errored before the runner could dispatch its turn phase, so it has +// nothing to add after a turn-phase attempt that resolved the same paths and +// failed on them — it would only re-derive the identical error, which the runner +// then joins onto the one already propagating, printing the same failure twice +// under two different phase names. func (h *Hook) commitsAt(phase agent.Phase) bool { if phase == commitPhase(h.Phase()) { return true } - return h.Phase() == api.CommitOnTurn && phase == agent.PhaseAgent + return h.Phase() == api.CommitOnTurn && phase == agent.PhaseAgent && h.failed == nil } // commit resolves and cuts one commit. Finding nothing to commit is a normal @@ -184,6 +193,11 @@ func (h *Hook) commit(hc *agent.HookContext, phase agent.Phase) error { if h.Do == nil && h.EffectiveGates() == api.CommitGatesFull { return fmt.Errorf("commit: gates: full runs the host's pre-commit pipeline, which captain does not have — supply Hook.Do or lower gates to cheap") } + // Announced before the cut, not only after it: a host pipeline on gates: full + // runs lint and pre-commit hooks, which is long enough that silence reads as a + // hang. A turn that resolved no paths says nothing at all — silence there + // already means "nothing happened", and a line per read-only turn is noise. + hc.Notify("[post-%s] committing %d file(s)", phase, len(paths)) sha, err := h.cut(hc, plan) return h.record(hc, plan, sha, err) } @@ -233,6 +247,10 @@ func (h *Hook) record(hc *agent.HookContext, plan Plan, sha string, err error) e return err } if sha == "" { + // The paths resolved but the pipeline staged nothing (an earlier phase + // already took them). Said out loud because the "committing" line above + // has already promised a commit. + hc.Notify("[post-%s] nothing left to stage", plan.Phase) return nil } message := plan.Subject @@ -245,9 +263,19 @@ func (h *Hook) record(hc *agent.HookContext, plan Plan, sha string, err error) e h.anchor = sha // amend rewrote the anchor } hc.Workspace().AddCommit(sha, message) + hc.Notify("[post-%s] committed %s: %s", plan.Phase, shortSHA(sha), message) return nil } +// shortSHA abbreviates to git's conventional display width. A host pipeline may +// hand back a short hash already, so this truncates rather than assuming 40. +func shortSHA(sha string) string { + if len(sha) <= 7 { + return sha + } + return sha[:7] +} + // squash collapses the fixup chain back into its anchor. A chain of zero fixups // is already one commit, so the rebase is skipped rather than run as a no-op // that could still fail. @@ -255,6 +283,7 @@ func (h *Hook) squash(hc *agent.HookContext) error { if h.fixups == 0 || h.anchor == "" || h.DryRun { return nil } + fixups := h.fixups dir, err := workDir(hc) if err != nil { return err @@ -274,6 +303,7 @@ func (h *Hook) squash(hc *agent.HookContext) error { return err } h.anchor = head + hc.Notify("[post-run] squashed %d fixup(s) into %s", fixups, shortSHA(head)) return nil } @@ -309,11 +339,12 @@ func (h *Hook) resolvePaths(hc *agent.HookContext, dir string) ([]string, api.Co if mode == api.CommitStageWorktree { return dirty, mode, nil } - changed := attributable(dir, dirty, hc.Workspace().Changed) + recorded := hc.Workspace().Changed + changed := attributable(dir, recordBase(hc), dirty, recorded) if len(changed) == 0 { // Staging the tree anyway would sweep the caller's own uncommitted work // into an agent commit. Refusing is the whole point of the changed mode. - return nil, mode, fmt.Errorf("commit: %s has uncommitted changes but none are attributable to this run (%d dirty path(s), 0 recorded as agent-modified); refusing to stage a tree that may hold your own work — run with an isolated worktree, or set stage: worktree to commit everything", dir, len(dirty)) + return nil, mode, unattributableErr(dir, dirty, recorded) } return changed, mode, nil } @@ -331,33 +362,67 @@ func (h *Hook) stageMode(hc *agent.HookContext) api.CommitStage { return api.CommitStageChanged } +// unattributableErr explains a refusal. The two ways of arriving here are +// indistinguishable in the tree but call for different investigations, so they +// get different messages: an agent that recorded nothing may have written +// through tools the runner does not track (a shell redirect, an MCP server), +// while an agent that recorded files none of which are dirty here needs those +// paths named — reporting them as "none recorded" sends the reader looking for +// a bug in the agent instead of at where its edits actually landed. +func unattributableErr(dir string, dirty, recorded []string) error { + const advice = "refusing to stage a tree that may hold your own work — run with an isolated worktree, or set stage: worktree to commit everything" + if len(recorded) == 0 { + return fmt.Errorf("commit: %s has %d uncommitted path(s) but the run recorded no file edits (an agent that writes through the shell or an MCP tool is not tracked); %s", dir, len(dirty), advice) + } + return fmt.Errorf("commit: %s has %d uncommitted path(s), none of them among the %d file(s) the run recorded editing (%s) — those edits are either in another tree or already committed; %s", + dir, len(dirty), len(recorded), strings.Join(elide(recorded, 3), ", "), advice) +} + +// elide renders at most limit entries, summarising the rest by count so a run +// that touched a hundred files still produces a readable one-line error. The +// capped slice keeps the append off the caller's backing array. +func elide(paths []string, limit int) []string { + if len(paths) <= limit { + return paths + } + return append(paths[:limit:limit], fmt.Sprintf("and %d more", len(paths)-limit)) +} + // attributable intersects git's dirty set with the paths the agent is recorded // as having modified, so a commit can never contain a file the run did not -// touch. Recorded paths may be absolute (the agent reports what it wrote), so -// each is normalized against the working dir first. -func attributable(dir string, dirty, changed []string) []string { +// touch. +// +// The two sides arrive in different namespaces and neither can be converted to +// the other's by string surgery: dirty paths are relative to root, while +// recorded paths are relative to recordBase — which is the same directory only +// when the run was launched from the top of its working tree. Resolving both to +// absolute paths is what lets a run started in a subdirectory attribute its own +// edits. +func attributable(root, base string, dirty, changed []string) []string { recorded := make(map[string]bool, len(changed)) for _, c := range changed { - recorded[normalizePath(dir, c)] = true + recorded[resolveAgainst(base, c)] = true } var out []string for _, p := range dirty { - if recorded[filepath.ToSlash(p)] { + if recorded[resolveAgainst(root, p)] { out = append(out, p) } } return out } -// normalizePath renders path relative to dir, matching git's repo-relative, -// forward-slashed form. -func normalizePath(dir, path string) string { - if filepath.IsAbs(path) { - if rel, err := filepath.Rel(dir, path); err == nil && !strings.HasPrefix(rel, "..") { - path = rel +// resolveAgainst renders path as an absolute, forward-slashed path anchored on +// base. A path that is already absolute stands on its own; base may be empty, +// in which case only absolute paths can ever match. +func resolveAgainst(base, path string) string { + if !filepath.IsAbs(path) { + if base == "" { + return filepath.ToSlash(filepath.Clean(path)) } + path = filepath.Join(base, path) } - return filepath.ToSlash(path) + return filepath.ToSlash(filepath.Clean(path)) } // resolveAnchor returns the SHA this commit fixes up onto, or "" when it is @@ -450,15 +515,34 @@ func firstLine(body string) string { return "" } -// workDir is the directory commits are cut in: the run's cwd (an isolated -// worktree when there is one), else the repo root. +// workDir is the directory commits are cut in: the root of the working tree the +// run's cwd sits in (an isolated worktree when there is one). It is the root +// rather than the cwd itself because every path that reaches staging comes from +// `git status`, which reports repo-relative paths however deep it is invoked — +// a run started in a monorepo subdirectory would otherwise stage pathspecs +// against the wrong base. func workDir(hc *agent.HookContext) (string, error) { ws := hc.Workspace() - if ws.Cwd != "" { - return ws.Cwd, nil + // Cwd leads: an isolated run keeps Repo pointing at the checkout it branched + // from, and committing there rather than in the worktree would defeat the + // isolation entirely. + dir := ws.Cwd + if dir == "" { + dir = ws.Repo } - if ws.Repo != "" { - return ws.Repo, nil + if dir == "" { + return "", fmt.Errorf("commit: no working directory on the run's workspace (set Runner.Cwd or Runner.Repo)") } - return "", fmt.Errorf("commit: no working directory on the run's workspace (set Runner.Cwd or Runner.Repo)") + return gitRoot(dir) +} + +// recordBase is the directory the workspace's recorded paths are relative to. +// Runner.recordEvent relativizes each edit against ws.Repo and leaves it +// absolute when there is none, so this is ws.Repo and nothing else — resolving +// against any other directory would silently mis-attribute every path. +func recordBase(hc *agent.HookContext) string { + if repo := hc.Workspace().Repo; repo != "" { + return canonicalDir(repo) + } + return "" } diff --git a/pkg/ai/agent/commit/commit_test.go b/pkg/ai/agent/commit/commit_test.go index 7fb7390e..1652ebfd 100644 --- a/pkg/ai/agent/commit/commit_test.go +++ b/pkg/ai/agent/commit/commit_test.go @@ -175,6 +175,99 @@ func TestAgentPhaseSweepsWorkFromAnErroredTurn(t *testing.T) { } } +// TestFailedTurnCommitIsReportedOnce is the counterpart to the sweep above: the +// turn phase resolved the paths and failed on them, so the sweep has nothing new +// to try. Letting it run again re-derives the identical error, and the runner +// joins that copy onto the one already propagating from the turn — the caller +// then prints the same failure twice, once as `turn hook "commit:turn"` and once +// as `agent hook "commit:turn"`. +func TestFailedTurnCommitIsReportedOnce(t *testing.T) { + dir := newRepo(t) + hc := isolated(dir) + h := New(api.Commit{On: api.CommitOnTurn, Message: "fix: attempt"}) + attempts := 0 + h.Do = func(*agent.HookContext, Plan) (string, error) { + attempts++ + return "", fmt.Errorf("pre-commit gate rejected the change set") + } + + write(t, dir, "fix.go", "package main\n") + changed(hc, "fix.go") + turnErr := h.Post(hc, agent.PhaseTurn) + if turnErr == nil { + t.Fatal("turn phase should surface the commit failure") + } + hc.Failed = true + + if err := h.Post(hc, agent.PhaseAgent); err != nil { + t.Errorf("agent sweep repeated the turn's failure: %v", err) + } + if attempts != 1 { + t.Errorf("commit attempted %d times, want 1 — the sweep must not retry a failure", attempts) + } +} + +// TestCommitsAreNarratedOnTheWorkspace covers the other half of the silence +// problem: hooks act between turns, where the provider transcript has nothing to +// say, so a run's commits used to leave no trace a reader could follow. +func TestCommitsAreNarratedOnTheWorkspace(t *testing.T) { + dir := newRepo(t) + hc := isolated(dir) + h := New(api.Commit{On: api.CommitOnTurn, Message: "feat: add greeting"}) + + turn(t, h, hc, dir, "one.go", "package main\n") + turn(t, h, hc, dir, "two.go", "package main\n") + if err := finish(t, h, hc); err != nil { + t.Fatalf("finish: %v", err) + } + + var got []string + for _, notice := range hc.Workspace().Notices { + got = append(got, notice.Text) + if notice.At.IsZero() { + t.Errorf("notice %q has no timestamp; it cannot be sorted back among the turns", notice.Text) + } + } + // The anchor turn, its fixup, and the squash that collapses them — the whole + // shape of a per-turn policy, readable from the notices alone. + want := []string{ + "[post-turn] committing 1 file(s)", + "[post-turn] committed", + "[post-turn] committing 1 file(s)", + "[post-turn] committed", + "[post-run] squashed 1 fixup(s) into ", + } + if len(got) != len(want) { + t.Fatalf("notices = %v, want %d entries shaped like %v", got, len(want), want) + } + for i, prefix := range want { + if !strings.HasPrefix(got[i], prefix) { + t.Errorf("notice %d = %q, want prefix %q", i, got[i], prefix) + } + } +} + +// TestAgentSweepStillRunsAfterAnUnrelatedTurnError guards the narrowing above: +// only a *commit* failure disarms the sweep. A turn that failed for any other +// reason never reached the hook, and its work must still be made durable. +func TestAgentSweepStillRunsAfterAnUnrelatedTurnError(t *testing.T) { + dir := newRepo(t) + hc := isolated(dir) + h := New(api.Commit{On: api.CommitOnTurn, Message: "feat: partial"}) + + // The provider errored mid-turn: the file landed, PhaseTurn never dispatched. + write(t, dir, "one.go", "package main\n") + changed(hc, "one.go") + hc.Failed = true + + if err := h.Post(hc, agent.PhaseAgent); err != nil { + t.Fatalf("agent sweep: %v", err) + } + if !isClean(t, dir) { + t.Errorf("agent-phase sweep left work uncommitted:\n%s", mustGit(t, dir, "status", "--porcelain")) + } +} + func TestOutcomeGates(t *testing.T) { cases := []struct { name string diff --git a/pkg/ai/agent/commit/git.go b/pkg/ai/agent/commit/git.go index a2865b81..34ce4b90 100644 --- a/pkg/ai/agent/commit/git.go +++ b/pkg/ai/agent/commit/git.go @@ -23,6 +23,34 @@ func git(dir string, args ...string) (string, error) { return strings.TrimSpace(res.Stdout), nil } +// gitRoot resolves the working tree dir belongs to. Every path this package +// handles is relative to that root — `git status` reports repo-relative paths +// whatever directory it is invoked from — so a run launched inside a +// subdirectory has to be lifted to the root before its paths mean anything. +func gitRoot(dir string) (string, error) { + root, err := git(dir, "rev-parse", "--show-toplevel") + if err != nil { + return "", fmt.Errorf("commit: %s is not inside a git working tree: %w", dir, err) + } + return canonicalDir(root), nil +} + +// canonicalDir renders dir absolute with its symlinks resolved, so two bases +// arrived at by different routes — git's own output versus the caller's cwd — +// compare and join equal. On macOS this is what keeps /tmp and /private/tmp +// from looking like different trees. +func canonicalDir(dir string) string { + abs, err := filepath.Abs(dir) + if err != nil { + return filepath.Clean(dir) + } + resolved, err := filepath.EvalSymlinks(abs) + if err != nil { + return abs + } + return resolved +} + // dirtyPaths lists every repo-relative path with uncommitted work in dir, in // git's own order. Ignored files are absent — `git status` excludes them — which // is what keeps build artifacts out of the chain regardless of gate level. diff --git a/pkg/ai/agent/commit/stage_test.go b/pkg/ai/agent/commit/stage_test.go index f0bb5f0b..dba3ece1 100644 --- a/pkg/ai/agent/commit/stage_test.go +++ b/pkg/ai/agent/commit/stage_test.go @@ -2,6 +2,8 @@ package commit import ( "fmt" + "os" + "path/filepath" "strings" "testing" @@ -33,25 +35,105 @@ func TestSharedTreeCommitsOnlyAgentFiles(t *testing.T) { } } -// TestSharedTreeRefusesUnattributableWork is the test that protects the user's -// working tree: when nothing in a dirty shared checkout can be attributed to the -// run, the only safe move is to fail loudly rather than sweep it all in. -func TestSharedTreeRefusesUnattributableWork(t *testing.T) { - dir := newRepo(t) - hc := shared(dir) - h := New(api.Commit{On: api.CommitOnAgent, Message: "feat: nothing of mine"}) +// TestSharedTreeCommitsFromSubdirectory: a run launched from a subdirectory of +// the repo — a monorepo package, say — records its edits relative to that +// subdirectory, while git reports dirt relative to the repo root whatever the +// cwd. Attribution has to reconcile the two bases; treating them as one +// namespace makes every such run refuse to commit work it demonstrably did. +// +// Both recorded forms are covered because claude.RelativePath emits either, +// depending on how far outside the run's directory the edited file sits. +func TestSharedTreeCommitsFromSubdirectory(t *testing.T) { + cases := []struct { + name string + sub string + // record builds the entry recordEvent would have stored for + // /packages/agent.go, given the repo root. + record func(root string) string + }{ + { + name: "one level up, recorded relative", + sub: "apps", + record: func(string) string { return "../packages/agent.go" }, + }, + { + name: "further away, recorded absolute", + sub: filepath.Join("apps", "storybook"), + record: func(root string) string { + return filepath.Join(root, "packages", "agent.go") + }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + dir := newRepo(t) + if err := os.MkdirAll(filepath.Join(dir, tc.sub), 0o755); err != nil { + t.Fatalf("create subdirectory: %v", err) + } + hc := shared(filepath.Join(dir, tc.sub)) // Repo and Cwd are the subdirectory + h := New(api.Commit{On: api.CommitOnAgent, Message: "feat: agent work"}) - write(t, dir, "mine.go", "// the user's own uncommitted work\n") + write(t, dir, "mine.go", "// the user's own uncommitted work\n") + write(t, dir, "packages/agent.go", "package main\n") + changed(hc, tc.record(canonicalDir(dir))) - err := h.Post(hc, agent.PhaseAgent) - if err == nil { - t.Fatalf("expected an error; instead the tree was committed as %v", subjects(t, dir)) + if err := h.Post(hc, agent.PhaseAgent); err != nil { + t.Fatalf("agent phase: %v", err) + } + if got := filesInHead(t, dir, "HEAD"); fmt.Sprint(got) != fmt.Sprint([]string{"packages/agent.go"}) { + t.Errorf("commit touched %v, want only packages/agent.go", got) + } + if status := mustGit(t, dir, "status", "--porcelain"); !strings.Contains(status, "mine.go") { + t.Errorf("the user's file should still be uncommitted, status:\n%s", status) + } + }) } - if !strings.Contains(err.Error(), "attributable") { - t.Errorf("error should explain why it refused, got: %v", err) +} + +// TestRefusalReportsWhatWasRecorded is the test that protects the user's +// working tree: when nothing in a dirty shared checkout can be attributed to +// the run, the only safe move is to fail loudly rather than sweep it all in. +// The two ways of coming up empty need different messages, because they call +// for different fixes. Nothing recorded is that safety refusal working as +// designed; files recorded but none of them dirty here means the run edited +// another tree, and reporting that as "0 recorded" sends the reader hunting for +// a bug in the agent instead. +func TestRefusalReportsWhatWasRecorded(t *testing.T) { + cases := []struct { + name string + record []string + wantErr string + }{ + { + name: "the agent recorded nothing", + wantErr: "recorded no file edits", + }, + { + name: "the agent edited a different tree", + record: []string{"/somewhere/else/agent.go"}, + wantErr: "/somewhere/else/agent.go", + }, } - if commitCount(t, dir) != 1 { - t.Errorf("nothing should have been committed, subjects: %v", subjects(t, dir)) + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + dir := newRepo(t) + hc := shared(dir) + h := New(api.Commit{On: api.CommitOnAgent, Message: "feat: nothing of mine"}) + + write(t, dir, "mine.go", "// the user's own uncommitted work\n") + changed(hc, tc.record...) + + err := h.Post(hc, agent.PhaseAgent) + if err == nil { + t.Fatalf("expected an error; instead the tree was committed as %v", subjects(t, dir)) + } + if !strings.Contains(err.Error(), tc.wantErr) { + t.Errorf("error should mention %q, got: %v", tc.wantErr, err) + } + if commitCount(t, dir) != 1 { + t.Errorf("nothing should have been committed, subjects: %v", subjects(t, dir)) + } + }) } } @@ -278,8 +360,10 @@ func TestDoCallbackReceivesAResolvedPlan(t *testing.T) { t.Fatalf("agent phase: %v", err) } - if got.Dir != dir || got.Phase != agent.PhaseAgent || got.Subject != "feat: hosted" { - t.Errorf("plan = %+v, want dir/phase/subject resolved", got) + // Dir is the resolved working-tree root, not whatever the caller passed: + // Plan.Paths are repo-relative, so the host needs the base they hang off. + if got.Dir != canonicalDir(dir) || got.Phase != agent.PhaseAgent || got.Subject != "feat: hosted" { + t.Errorf("plan = %+v, want dir %s with phase/subject resolved", got, canonicalDir(dir)) } if got.Gates != api.CommitGatesFull || got.Stage != api.CommitStageWorktree { t.Errorf("plan policy = %+v, want gates full and worktree staging", got) diff --git a/pkg/ai/agent/runner.go b/pkg/ai/agent/runner.go index 450eda3e..ad240e6a 100644 --- a/pkg/ai/agent/runner.go +++ b/pkg/ai/agent/runner.go @@ -16,6 +16,7 @@ import ( "errors" "fmt" "strings" + "time" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/ai/history" @@ -110,6 +111,25 @@ type HookContext struct { // Failed is true when the generate/verify run itself returned an error // (a provider failure, not a failing verdict). Failed bool + + // emit publishes an event on the run's stream, so what a hook does between + // turns reaches the same renderers as what the model does during them. Set by + // Runner.Run; nil in a hand-built context, which Notify tolerates. + emit func(ai.Event) +} + +// Notify reports one thing this hook did, in the run's own voice. It reaches the +// live stream as an ai.EventSystem and is buffered on the workspace with its +// timestamp, so a caller can persist it into the transcript once the run's +// session id is known — a hook firing mid-turn cannot know it yet. +// +// Purely informational: a hook that failed returns an error, it does not Notify. +func (hc *HookContext) Notify(format string, args ...any) { + text := fmt.Sprintf(format, args...) + hc.Workspace().AddNotice(time.Now(), string(hc.Phase), text) + if hc.emit != nil { + hc.emit(ai.Event{Kind: ai.EventSystem, Text: text}) + } } // Workspace returns the run's working-dir state, allocating it if needed (so it @@ -231,6 +251,14 @@ func (r *Runner[T]) Run(ctx context.Context) (Result[T], error) { Scope: scope, Hooks: r.Hooks, } + // Hook notices join the model's own events on one stream, attributed to the + // iteration whose boundary the hook is standing on — hc.Iteration still names + // the turn that just completed when PhaseTurn dispatches. + hc.emit = func(ev ai.Event) { + if r.OnEvent != nil { + r.OnEvent(hc.Iteration, ev) + } + } result := Result[T]{Response: resp} for _, h := range r.Hooks { diff --git a/pkg/ai/agent/verify/workflow.go b/pkg/ai/agent/verify/workflow.go index b730ed48..86013b7a 100644 --- a/pkg/ai/agent/verify/workflow.go +++ b/pkg/ai/agent/verify/workflow.go @@ -81,7 +81,7 @@ func rejectJudgeOverrides(path string, tmpl *prompt.Template, provider ai.Provid if probe.Sandbox != nil { return fmt.Errorf("verify prompt %q declares a sandbox; judge hooks run on the run's provider and cannot relocate", path) } - if declared := strings.TrimSpace(probe.Model.Name); declared != "" && declared != provider.GetModel() { + if declared := strings.TrimSpace(probe.Name); declared != "" && declared != provider.GetModel() { return fmt.Errorf("verify prompt %q declares model %q but judge hooks run on the run's provider (%s); remove the model or match it", path, declared, provider.GetModel()) } diff --git a/pkg/ai/availability.go b/pkg/ai/availability.go new file mode 100644 index 00000000..0f8524ed --- /dev/null +++ b/pkg/ai/availability.go @@ -0,0 +1,107 @@ +package ai + +import ( + "fmt" + "strings" + "time" + + "github.com/flanksource/captain/pkg/api" +) + +// AvailabilityForAdapter projects provider readiness into presentation-safe +// status without exposing credentials or authentication details. +func AvailabilityForAdapter(status AdapterStatus) api.Availability { + backend := Backend(status.Backend) + label := runtimeLabel(backend) + switch { + case status.Disabled: + return api.Availability{State: api.AvailabilityDisabled, + Reason: "Disabled by " + status.DisabledReason + " in Captain configuration.", + Remediation: "Enable " + status.DisabledReason + " on the Whoami page, then refresh."} + case status.RuntimeError != "": + return api.Availability{State: api.AvailabilityUnavailable, + Reason: fmt.Sprintf("%s prerequisites could not be inspected.", label), + Remediation: "Review the runtime diagnostics on the Whoami page, then refresh."} + case status.DependencyMissing != "": + return api.Availability{State: api.AvailabilityMissingDependency, + Reason: fmt.Sprintf("%s requires %s, which is not available.", label, quoted(status.DependencyMissing)), + Remediation: fmt.Sprintf("Install %s or add it to PATH, then refresh.", quoted(status.DependencyMissing))} + case status.Type == "cli" && status.Binary == "": + binary := status.BinaryMissing + if binary == "" { + binary = strings.TrimSpace(string(backend)) + } + return api.Availability{State: api.AvailabilityMissingExecutable, + Reason: fmt.Sprintf("%s was not found on PATH.", quoted(binary)), + Remediation: fmt.Sprintf("Install %s or add it to PATH, then refresh.", label)} + case !status.Authenticated && status.Type == "cli": + return api.Availability{State: api.AvailabilityNotAuthenticated, + Reason: label + " is installed but not authenticated.", + Remediation: "Authenticate with " + loginCommand(backend) + ", then refresh."} + case !status.Authenticated: + return api.Availability{State: api.AvailabilityMissingCredential, + Reason: "No " + label + " credentials are configured.", + Remediation: "Configure credentials on the Whoami page, then refresh."} + default: + return api.Available() + } +} + +// LiveRuntimeCatalog annotates the registry runtime catalog with host readiness. +func LiveRuntimeCatalog() ([]api.RuntimeFamily, error) { + adapters, err := CachedAdapters(time.Now()) + if err != nil { + return nil, err + } + byBackend := make(map[string]AdapterStatus, len(adapters)) + for _, adapter := range adapters { + byBackend[adapter.Backend] = adapter + } + runtimes := api.RuntimeCatalog() + for familyIndex := range runtimes { + for modeIndex := range runtimes[familyIndex].Modes { + mode := &runtimes[familyIndex].Modes[modeIndex] + backend := Backend(mode.Backend) + if mode.Disabled { + continue + } + adapter, ok := byBackend[mode.Backend] + if !ok { + mode.Availability = api.Availability{State: api.AvailabilityUnavailable, + Reason: runtimeLabel(backend) + " readiness was not reported.", + Remediation: "Review the runtime diagnostics on the Whoami page, then refresh."} + continue + } + mode.Availability = AvailabilityForAdapter(adapter) + } + } + return runtimes, nil +} + +func runtimeLabel(backend Backend) string { + family := backend.Family() + if family != "" { + family = strings.ToUpper(family[:1]) + family[1:] + } + switch backend.Mode() { + case api.ModeAPI: + return family + " API" + case api.ModeCLI: + return family + " CLI" + case api.ModeAgent: + return family + " Agent" + case api.ModeCmux: + return family + " cmux" + default: + return string(backend) + } +} + +func loginCommand(backend Backend) string { + if adapter, ok := cliAdapters()[backend]; ok && len(adapter.logins) > 0 { + return quoted(adapter.logins[0].label) + } + return quoted(string(backend) + " login") +} + +func quoted(value string) string { return "`" + value + "`" } diff --git a/pkg/ai/availability_test.go b/pkg/ai/availability_test.go new file mode 100644 index 00000000..7380c093 --- /dev/null +++ b/pkg/ai/availability_test.go @@ -0,0 +1,107 @@ +package ai + +import ( + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/api" +) + +var _ = Describe("runtime availability", func() { + DescribeTable("classifies adapter readiness", + func(input AdapterStatus, state api.AvailabilityState, text string) { + got := AvailabilityForAdapter(input) + Expect(got.State).To(Equal(state)) + if text != "" { + Expect(strings.ToLower(got.Reason)).To(ContainSubstring(strings.ToLower(text))) + } + if !got.IsAvailable() { + Expect(got.Remediation).NotTo(BeEmpty()) + } + }, + Entry("disabled mode", + AdapterStatus{Backend: string(BackendClaudeCmux), Type: "cli", Authenticated: true, Binary: "/bin/claude", + Disabled: true, DisabledReason: "mode cmux"}, + api.AvailabilityDisabled, "mode cmux"), + Entry("missing API credentials", + AdapterStatus{Backend: string(BackendOpenAI), Type: "api"}, + api.AvailabilityMissingCredential, "credentials"), + Entry("local CLI not authenticated", + AdapterStatus{Backend: string(BackendCodexCLI), Type: "cli", Binary: "/bin/codex"}, + api.AvailabilityNotAuthenticated, "authenticated"), + Entry("missing executable", + AdapterStatus{Backend: string(BackendGeminiCLI), Type: "cli", Authenticated: true, BinaryMissing: "gemini"}, + api.AvailabilityMissingExecutable, "gemini"), + Entry("missing runtime dependency", + AdapterStatus{Backend: string(BackendClaudeAgent), Type: "cli", Authenticated: true, DependencyMissing: "npm"}, + api.AvailabilityMissingDependency, "npm"), + Entry("ready", + AdapterStatus{Backend: string(BackendCodexCLI), Type: "cli", Authenticated: true, Binary: "/bin/codex"}, + api.AvailabilityAvailable, ""), + ) + + It("does not expose runtime probe error details", func() { + availability := AvailabilityForAdapter(AdapterStatus{ + Backend: string(BackendCodexCLI), + Type: "cli", + RuntimeError: "inspect token=super-secret: permission denied", + }) + + Expect(availability.State).To(Equal(api.AvailabilityUnavailable)) + Expect(availability.Reason).To(Equal("Codex CLI prerequisites could not be inspected.")) + Expect(availability.Reason).NotTo(ContainSubstring("super-secret")) + }) + + It("uses each adapter's live readiness in the runtime catalog", func() { + previousProbe := adapterProbe + previousCache, previousAt := adapterCache, adapterCacheAt + DeferCleanup(func() { + adapterProbe = previousProbe + adapterCache, adapterCacheAt = previousCache, previousAt + }) + adapterCache, adapterCacheAt = nil, time.Time{} + adapterProbe = func() ([]AdapterStatus, error) { + return []AdapterStatus{ + {Backend: string(BackendOpenAI), Type: "api"}, + {Backend: string(BackendCodexCLI), Type: "cli", Binary: "/bin/codex"}, + {Backend: string(BackendClaudeAgent), Type: "cli", Authenticated: true, DependencyMissing: "npm"}, + }, nil + } + + runtimes, err := LiveRuntimeCatalog() + Expect(err).NotTo(HaveOccurred()) + Expect(runtimeEntry(runtimes, BackendClaudeAgent).CatalogProvider).To(Equal("claude-agent")) + Expect(runtimeEntry(runtimes, BackendOpenAI).CatalogProvider).To(Equal("openai")) + for _, test := range []struct { + backend Backend + state api.AvailabilityState + }{ + {backend: BackendOpenAI, state: api.AvailabilityMissingCredential}, + {backend: BackendCodexCLI, state: api.AvailabilityNotAuthenticated}, + {backend: BackendClaudeAgent, state: api.AvailabilityMissingDependency}, + } { + availability := runtimeAvailability(runtimes, test.backend) + Expect(availability.State).To(Equal(test.state), string(test.backend)) + Expect(availability.Remediation).NotTo(BeEmpty(), string(test.backend)) + } + }) +}) + +func runtimeAvailability(families []api.RuntimeFamily, backend Backend) api.Availability { + return runtimeEntry(families, backend).Availability +} + +func runtimeEntry(families []api.RuntimeFamily, backend Backend) api.RuntimeModeEntry { + for _, family := range families { + for _, mode := range family.Modes { + if mode.Backend == string(backend) { + return mode + } + } + } + Fail("runtime backend " + string(backend) + " not found") + return api.RuntimeModeEntry{} +} diff --git a/pkg/ai/callertools/callertools_suite_test.go b/pkg/ai/callertools/callertools_suite_test.go new file mode 100644 index 00000000..72c2c3f7 --- /dev/null +++ b/pkg/ai/callertools/callertools_suite_test.go @@ -0,0 +1,13 @@ +package callertools_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCallerTools(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Caller Tools Suite") +} diff --git a/pkg/ai/callertools/credential_ginkgo_test.go b/pkg/ai/callertools/credential_ginkgo_test.go new file mode 100644 index 00000000..203caf79 --- /dev/null +++ b/pkg/ai/callertools/credential_ginkgo_test.go @@ -0,0 +1,54 @@ +package callertools_test + +import ( + "context" + "errors" + "sync/atomic" + + "github.com/flanksource/captain/pkg/ai/callertools" + "github.com/flanksource/captain/pkg/api" + "github.com/mark3labs/mcp-go/mcp" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Caller-tool credential lease", func() { + It("exposes only the credential hash and revalidates the persisted lease", func(ctx SpecContext) { + var revoked atomic.Bool + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "account_edit", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { + return map[string]any{"updated": true}, nil + }, + }}, + SessionID: "captain-session-1", + ValidateCredential: func(context.Context) error { + if revoked.Load() { + return errors.New("caller-tool credential is revoked") + } + return nil + }, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(runtime.Close) + + hash := runtime.CredentialHash() + Expect(hash).To(HaveLen(32)) + Expect(string(hash)).NotTo(ContainSubstring("cap_captain-session-1")) + + client := authenticatedClient(ctx, runtime.Endpoint()) + DeferCleanup(client.Close) + request := mcp.CallToolRequest{} + request.Params.Name = "account_edit" + request.Params.Arguments = map[string]any{"id": "acc-1"} + result, err := client.CallTool(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(result.IsError).To(BeFalse()) + + revoked.Store(true) + _, err = client.ListTools(ctx, mcp.ListToolsRequest{}) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/pkg/ai/callertools/runtime.go b/pkg/ai/callertools/runtime.go new file mode 100644 index 00000000..acbf2be2 --- /dev/null +++ b/pkg/ai/callertools/runtime.go @@ -0,0 +1,416 @@ +// Package callertools exposes caller-owned Go tool handlers to out-of-process +// agent runtimes through a private authenticated MCP endpoint. +package callertools + +import ( + "bytes" + "context" + "crypto/rand" + "crypto/sha256" + "crypto/subtle" + "encoding/base64" + "encoding/json" + "fmt" + "net" + "net/http" + "strings" + "sync" + "sync/atomic" + "time" + + aitools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + "github.com/santhosh-tekuri/jsonschema/v6" +) + +const ( + endpointPath = "/mcp" + serverName = "captain" + defaultApprovalTimeout = 5 * time.Minute + // ToolUseIDInputKey carries an out-of-process provider's tool call ID + // through clients that cannot attach MCP request metadata. The runtime + // removes it before policy, schema validation, and handler execution. + ToolUseIDInputKey = "__captain_tool_use_id" +) + +// Options defines one private caller-tool capability. +type Options struct { + Definitions []api.ToolDefinition + Preferences api.ToolPreferences + CanUseTool api.PermissionFunc + SessionID string + ExpiresAt time.Time + // ValidateCredential rechecks the persisted lease on every request and + // immediately before a tool handler executes. + ValidateCredential func(context.Context) error + + ApprovalTimeout time.Duration +} + +// Runtime owns one loopback-only MCP server and its in-memory bearer +// credential. Closing it revokes the capability by shutting down the listener. +type Runtime struct { + definitions map[string]api.ToolDefinition + schemas map[string]*jsonschema.Schema + canUseTool api.PermissionFunc + validate func(context.Context) error + sessionID string + token string + tokenHash [sha256.Size]byte + expiresAt time.Time + + approvalTimeout time.Duration + ctx context.Context + cancel context.CancelFunc + revoked atomic.Bool + + endpoint api.CallerToolEndpoint + server *http.Server + listener net.Listener + + closeOnce sync.Once + closeErr error +} + +// New validates and resolves the tool policy before starting a private server. +func New(options Options) (*Runtime, error) { + if !options.ExpiresAt.IsZero() && !options.ExpiresAt.After(time.Now()) { + return nil, fmt.Errorf("caller-tool credential expiry must be in the future") + } + if options.ApprovalTimeout < 0 { + return nil, fmt.Errorf("caller-tool approval timeout cannot be negative") + } + if options.ApprovalTimeout == 0 { + options.ApprovalTimeout = defaultApprovalTimeout + } + definitions, err := aitools.ResolveDefinitions(options.Definitions, options.Preferences) + if err != nil { + return nil, err + } + if len(definitions) == 0 { + return nil, fmt.Errorf("caller-tool runtime requires at least one enabled tool") + } + token, err := capabilityToken(options.SessionID) + if err != nil { + return nil, err + } + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("listen for caller tools: %w", err) + } + ctx, cancel := context.WithCancel(context.Background()) + runtime := &Runtime{ + definitions: make(map[string]api.ToolDefinition, len(definitions)), + schemas: make(map[string]*jsonschema.Schema, len(definitions)), + canUseTool: options.CanUseTool, + validate: options.ValidateCredential, + sessionID: options.SessionID, + token: token, + tokenHash: sha256.Sum256([]byte(token)), + expiresAt: options.ExpiresAt, + approvalTimeout: options.ApprovalTimeout, + ctx: ctx, + cancel: cancel, + listener: listener, + } + mcpServer := server.NewMCPServer( + "captain-caller-tools", + "1.0.0", + server.WithToolCapabilities(false), + server.WithToolFilter(runtime.filterTools), + server.WithInputSchemaValidation(), + ) + for _, definition := range definitions { + runtime.definitions[definition.Name] = definition + tool, schema, err := mcpTool(definition) + if err != nil { + cancel() + _ = listener.Close() + return nil, err + } + runtime.schemas[definition.Name] = schema + mcpServer.AddTool(tool, runtime.handler(definition)) + } + handler := server.NewStreamableHTTPServer( + mcpServer, + server.WithStateLess(true), + server.WithEndpointPath(endpointPath), + ) + runtime.server = &http.Server{ + Handler: runtime.authorize(handler), + ReadHeaderTimeout: 5 * time.Second, + } + runtime.endpoint = api.CallerToolEndpoint{ + Name: serverName, + URL: "http://" + listener.Addr().String() + endpointPath, + Headers: map[string]string{ + "Authorization": "Bearer " + token, + }, + } + go func() { + _ = runtime.server.Serve(listener) + }() + return runtime, nil +} + +// Endpoint returns a copy so callers cannot mutate the runtime's credential. +func (r *Runtime) Endpoint() api.CallerToolEndpoint { + endpoint := r.endpoint + endpoint.Headers = cloneHeaders(r.endpoint.Headers) + return endpoint +} + +// CredentialHash returns the SHA-256 bearer hash persisted by an authority. +// The plaintext capability remains confined to Endpoint headers. +func (r *Runtime) CredentialHash() []byte { + hash := make([]byte, len(r.tokenHash)) + copy(hash, r.tokenHash[:]) + return hash +} + +// Close revokes the endpoint and is safe to call more than once. +func (r *Runtime) Close() error { + r.closeOnce.Do(func() { + r.Revoke() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + r.closeErr = r.server.Shutdown(ctx) + }) + return r.closeErr +} + +// Revoke invalidates the capability immediately and cancels active calls. +func (r *Runtime) Revoke() { + if r.revoked.CompareAndSwap(false, true) { + r.cancel() + } +} + +func (r *Runtime) filterTools(_ context.Context, tools []mcp.Tool) []mcp.Tool { + filtered := make([]mcp.Tool, 0, len(tools)) + for _, tool := range tools { + if _, ok := r.definitions[tool.Name]; ok { + filtered = append(filtered, tool) + } + } + return filtered +} + +func (r *Runtime) handler(definition api.ToolDefinition) server.ToolHandlerFunc { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + if _, ok := r.definitions[definition.Name]; !ok { + return nil, fmt.Errorf("caller tool %q is not authorized", definition.Name) + } + callCtx, cancel := context.WithCancel(ctx) + stop := context.AfterFunc(r.ctx, cancel) + defer stop() + defer cancel() + input := request.GetArguments() + if input == nil { + input = map[string]any{} + } + toolUseID, generatedToolUseID, err := toolUseID(request, input) + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + if definition.NeedsApproval() { + if r.canUseTool == nil { + return mcp.NewToolResultError("tool approval is required but no approval broker is configured"), nil + } + approvalCtx, approvalCancel := context.WithTimeout(callCtx, r.approvalTimeout) + decision, err := r.canUseTool(approvalCtx, api.PermissionRequest{ + Tool: definition.Name, Input: input, ToolUseID: toolUseID, + ToolUseIDGenerated: generatedToolUseID, SessionID: r.sessionID, + }) + approvalCancel() + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + if !decision.Allow { + message := decision.Message + if message == "" { + message = "tool call denied" + } + return mcp.NewToolResultError(message), nil + } + if decision.UpdatedInput != nil { + input = decision.UpdatedInput + } + } + if err := r.validateActive(callCtx); err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + if err := r.validateInput(definition.Name, input); err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + output, err := definition.Handler(callCtx, input) + if err != nil { + return mcp.NewToolResultError(err.Error()), nil + } + result, err := mcp.NewToolResultJSON(output) + if err != nil { + return mcp.NewToolResultErrorf("marshal caller tool %q result: %v", definition.Name, err), nil + } + return result, nil + } +} + +func (r *Runtime) authorize(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) { + host, _, err := net.SplitHostPort(request.RemoteAddr) + if err != nil || !net.ParseIP(host).IsLoopback() { + http.Error(w, "caller-tool endpoint requires loopback access", http.StatusForbidden) + return + } + if strings.TrimSpace(request.Header.Get("Origin")) != "" { + http.Error(w, "caller-tool endpoint does not accept browser origins", http.StatusForbidden) + return + } + actual := request.Header.Get("Authorization") + expected := "Bearer " + r.token + if subtle.ConstantTimeCompare([]byte(actual), []byte(expected)) != 1 || + r.validateActive(request.Context()) != nil { + w.Header().Set("WWW-Authenticate", "Bearer") + http.Error(w, "invalid caller-tool credential", http.StatusUnauthorized) + return + } + next.ServeHTTP(w, request) + }) +} + +func (r *Runtime) active() bool { + return !r.revoked.Load() && (r.expiresAt.IsZero() || time.Now().Before(r.expiresAt)) +} + +func (r *Runtime) validateActive(ctx context.Context) error { + if !r.active() { + return fmt.Errorf("caller-tool credential is inactive") + } + if r.validate != nil { + if err := r.validate(ctx); err != nil { + return fmt.Errorf("validate caller-tool credential: %w", err) + } + } + return nil +} + +func mcpTool(definition api.ToolDefinition) (mcp.Tool, *jsonschema.Schema, error) { + schema := definition.InputSchema + if schema == nil { + schema = map[string]any{"type": "object", "properties": map[string]any{}} + } + raw, err := json.Marshal(schema) + if err != nil { + return mcp.Tool{}, nil, fmt.Errorf("marshal caller tool %q schema: %w", definition.Name, err) + } + document, err := jsonschema.UnmarshalJSON(bytes.NewReader(raw)) + if err != nil { + return mcp.Tool{}, nil, fmt.Errorf("decode caller tool %q schema: %w", definition.Name, err) + } + compiler := jsonschema.NewCompiler() + resourceURL := "mem:///captain/caller-tools/" + definition.Name + "/input-schema.json" + if err := compiler.AddResource(resourceURL, document); err != nil { + return mcp.Tool{}, nil, fmt.Errorf("register caller tool %q schema: %w", definition.Name, err) + } + compiled, err := compiler.Compile(resourceURL) + if err != nil { + return mcp.Tool{}, nil, fmt.Errorf("compile caller tool %q schema: %w", definition.Name, err) + } + tool := mcp.NewToolWithRawSchema(definition.Name, definition.Description, raw) + tool.Annotations = mcp.ToolAnnotation{ + ReadOnlyHint: definition.ReadOnlyHint, DestructiveHint: definition.DestructiveHint, + IdempotentHint: definition.IdempotentHint, + } + return tool, compiled, nil +} + +func (r *Runtime) validateInput(toolName string, input map[string]any) error { + raw, err := json.Marshal(input) + if err != nil { + return fmt.Errorf("marshal caller tool %q input: %w", toolName, err) + } + document, err := jsonschema.UnmarshalJSON(bytes.NewReader(raw)) + if err != nil { + return fmt.Errorf("decode caller tool %q input: %w", toolName, err) + } + if err := r.schemas[toolName].Validate(document); err != nil { + return fmt.Errorf("caller tool %q input is invalid: %w", toolName, err) + } + return nil +} + +func capabilityToken(sessionID string) (string, error) { + secret, err := randomID(32) + if err != nil { + return "", fmt.Errorf("generate caller-tool credential: %w", err) + } + return "cap_" + capabilityIdentity(sessionID) + "." + secret, nil +} + +func toolUseID(request mcp.CallToolRequest, input map[string]any) (string, bool, error) { + inputID := "" + if value, exists := input[ToolUseIDInputKey]; exists { + delete(input, ToolUseIDInputKey) + var ok bool + inputID, ok = value.(string) + if !ok || strings.TrimSpace(inputID) == "" { + return "", false, fmt.Errorf("caller-tool provider ID must be a non-empty string") + } + } + metadataID := "" + if request.Params.Meta != nil { + if value, ok := request.Params.Meta.AdditionalFields["toolUseId"].(string); ok && strings.TrimSpace(value) != "" { + metadataID = value + } + } + if inputID != "" && metadataID != "" && inputID != metadataID { + return "", false, fmt.Errorf("caller-tool provider ID conflicts with MCP metadata") + } + if metadataID != "" { + return metadataID, false, nil + } + if inputID != "" { + return inputID, false, nil + } + id, err := randomID(16) + if err != nil { + return "", false, fmt.Errorf("generate caller-tool call ID: %w", err) + } + return "mcp_" + id, true, nil +} + +func randomID(size int) (string, error) { + value := make([]byte, size) + if _, err := rand.Read(value); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(value), nil +} + +func capabilityIdentity(sessionID string) string { + identity := strings.Map(func(value rune) rune { + switch { + case value >= 'a' && value <= 'z', value >= 'A' && value <= 'Z', value >= '0' && value <= '9', value == '-', value == '_': + return value + default: + return '_' + } + }, strings.TrimSpace(sessionID)) + if identity == "" { + return "run" + } + return identity +} + +func cloneHeaders(headers map[string]string) map[string]string { + if headers == nil { + return nil + } + cloned := make(map[string]string, len(headers)) + for key, value := range headers { + cloned[key] = value + } + return cloned +} diff --git a/pkg/ai/callertools/runtime_ginkgo_test.go b/pkg/ai/callertools/runtime_ginkgo_test.go new file mode 100644 index 00000000..aaa1c938 --- /dev/null +++ b/pkg/ai/callertools/runtime_ginkgo_test.go @@ -0,0 +1,352 @@ +package callertools_test + +import ( + "context" + "errors" + "net/http" + "sync/atomic" + "time" + + "github.com/flanksource/captain/pkg/ai/callertools" + "github.com/flanksource/captain/pkg/api" + mcpclient "github.com/mark3labs/mcp-go/client" + "github.com/mark3labs/mcp-go/client/transport" + "github.com/mark3labs/mcp-go/mcp" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Authenticated caller-tool runtime", func() { + It("omits denied tools and rejects unauthenticated requests", func(ctx SpecContext) { + var hiddenCalls atomic.Int32 + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{ + { + Name: "invoice_get", Description: "Read an invoice", + InputSchema: map[string]any{"type": "object", "properties": map[string]any{"id": map[string]any{"type": "string"}}}, + DefaultPermission: api.ToolModeOn, + Handler: func(_ context.Context, input map[string]any) (any, error) { + return map[string]any{"id": input["id"], "status": "draft"}, nil + }, + }, + { + Name: "invoice_delete", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { + hiddenCalls.Add(1) + return "deleted", nil + }, + }, + }, + Preferences: api.ToolPreferences{"invoice_delete": api.ToolModeOff}, + SessionID: "captain-session-1", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(runtime.Close) + + response, err := http.Post(runtime.Endpoint().URL, "application/json", nil) + Expect(err).NotTo(HaveOccurred()) + Expect(response.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(response.Body.Close()).To(Succeed()) + + client := authenticatedClient(ctx, runtime.Endpoint()) + DeferCleanup(client.Close) + tools, err := client.ListTools(ctx, mcp.ListToolsRequest{}) + Expect(err).NotTo(HaveOccurred()) + Expect(tools.Tools).To(HaveLen(1)) + Expect(tools.Tools[0].Name).To(Equal("invoice_get")) + + request := mcp.CallToolRequest{} + request.Params.Name = "invoice_delete" + request.Params.Arguments = map[string]any{} + _, err = client.CallTool(ctx, request) + Expect(err).To(HaveOccurred()) + Expect(hiddenCalls.Load()).To(BeZero()) + }) + + It("brokers ask tools and applies updated input", func(ctx SpecContext) { + var calls atomic.Int32 + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "invoice_update", DefaultPermission: api.ToolModeAsk, + Handler: func(_ context.Context, input map[string]any) (any, error) { + calls.Add(1) + return input, nil + }, + }}, + CanUseTool: func(_ context.Context, request api.PermissionRequest) (api.PermissionDecision, error) { + Expect(request.Tool).To(Equal("invoice_update")) + Expect(request.SessionID).To(Equal("captain-session-2")) + Expect(request.ToolUseID).To(Equal("approval-call-1")) + return api.PermissionDecision{Allow: true, UpdatedInput: map[string]any{"status": "approved"}}, nil + }, + SessionID: "captain-session-2", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(runtime.Close) + + client := authenticatedClient(ctx, runtime.Endpoint()) + DeferCleanup(client.Close) + request := mcp.CallToolRequest{} + request.Params.Name = "invoice_update" + request.Params.Arguments = map[string]any{"status": "draft"} + request.Params.Meta = &mcp.Meta{AdditionalFields: map[string]any{"toolUseId": "approval-call-1"}} + result, err := client.CallTool(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(result.IsError).To(BeFalse()) + Expect(result.StructuredContent).To(Equal(map[string]any{"status": "approved"})) + Expect(calls.Load()).To(Equal(int32(1))) + }) + + It("uses the provider tool-use ID without exposing transport input to the handler", func(ctx SpecContext) { + var handledInput map[string]any + permissionRequests := make(chan api.PermissionRequest, 1) + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "invoice_update", DefaultPermission: api.ToolModeAsk, + Handler: func(_ context.Context, input map[string]any) (any, error) { + handledInput = input + return input, nil + }, + }}, + CanUseTool: func(_ context.Context, request api.PermissionRequest) (api.PermissionDecision, error) { + permissionRequests <- request + return api.PermissionDecision{Allow: true}, nil + }, + SessionID: "provider-correlation-session", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(runtime.Close) + + client := authenticatedClient(ctx, runtime.Endpoint()) + DeferCleanup(client.Close) + request := mcp.CallToolRequest{} + request.Params.Name = "invoice_update" + request.Params.Arguments = map[string]any{ + "status": "draft", "__captain_tool_use_id": "claude-tool-use-1", + } + result, err := client.CallTool(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(result.IsError).To(BeFalse()) + var permission api.PermissionRequest + Eventually(permissionRequests).Should(Receive(&permission)) + Expect(permission.ToolUseID).To(Equal("claude-tool-use-1")) + Expect(permission.Input).To(Equal(map[string]any{"status": "draft"})) + Expect(handledInput).To(Equal(map[string]any{"status": "draft"})) + }) + + It("rejects wrong-session credentials and browser origins", func() { + first := newRuntime("captain-session-1", "first") + DeferCleanup(first.Close) + second := newRuntime("captain-session-2", "second") + DeferCleanup(second.Close) + + request, err := http.NewRequest(http.MethodPost, first.Endpoint().URL, nil) + Expect(err).NotTo(HaveOccurred()) + for name, value := range second.Endpoint().Headers { + request.Header.Set(name, value) + } + response, err := http.DefaultClient.Do(request) + Expect(err).NotTo(HaveOccurred()) + Expect(response.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(response.Body.Close()).To(Succeed()) + + request, err = http.NewRequest(http.MethodPost, first.Endpoint().URL, nil) + Expect(err).NotTo(HaveOccurred()) + for name, value := range first.Endpoint().Headers { + request.Header.Set(name, value) + } + request.Header.Set("Origin", "https://example.com") + response, err = http.DefaultClient.Do(request) + Expect(err).NotTo(HaveOccurred()) + Expect(response.StatusCode).To(Equal(http.StatusForbidden)) + Expect(response.Body.Close()).To(Succeed()) + }) + + It("expires and explicitly revokes capabilities", func() { + expiring, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "lookup", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return "ok", nil }, + }}, + SessionID: "expiring-session", + ExpiresAt: time.Now().Add(25 * time.Millisecond), + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(expiring.Close) + Eventually(func() int { + return authenticatedStatus(expiring.Endpoint()) + }).Should(Equal(http.StatusUnauthorized)) + + revoked := newRuntime("revoked-session", "revoked") + DeferCleanup(revoked.Close) + revoked.Revoke() + Expect(authenticatedStatus(revoked.Endpoint())).To(Equal(http.StatusUnauthorized)) + }) + + It("times out approvals and returns handler failures without executing past the boundary", func(ctx SpecContext) { + var calls atomic.Int32 + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "invoice_update", DefaultPermission: api.ToolModeAsk, + Handler: func(context.Context, map[string]any) (any, error) { + calls.Add(1) + return nil, errors.New("must not execute") + }, + }}, + CanUseTool: func(ctx context.Context, _ api.PermissionRequest) (api.PermissionDecision, error) { + <-ctx.Done() + return api.PermissionDecision{}, ctx.Err() + }, + SessionID: "approval-session", + ApprovalTimeout: 25 * time.Millisecond, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(runtime.Close) + + client := authenticatedClient(ctx, runtime.Endpoint()) + DeferCleanup(client.Close) + request := mcp.CallToolRequest{} + request.Params.Name = "invoice_update" + result, err := client.CallTool(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(result.IsError).To(BeTrue()) + Expect(calls.Load()).To(BeZero()) + }) + + It("returns handler failures as MCP tool errors", func(ctx SpecContext) { + var calls atomic.Int32 + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { + calls.Add(1) + return nil, errors.New("invoice unavailable") + }, + }}, + SessionID: "failure-session", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(runtime.Close) + + client := authenticatedClient(ctx, runtime.Endpoint()) + DeferCleanup(client.Close) + request := mcp.CallToolRequest{} + request.Params.Name = "invoice_get" + result, err := client.CallTool(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(result.IsError).To(BeTrue()) + Expect(calls.Load()).To(Equal(int32(1))) + }) + + It("rejects approval-updated input that violates the tool schema", func(ctx SpecContext) { + var calls atomic.Int32 + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "invoice_update", DefaultPermission: api.ToolModeAsk, + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "id": map[string]any{"type": "string"}, + }, + "required": []string{"id"}, + }, + Handler: func(context.Context, map[string]any) (any, error) { + calls.Add(1) + return "updated", nil + }, + }}, + CanUseTool: func(context.Context, api.PermissionRequest) (api.PermissionDecision, error) { + return api.PermissionDecision{Allow: true, UpdatedInput: map[string]any{"id": 42}}, nil + }, + SessionID: "validation-session", + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(runtime.Close) + + client := authenticatedClient(ctx, runtime.Endpoint()) + DeferCleanup(client.Close) + request := mcp.CallToolRequest{} + request.Params.Name = "invoice_update" + request.Params.Arguments = map[string]any{"id": "inv-1"} + result, err := client.CallTool(ctx, request) + Expect(err).NotTo(HaveOccurred()) + Expect(result.IsError).To(BeTrue()) + Expect(calls.Load()).To(BeZero()) + }) + + It("isolates concurrently active session capabilities", func(ctx SpecContext) { + first := newRuntime("captain-session-1", "first") + DeferCleanup(first.Close) + second := newRuntime("captain-session-2", "second") + DeferCleanup(second.Close) + firstClient := authenticatedClient(ctx, first.Endpoint()) + DeferCleanup(firstClient.Close) + secondClient := authenticatedClient(ctx, second.Endpoint()) + DeferCleanup(secondClient.Close) + + type outcome struct { + result *mcp.CallToolResult + err error + } + outcomes := make(chan outcome, 2) + call := func(client *mcpclient.Client) { + request := mcp.CallToolRequest{} + request.Params.Name = "identity" + result, err := client.CallTool(ctx, request) + outcomes <- outcome{result: result, err: err} + } + go call(firstClient) + go call(secondClient) + + values := make([]string, 0, 2) + for range 2 { + outcome := <-outcomes + Expect(outcome.err).NotTo(HaveOccurred()) + Expect(outcome.result.IsError).To(BeFalse()) + values = append(values, outcome.result.StructuredContent.(map[string]any)["session"].(string)) + } + Expect(values).To(ConsistOf("first", "second")) + }) +}) + +func authenticatedClient(ctx context.Context, endpoint api.CallerToolEndpoint) *mcpclient.Client { + channel, err := transport.NewStreamableHTTP(endpoint.URL, transport.WithHTTPHeaders(endpoint.Headers)) + Expect(err).NotTo(HaveOccurred()) + client := mcpclient.NewClient(channel) + Expect(client.Start(ctx)).To(Succeed()) + request := mcp.InitializeRequest{} + request.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION + request.Params.ClientInfo = mcp.Implementation{Name: "captain-test", Version: "1.0.0"} + _, err = client.Initialize(ctx, request) + Expect(err).NotTo(HaveOccurred()) + return client +} + +func authenticatedStatus(endpoint api.CallerToolEndpoint) int { + request, err := http.NewRequest(http.MethodPost, endpoint.URL, nil) + Expect(err).NotTo(HaveOccurred()) + for name, value := range endpoint.Headers { + request.Header.Set(name, value) + } + response, err := http.DefaultClient.Do(request) + Expect(err).NotTo(HaveOccurred()) + defer func() { + Expect(response.Body.Close()).To(Succeed()) + }() + return response.StatusCode +} + +func newRuntime(sessionID, marker string) *callertools.Runtime { + runtime, err := callertools.New(callertools.Options{ + Definitions: []api.ToolDefinition{{ + Name: "identity", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { + return map[string]any{"session": marker}, nil + }, + }}, + SessionID: sessionID, + }) + Expect(err).NotTo(HaveOccurred()) + return runtime +} diff --git a/pkg/ai/catalog.go b/pkg/ai/catalog.go index 618cf755..53a71c0d 100644 --- a/pkg/ai/catalog.go +++ b/pkg/ai/catalog.go @@ -79,12 +79,10 @@ var ( // the registry at package init — long before captainconfig.Load installs the // disabled set — so anything baked in at build time would never see it. func Catalog() []Model { - modelRegistryMu.RLock() - defer modelRegistryMu.RUnlock() - disabled := Disabled() - out := make([]Model, 0, len(catalog)) - for _, m := range catalog { + models := catalogSnapshot() + out := make([]Model, 0, len(models)) + for _, m := range models { if disabled.Model(m.Backend, bareProviderModelID(m.ID)) { continue } @@ -94,6 +92,15 @@ func Catalog() []Model { return out } +// catalogSnapshot returns every registered model before user opt-outs are +// applied. Execution paths use Catalog; descriptive menus use this snapshot so +// they can retain disabled choices and explain how to re-enable them. +func catalogSnapshot() []Model { + modelRegistryMu.RLock() + defer modelRegistryMu.RUnlock() + return append([]Model(nil), catalog...) +} + // RegisterModel adds a model to the global catalog, or replaces the existing // entry with the same ID while preserving its position in the menu. func RegisterModel(model Model) error { diff --git a/pkg/ai/catalog_disabled_ginkgo_test.go b/pkg/ai/catalog_disabled_ginkgo_test.go index 71ad3fec..3adc1344 100644 --- a/pkg/ai/catalog_disabled_ginkgo_test.go +++ b/pkg/ai/catalog_disabled_ginkgo_test.go @@ -1,6 +1,8 @@ package ai import ( + "time" + . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -106,6 +108,33 @@ var _ = Describe("catalog opt-out filtering", func() { } }) + It("retains disabled models with remediation in the live descriptive menu", func() { + previousProbe := adapterProbe + previousCache, previousAt := adapterCache, adapterCacheAt + DeferCleanup(func() { + adapterProbe = previousProbe + adapterCache, adapterCacheAt = previousCache, previousAt + }) + adapterCache, adapterCacheAt = nil, time.Time{} + adapterProbe = func() ([]AdapterStatus, error) { return nil, nil } + disable(nil, []string{"deepseek"}, nil, nil, nil) + + infos, err := LiveCatalogInfo(nil) + Expect(err).NotTo(HaveOccurred()) + var deepseek []ModelInfo + for _, info := range infos { + if info.Provider == "deepseek" { + deepseek = append(deepseek, info) + } + } + Expect(deepseek).NotTo(BeEmpty()) + for _, info := range deepseek { + Expect(info.Configured).To(BeFalse()) + Expect(info.Availability.State).To(Equal(api.AvailabilityDisabled)) + Expect(info.Availability.Remediation).NotTo(BeEmpty()) + } + }) + // The menu names the default so no client hardcodes a model id. Disabling // that model has to leave the menu with none marked rather than pointing a // picker at a row the user switched off. @@ -125,4 +154,21 @@ var _ = Describe("catalog opt-out filtering", func() { Expect(defaults()).To(BeEmpty()) }) + + It("serves the exact Captain runtime selection for every model row", func() { + models := Catalog() + info := CatalogInfo(nil) + + Expect(info).To(HaveLen(len(models))) + for index, model := range models { + want := api.Model{ + Name: model.BareID(), + Backend: model.Backend, + }.Capabilities() + if model.ID != want.Name { + want.ID = model.ID + } + Expect(info[index].Runtime).To(Equal(want)) + } + }) }) diff --git a/pkg/ai/catalog_info.go b/pkg/ai/catalog_info.go index 84da1217..3f80bb82 100644 --- a/pkg/ai/catalog_info.go +++ b/pkg/ai/catalog_info.go @@ -4,6 +4,7 @@ import ( "os/exec" "slices" + "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/api/registry" ) @@ -11,12 +12,14 @@ import ( // 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"` - Temperature bool `json:"temperature"` - Configured bool `json:"configured"` + ID string `json:"id"` + Provider string `json:"provider"` + Label string `json:"label"` + Runtime api.Model `json:"runtime"` + Reasoning bool `json:"reasoning"` + Temperature bool `json:"temperature"` + Configured bool `json:"configured"` + Availability api.Availability `json:"availability"` // Default marks captain's declared default model, so a client seeds its // picker from the menu instead of hardcoding an id that rots on the next // release. At most one row carries it, and none does when that model is @@ -50,31 +53,40 @@ func BackendToProvider(b Backend) string { // backend binary is installed. Order mirrors the catalog so the client renders a // stable, grouped menu. func CatalogInfo(configuredProviders []string) []ModelInfo { - return catalogInfoFrom(Catalog(), configuredProviders) + return catalogInfoFrom(Catalog(), catalogInfoOptions{ConfiguredProviders: configuredProviders}) +} + +type catalogInfoOptions struct { + ConfiguredProviders []string + Adapters []AdapterStatus } // catalogInfoFrom annotates an arbitrary model list with selectability, shared // by CatalogInfo (static catalog) and LiveCatalogInfo (whoami-probed catalog). // -// Both inputs arrive already filtered — Catalog() drops disabled models at read -// time and mergeLiveCatalog drops disabled probe rows — so this only annotates. -func catalogInfoFrom(models []Model, configuredProviders []string) []ModelInfo { +// Static execution menus arrive filtered; live descriptive menus may retain +// disabled rows so this seam also supplies their reason and remediation. +func catalogInfoFrom(models []Model, options catalogInfoOptions) []ModelInfo { out := make([]ModelInfo, 0, len(models)) for _, m := range models { - configured := false - if m.IsAgent() { - configured = agentBackendAvailable(m.Backend) - } else { - configured = slices.Contains(configuredProviders, BackendToProvider(m.Backend)) + availability := modelAvailability(m, options) + runtime := api.Model{ + Name: m.BareID(), + Backend: m.Backend, + }.Capabilities() + if m.ID != runtime.Name { + runtime.ID = m.ID } out = append(out, ModelInfo{ ID: m.ID, Provider: BackendToProvider(m.Backend), Label: m.Label, + Runtime: runtime, Reasoning: m.Reasoning, Temperature: m.Temperature, - Configured: configured, - Default: m.ID == registry.DefaultModelID, + Configured: availability.IsAvailable(), + Availability: availability, + Default: m.ID == registry.DefaultModelID && availability.State != api.AvailabilityDisabled, ContextWindow: m.ContextWindow, InputMediaTypes: append([]string(nil), m.InputMediaTypes...), }) @@ -82,6 +94,49 @@ func catalogInfoFrom(models []Model, configuredProviders []string) []ModelInfo { return out } +func modelAvailability(model Model, options catalogInfoOptions) api.Availability { + disabled := Disabled() + if disabled.Model(model.Backend, model.BareID()) { + reason := disabled.Reason(model.Backend) + if reason == "" { + reason = "model " + model.ID + } + return api.Availability{State: api.AvailabilityDisabled, + Reason: "Disabled by " + reason + " in Captain configuration.", + Remediation: "Enable " + reason + " on the Whoami page, then refresh."} + } + if slices.Contains(options.ConfiguredProviders, BackendToProvider(model.Backend)) { + return api.Available() + } + for _, adapter := range options.Adapters { + if adapter.Backend == string(model.Backend) { + return AvailabilityForAdapter(adapter) + } + } + if model.IsAgent() && agentBackendAvailable(model.Backend) { + return api.Available() + } + if model.IsAgent() { + return AvailabilityForAdapter(AdapterStatus{Backend: string(model.Backend), Type: "cli", BinaryMissing: requiredBinary(model.Backend)}) + } + return AvailabilityForAdapter(AdapterStatus{Backend: string(model.Backend), Type: "api"}) +} + +func requiredBinary(backend Backend) string { + switch backend { + case BackendCodexCLI, BackendCodexAgent, BackendCodexCmux: + return "codex" + case BackendClaudeCLI, BackendClaudeCmux: + return "claude" + case BackendClaudeAgent: + return "tsx" + case BackendGeminiCLI: + return "gemini" + default: + return string(backend) + } +} + // agentBackendAvailable reports whether an agent backend's local binary is // 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 diff --git a/pkg/ai/catalog_resolve.go b/pkg/ai/catalog_resolve.go index 549fc07c..8b0cb7a1 100644 --- a/pkg/ai/catalog_resolve.go +++ b/pkg/ai/catalog_resolve.go @@ -102,7 +102,9 @@ func cachedRows(opts ResolveOptions, fp string) ([]ResolvedModel, bool) { } // resolveFresh seeds the catalog filtered by backend, unions live API models -// when tokens are present, and joins each row to pricing. +// when tokens are present, and joins each row to pricing. opts.Refresh reaches +// the pricing snapshot too: bypassing the model cache while still pricing from a +// day-old OpenRouter snapshot would only half-honour --no-cache. func resolveFresh(ctx context.Context, opts ResolveOptions) ([]ResolvedModel, error) { rows, index := seedCatalog(opts.Backend) @@ -112,7 +114,7 @@ func resolveFresh(ctx context.Context, opts ResolveOptions) ([]ResolvedModel, er } } - pricing.EnsureLoaded() + pricing.EnsureLoaded(pricing.LoadOptions{Refresh: opts.Refresh}) for i := range rows { if info, ok := lookupPricing(rows[i].Backend, rows[i].BareID()); ok { rows[i].Price = info diff --git a/pkg/ai/client.go b/pkg/ai/client.go index dee1499a..1fb579d4 100644 --- a/pkg/ai/client.go +++ b/pkg/ai/client.go @@ -88,7 +88,7 @@ func suggestModelName(err error, model string) error { candidates = append(candidates, id) } } - pricing.EnsureLoaded() + pricing.EnsureLoaded(pricing.LoadOptions{}) for _, mi := range pricing.ListModels("") { candidates = append(candidates, mi.ModelID) } diff --git a/pkg/ai/fixture/stream.go b/pkg/ai/fixture/stream.go index 166bee0e..b32eab1f 100644 --- a/pkg/ai/fixture/stream.go +++ b/pkg/ai/fixture/stream.go @@ -9,6 +9,8 @@ import ( "encoding/json" "regexp" "strings" + + "github.com/flanksource/captain/pkg/api" ) type Summary struct { @@ -32,22 +34,50 @@ type Summary struct { MCPAPILog []MCPAPIEntry ToolCallLog []ToolCallEntry ToolCounts map[string]int + + // usageFromResult records that the token counts came from the stream's + // result line — the provider's own total — rather than being rebuilt from + // per-message usage. + usageFromResult bool + // responses deduplicates the per-content-block lines one response spans, + // for streams that report no result usage. See api.ResponseSet. + responses api.ResponseSet +} + +// UsageFromResult reports whether the token counts are the provider's reported +// total rather than a per-message reconstruction. +func (s *Summary) UsageFromResult() bool { return s.usageFromResult } + +func firstNonZeroCost(values ...float64) float64 { + for _, value := range values { + if value != 0 { + return value + } + } + return 0 } type streamEvent struct { - Type string `json:"type,omitempty"` - Subtype string `json:"subtype,omitempty"` - SessionID string `json:"session_id,omitempty"` - Message json.RawMessage `json:"message,omitempty"` - Result string `json:"result,omitempty"` - Error string `json:"error,omitempty"` - IsError bool `json:"is_error,omitempty"` - CostUSD float64 `json:"cost_usd,omitempty"` - DurationMS float64 `json:"duration_ms,omitempty"` - Usage *streamUsage `json:"usage,omitempty"` + Type string `json:"type,omitempty"` + Subtype string `json:"subtype,omitempty"` + SessionID string `json:"session_id,omitempty"` + Message json.RawMessage `json:"message,omitempty"` + Result string `json:"result,omitempty"` + Error string `json:"error,omitempty"` + IsError bool `json:"is_error,omitempty"` + // Two real producers spell the result's cost differently: the claude CLI's + // stream-json result line reports total_cost_usd, while the claude-agent + // SDK's turn-done notification reports cost_usd. Accept both. + CostUSD float64 `json:"cost_usd,omitempty"` + TotalCostUSD float64 `json:"total_cost_usd,omitempty"` + DurationMS float64 `json:"duration_ms,omitempty"` + Usage *streamUsage `json:"usage,omitempty"` } type streamMessage struct { + // ID identifies the API response; a response spanning several content + // blocks is written as several lines that all repeat the same usage. + ID string `json:"id,omitempty"` Content []streamContent `json:"content,omitempty"` Usage *streamUsage `json:"usage,omitempty"` } @@ -107,6 +137,8 @@ type Event struct { Content []streamContent // Raw usage reported on the message (if any). MessageUsage *streamUsage + // MessageID identifies the API response the message belongs to. + MessageID string // For result events. Result string @@ -133,7 +165,7 @@ func ParseLine(line []byte) (Event, bool) { Result: ev.Result, Error: ev.Error, DurationMS: ev.DurationMS, - CostUSD: ev.CostUSD, + CostUSD: firstNonZeroCost(ev.TotalCostUSD, ev.CostUSD), Usage: ev.Usage, } if len(ev.Message) > 0 { @@ -141,6 +173,7 @@ func ParseLine(line []byte) (Event, bool) { if err := json.Unmarshal(ev.Message, &msg); err == nil { out.Content = msg.Content out.MessageUsage = msg.Usage + out.MessageID = msg.ID } } return out, true @@ -153,7 +186,7 @@ func (s *Summary) Apply(ev Event) { if ev.SessionID != "" { s.SessionID = ev.SessionID } - if ev.MessageUsage != nil { + if ev.MessageUsage != nil && !s.usageFromResult && s.responses.First(ev.MessageID) { s.Input += ev.MessageUsage.InputTokens s.Output += ev.MessageUsage.OutputTokens s.CacheRead += ev.MessageUsage.CacheReadInputTokens @@ -175,10 +208,20 @@ func (s *Summary) Apply(ev Event) { } } } - // Note: deliberately do not consume ev.Usage on the top-level result event. - // That field's semantics (cumulative vs final-turn) are unclear and used to - // overwrite the per-message accumulation, dropping earlier turns' tokens. - // Per-message usage from streamMessage.Usage above is the source of truth. + // The result line's usage is the whole invocation's total, not the final + // turn's — pkg/ai/provider/claude_cli.go reads it that way, and its test + // pins a verbatim CLI result line proving it. So it replaces the per-message + // accumulation rather than adding to it: a reported total is exact, where a + // total rebuilt from per-message lines has to be deduplicated and can still + // drift. The accumulation above remains the fallback for streams that report + // no result usage. + if ev.Usage != nil { + s.Input = ev.Usage.InputTokens + s.Output = ev.Usage.OutputTokens + s.CacheRead = ev.Usage.CacheReadInputTokens + s.CacheWrite = ev.Usage.CacheCreationInputTokens + s.usageFromResult = true + } if ev.Type == "result" || ev.Result != "" || ev.Error != "" { if ev.Result != "" { s.Result = ev.Result diff --git a/pkg/ai/fixture/stream_result_test.go b/pkg/ai/fixture/stream_result_test.go new file mode 100644 index 00000000..da6bfc02 --- /dev/null +++ b/pkg/ai/fixture/stream_result_test.go @@ -0,0 +1,88 @@ +package fixture + +import "testing" + +// resultLine is a verbatim `type: result` line from claude 2.1.220 — the same +// fixture pkg/ai/provider/claude_cli.go pins. The CLI reports the invocation's +// cost as total_cost_usd and its token totals at the top level, neither of which +// this parser used to read: cost came back 0 and tokens were rebuilt by summing +// per-message lines. +const resultLine = `{"is_error":false,"duration_api_ms":5,"num_turns":1,"stop_reason":"end_turn",` + + `"session_id":"a32c3053-b70a-4f4f-9e0e-9d1d4fb48e8e","total_cost_usd":0.000162,` + + `"usage":{"input_tokens":14,"cache_creation_input_tokens":2,"cache_read_input_tokens":3,"output_tokens":8},` + + `"terminal_reason":"completed","subtype":"success","result":"The capital of France is Paris.","type":"result"}` + +func applyLines(t *testing.T, lines ...string) Summary { + t.Helper() + var summary Summary + for _, line := range lines { + ev, ok := ParseLine([]byte(line)) + if !ok { + t.Fatalf("ParseLine(%q) returned not-ok", line) + } + summary.Apply(ev) + } + return summary +} + +func TestSummaryReadsCostFromTheRealCLIResultKey(t *testing.T) { + summary := applyLines(t, resultLine) + + if summary.CostUSD != 0.000162 { + t.Errorf("CostUSD = %v, want 0.000162 from total_cost_usd", summary.CostUSD) + } +} + +func TestSummaryStillReadsTheAgentSDKCostKey(t *testing.T) { + // The claude-agent SDK spells the same figure cost_usd on its turn-done + // notification, so both producers must parse. + summary := applyLines(t, `{"type":"result","subtype":"success","cost_usd":0.25,"result":"ok"}`) + + if summary.CostUSD != 0.25 { + t.Errorf("CostUSD = %v, want 0.25 from cost_usd", summary.CostUSD) + } +} + +func TestSummaryTakesTokensFromTheResultTotal(t *testing.T) { + summary := applyLines(t, resultLine) + + if summary.Input != 14 || summary.Output != 8 || summary.CacheRead != 3 || summary.CacheWrite != 2 { + t.Errorf("usage = in:%d out:%d cacheRead:%d cacheWrite:%d, want 14/8/3/2", + summary.Input, summary.Output, summary.CacheRead, summary.CacheWrite) + } + if !summary.UsageFromResult() { + t.Error("UsageFromResult() = false, want true when the result reports usage") + } +} + +// The result's usage is the invocation total, so it replaces the per-message +// accumulation instead of adding to it — otherwise every turn is counted twice. +func TestSummaryResultTotalReplacesPerMessageAccumulation(t *testing.T) { + assistant := `{"type":"assistant","message":{"id":"msg_a","usage":` + + `{"input_tokens":14,"output_tokens":8,"cache_read_input_tokens":3,"cache_creation_input_tokens":2}}}` + + summary := applyLines(t, assistant, resultLine) + + if summary.Input != 14 || summary.Output != 8 { + t.Errorf("usage = in:%d out:%d, want the reported total 14/8, not 28/16", + summary.Input, summary.Output) + } +} + +// Without a result total the parser falls back to per-message usage, which must +// count a response once however many content-block lines it spans. +func TestSummaryFallsBackToDedupedPerMessageUsage(t *testing.T) { + block := `{"type":"assistant","message":{"id":"msg_a","usage":` + + `{"input_tokens":10,"output_tokens":4}}}` + other := `{"type":"assistant","message":{"id":"msg_b","usage":` + + `{"input_tokens":10,"output_tokens":4}}}` + + summary := applyLines(t, block, block, block, other) + + if summary.Input != 20 || summary.Output != 8 { + t.Errorf("usage = in:%d out:%d, want 20/8 for two responses", summary.Input, summary.Output) + } + if summary.UsageFromResult() { + t.Error("UsageFromResult() = true, want false when no result usage was reported") + } +} diff --git a/pkg/ai/history/codex_events.go b/pkg/ai/history/codex_events.go index 9237a514..221a6aad 100644 --- a/pkg/ai/history/codex_events.go +++ b/pkg/ai/history/codex_events.go @@ -75,11 +75,32 @@ func buildCodexEventUse(event CodexEvent, cwd, sessionID string) ToolUse { ReasoningTokens: usage.ReasoningOutputTokens, CacheReadTokens: usage.CachedInputTokens, TotalTokens: usage.TotalTokens, + CumulativeUsage: cumulativeEventUsage(event), ContextWindow: eventContextWindow(event), RecordType: "event_msg." + event.Payload.Type, } } +// cumulativeEventUsage nets the session-to-date totals codex reports alongside +// each event's own delta. Same netting as the per-event fields so the two are +// directly comparable: codex reports input inclusive of the cached prefix and +// output inclusive of reasoning. +func cumulativeEventUsage(event CodexEvent) *api.Usage { + if event.Payload.Info == nil { + return nil + } + total := event.Payload.Info.TotalTokenUsage + if total == (CodexTokenUsage{}) { + return nil + } + return &api.Usage{ + InputTokens: codexNonCachedInputTokens(total), + OutputTokens: api.NetOutputTokens(total.OutputTokens, total.ReasoningOutputTokens), + ReasoningTokens: total.ReasoningOutputTokens, + CacheReadTokens: total.CachedInputTokens, + } +} + func addCodexEventValue(input map[string]any, key string, value any) { switch typed := value.(type) { case string: diff --git a/pkg/ai/history/codex_normalize.go b/pkg/ai/history/codex_normalize.go index 70a34b06..19cec527 100644 --- a/pkg/ai/history/codex_normalize.go +++ b/pkg/ai/history/codex_normalize.go @@ -6,6 +6,7 @@ import ( "strings" "time" + "github.com/flanksource/captain/pkg/bash" "github.com/flanksource/captain/pkg/claude/tools" "github.com/segmentio/encoding/json" ) @@ -130,6 +131,9 @@ func normalizeCodexCall(call CodexToolCall, input map[string]any) ToolUse { if tool == "" { tool = "Bash" } + if tool == "Bash" { + input = bash.TransformBashInput(input) + } use := ToolUse{ Tool: tool, diff --git a/pkg/ai/history/codex_normalize_ginkgo_test.go b/pkg/ai/history/codex_normalize_ginkgo_test.go index 197759c9..235f055d 100644 --- a/pkg/ai/history/codex_normalize_ginkgo_test.go +++ b/pkg/ai/history/codex_normalize_ginkgo_test.go @@ -16,6 +16,20 @@ func TestCodexNormalization(t *testing.T) { } var _ = Describe("NormalizeCodexToolCall", func() { + It("transforms shell wrappers before returning canonical history input", func() { + use := NormalizeCodexToolCall(CodexToolCall{ + Command: `/bin/zsh -lc 'gavel pr status 50 --logs'`, + ID: "call-shell", + }) + + Expect(use.Tool).To(Equal("Bash")) + Expect(use.Input).To(Equal(map[string]any{ + "command": "gavel pr status 50 --logs", + "shell": "zsh", + "shellFlags": []string{"-l"}, + })) + }) + It("normalizes command calls from every Codex transport to Bash", func() { uses := []ToolUse{ NormalizeCodexToolCall(CodexToolCall{ diff --git a/pkg/ai/history/footprint.go b/pkg/ai/history/footprint.go new file mode 100644 index 00000000..43ce2e79 --- /dev/null +++ b/pkg/ai/history/footprint.go @@ -0,0 +1,123 @@ +package history + +import ( + "github.com/flanksource/captain/pkg/bash" + "github.com/flanksource/captain/pkg/claude/tools" +) + +// Footprint is the file set one tool use touched. Paths are returned exactly as +// the tool expressed them — relative to that tool's own cwd, or absolute — so +// each caller anchors them against the base it needs: `captain changes` resolves +// them absolutely, the session builders relativise against the project root. +type Footprint struct { + Read []string + Written []string +} + +// ToolFootprint is the single definition of which files a tool use read and +// wrote. It exists because the answer used to be computed four different ways — +// once for `captain changes`, once per session builder, and once more in a dead +// row projection — so the same session reported different changed files +// depending on which surface asked, and Claude and Codex sessions disagreed on +// the same stored field. +// +// A write always outranks a read: a command that reads a file and then rewrites +// it is a modification, and listing it under both is how the same path ended up +// on both sides of the same session. +func ToolFootprint(tu ToolUse) Footprint { + var footprint Footprint + switch tu.Tool { + case "Read": + footprint.Read = stringInput(tu, "file_path") + case "Grep", "Glob": + footprint.Read = stringInput(tu, "path") + case "ApplyPatch": + // A parsed patch keeps its operations on the row, so the paths come from + // there rather than from re-scanning the payload. A rename writes both + // ends: git reports both as dirty, so a commit has to attribute both. + for _, file := range tools.ApplyPatchFiles(tu.Input) { + footprint.Written = append(footprint.Written, file.Path, file.MoveTo) + } + case "CodexExecScript": + footprint.Written = patchPaths(tu, "script") + case "Bash": + footprint = bashFootprint(tu) + } + if key, ok := fileMutatingTools[tu.Tool]; ok { + footprint.Written = append(footprint.Written, stringInput(tu, key)...) + } + // Patch payloads are checked for every tool that can carry one, including the + // shapes handled above: an ApplyPatch row that was never parsed into + // operations still has its payload, and an agent can pipe a patch through the + // shell. + if _, ok := applyPatchInputs[tu.Tool]; ok { + footprint.Written = append(footprint.Written, patchPaths(tu, applyPatchInputs[tu.Tool])...) + } + footprint.Written = cleanPaths(footprint.Written, nil) + footprint.Read = cleanPaths(footprint.Read, footprint.Written) + return footprint +} + +// bashFootprint covers the writes only a shell command can express — redirects, +// sed -i, mv, rm — which no tool-input lookup can see. +func bashFootprint(tu ToolUse) Footprint { + command, _ := tu.Input["command"].(string) + if command == "" { + return Footprint{} + } + result, err := bash.Analyze(command) + if err != nil || result == nil { + return Footprint{} + } + var footprint Footprint + for _, operation := range result.Operations { + footprint.Written = append(footprint.Written, operation.Path) + } + footprint.Read = append(footprint.Read, result.ReferencedPaths...) + return footprint +} + +func patchPaths(tu ToolUse, key string) []string { + payload, _ := tu.Input[key].(string) + if payload == "" { + return nil + } + return tools.ExtractApplyPatchPaths(payload) +} + +func stringInput(tu ToolUse, key string) []string { + value, _ := tu.Input[key].(string) + if value == "" { + return nil + } + return []string{value} +} + +// cleanPaths drops empties, duplicates, and paths present in exclude, keeping +// first-seen order. /dev/null is dropped outright: a patch that creates or +// deletes a file names it as the empty side of the hunk, and it is not a file +// any surface should report as changed. +func cleanPaths(paths []string, exclude []string) []string { + if len(paths) == 0 { + return nil + } + seen := make(map[string]struct{}, len(paths)+len(exclude)) + for _, path := range exclude { + seen[path] = struct{}{} + } + cleaned := make([]string, 0, len(paths)) + for _, path := range paths { + if path == "" || path == "/dev/null" { + continue + } + if _, dup := seen[path]; dup { + continue + } + seen[path] = struct{}{} + cleaned = append(cleaned, path) + } + if len(cleaned) == 0 { + return nil + } + return cleaned +} diff --git a/pkg/ai/history/footprint_ginkgo_test.go b/pkg/ai/history/footprint_ginkgo_test.go new file mode 100644 index 00000000..4691978d --- /dev/null +++ b/pkg/ai/history/footprint_ginkgo_test.go @@ -0,0 +1,89 @@ +package history + +import ( + "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = ginkgo.Describe("ToolFootprint", func() { + footprint := func(tool string, input map[string]any) Footprint { + return ToolFootprint(ToolUse{Tool: tool, Input: input}) + } + + ginkgo.DescribeTable("attributes a tool's writes", + func(tool string, input map[string]any, expected []string) { + Expect(footprint(tool, input).Written).To(Equal(expected)) + }, + ginkgo.Entry("Edit names its file", + "Edit", map[string]any{"file_path": "pkg/a.go"}, []string{"pkg/a.go"}), + ginkgo.Entry("Write names its file", + "Write", map[string]any{"file_path": "pkg/a.go"}, []string{"pkg/a.go"}), + ginkgo.Entry("MultiEdit names one file for many edits", + "MultiEdit", map[string]any{"file_path": "pkg/a.go"}, []string{"pkg/a.go"}), + ginkgo.Entry("NotebookEdit names notebook_path, which file_path lookups miss", + "NotebookEdit", map[string]any{"notebook_path": "nb.ipynb"}, []string{"nb.ipynb"}), + ginkgo.Entry("a shell redirect is a write no input lookup can see", + "Bash", map[string]any{"command": "echo hi > out.txt"}, []string{"out.txt"}), + ginkgo.Entry("sed -i rewrites in place", + "Bash", map[string]any{"command": "sed -i '' s/a/b/ pkg/a.go"}, []string{"pkg/a.go"}), + ginkgo.Entry("a patch piped through the shell touches every file in it", + "Bash", map[string]any{"command": "apply_patch <<'EOF'\n*** Begin Patch\n*** Update File: pkg/a.go\n*** Update File: pkg/b.go\n*** End Patch\nEOF"}, + []string{"pkg/a.go", "pkg/b.go"}), + ginkgo.Entry("codex exec carries the patch in input", + "exec", map[string]any{"input": "*** Begin Patch\n*** Update File: pkg/a.go\n*** End Patch"}, + []string{"pkg/a.go"}), + ginkgo.Entry("a codex script carries the patch in script", + "CodexExecScript", map[string]any{"script": "*** Begin Patch\n*** Update File: pkg/a.go\n*** End Patch"}, + []string{"pkg/a.go"}), + ) + + ginkgo.DescribeTable("attributes a tool's reads", + func(tool string, input map[string]any, expected []string) { + Expect(footprint(tool, input).Read).To(Equal(expected)) + }, + ginkgo.Entry("Read names its file", "Read", map[string]any{"file_path": "pkg/a.go"}, []string{"pkg/a.go"}), + ginkgo.Entry("Grep names its search root", "Grep", map[string]any{"path": "pkg"}, []string{"pkg"}), + ginkgo.Entry("Glob names its search root", "Glob", map[string]any{"path": "pkg"}, []string{"pkg"}), + ) + + ginkgo.It("drops /dev/null, which a patch names as the empty side of a hunk", func() { + written := footprint("exec", map[string]any{ + "input": "*** Begin Patch\n*** Delete File: pkg/gone.go\n*** End Patch", + }).Written + + Expect(written).NotTo(ContainElement("/dev/null")) + }) + + ginkgo.It("counts a renamed file's destination as written", func() { + Expect(footprint("exec", map[string]any{ + "input": "*** Begin Patch\n*** Update File: pkg/old.go\n*** Move to: pkg/new.go\n*** End Patch", + }).Written).To(ContainElements("pkg/old.go", "pkg/new.go")) + }) + + ginkgo.It("classifies a file that is read and then rewritten as a write only", func() { + result := footprint("Bash", map[string]any{"command": "sed -i '' s/a/b/ pkg/a.go"}) + + Expect(result.Written).To(ContainElement("pkg/a.go")) + Expect(result.Read).NotTo(ContainElement("pkg/a.go")) + }) + + ginkgo.It("reports the same path once however many times a tool names it", func() { + Expect(footprint("Bash", map[string]any{ + "command": "echo a > out.txt; echo b >> out.txt", + }).Written).To(Equal([]string{"out.txt"})) + }) + + ginkgo.DescribeTable("touches nothing", + func(tool string, input map[string]any) { + result := footprint(tool, input) + + Expect(result.Written).To(BeNil()) + Expect(result.Read).To(BeNil()) + }, + ginkgo.Entry("a tool with no file surface", "WebSearch", map[string]any{"query": "x"}), + ginkgo.Entry("an empty input", "Edit", map[string]any{}), + ginkgo.Entry("a non-string path", "Edit", map[string]any{"file_path": 42}), + ginkgo.Entry("an empty bash command", "Bash", map[string]any{"command": ""}), + ginkgo.Entry("a nil input", "Edit", nil), + ) +}) diff --git a/pkg/ai/history/modified_files.go b/pkg/ai/history/modified_files.go index 5d985c17..6f06c6a5 100644 --- a/pkg/ai/history/modified_files.go +++ b/pkg/ai/history/modified_files.go @@ -16,27 +16,33 @@ var fileMutatingTools = map[string]string{ "NotebookEdit": "notebook_path", } +// applyPatchInputs are the tools that carry a whole patch rather than naming a +// single file. Both the raw codex names and the normalized one are listed +// because this table is consulted from either side of normalization; a `Bash` +// command is included because an agent can pipe a patch through the shell. +// +// A patch is the only tool shape that can touch several files at once, and the +// only one that can rename or delete, so it is parsed rather than looked up. +var applyPatchInputs = map[string]string{ + "ApplyPatch": "input", + "apply_patch": "input", + "exec": "input", + "Bash": "command", +} + // ModifiedFiles returns the distinct files an agent wrote to across the given -// tool uses, in first-seen order. Only Edit/Write/MultiEdit/NotebookEdit count; -// the path is read from the tool's input key (file_path / notebook_path). Empty -// or non-string paths are skipped. +// tool uses, in first-seen order. Empty or non-string paths are skipped. func ModifiedFiles(toolUses []ToolUse) []string { var files []string seen := make(map[string]struct{}, len(toolUses)) for _, tu := range toolUses { - key, ok := fileMutatingTools[tu.Tool] - if !ok { - continue - } - path, _ := tu.Input[key].(string) - if path == "" { - continue - } - if _, dup := seen[path]; dup { - continue + for _, path := range ToolFootprint(tu).Written { + if _, dup := seen[path]; dup { + continue + } + seen[path] = struct{}{} + files = append(files, path) } - seen[path] = struct{}{} - files = append(files, path) } return files } diff --git a/pkg/ai/history/modified_files_test.go b/pkg/ai/history/modified_files_test.go index e0d03af5..f0e5472f 100644 --- a/pkg/ai/history/modified_files_test.go +++ b/pkg/ai/history/modified_files_test.go @@ -27,6 +27,59 @@ func TestModifiedFiles(t *testing.T) { } } +// TestModifiedFilesFromPatches: codex expresses every edit as a patch, and +// normalization only rewrites the single-file add/update shapes into Write/Edit +// rows. Multi-file patches, deletes and renames keep the ApplyPatch name, and +// dropping them leaves a commit unable to attribute work the agent plainly did +// — a deleted or renamed file is dirty in git exactly like an edited one. +func TestModifiedFilesFromPatches(t *testing.T) { + multi := "*** Begin Patch\n" + + "*** Update File: pkg/a.go\n@@\n-old\n+new\n" + + "*** Add File: pkg/b.go\n+package b\n" + + "*** End Patch\n" + deletion := "*** Begin Patch\n*** Delete File: pkg/gone.go\n*** End Patch\n" + rename := "*** Begin Patch\n*** Update File: pkg/from.go\n*** Move to: pkg/to.go\n@@\n-old\n+new\n*** End Patch\n" + + cases := []struct { + name string + use ToolUse + want []string + }{ + { + name: "normalized multi-file patch", + use: ToolUse{Tool: "ApplyPatch", Input: map[string]any{"input": multi}}, + want: []string{"pkg/a.go", "pkg/b.go"}, + }, + { + name: "raw codex tool name", + use: ToolUse{Tool: "apply_patch", Input: map[string]any{"input": multi}}, + want: []string{"pkg/a.go", "pkg/b.go"}, + }, + { + name: "deletion", + use: ToolUse{Tool: "ApplyPatch", Input: map[string]any{"input": deletion}}, + want: []string{"pkg/gone.go"}, + }, + { + name: "rename reports both ends", + use: ToolUse{Tool: "ApplyPatch", Input: map[string]any{"input": rename}}, + want: []string{"pkg/from.go", "pkg/to.go"}, + }, + { + name: "patch piped through the shell", + use: ToolUse{Tool: "Bash", Input: map[string]any{"command": "apply_patch <<'EOF'\n" + deletion + "EOF"}}, + want: []string{"pkg/gone.go"}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := ModifiedFiles([]ToolUse{tc.use}); !reflect.DeepEqual(got, tc.want) { + t.Errorf("ModifiedFiles = %v, want %v", got, tc.want) + } + }) + } +} + func TestSessionModifiedFiles(t *testing.T) { dir := t.TempDir() session := filepath.Join(dir, "session.jsonl") diff --git a/pkg/ai/history/types.go b/pkg/ai/history/types.go index d4a90da7..4ac8c365 100644 --- a/pkg/ai/history/types.go +++ b/pkg/ai/history/types.go @@ -1,6 +1,10 @@ package history -import "time" +import ( + "time" + + "github.com/flanksource/captain/pkg/api" +) type ToolUse struct { Tool string `json:"tool,omitempty"` @@ -19,15 +23,21 @@ type ToolUse struct { // ReasoningTokens is disjoint from OutputTokens, per the api.Usage contract: // OpenAI reports reasoning as a subset of output, so it is netted out at this // parse boundary the way the live providers already net it. - ReasoningTokens int `json:"reasoning_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:"-"` + ReasoningTokens int `json:"reasoning_tokens,omitempty"` + CacheReadTokens int `json:"cache_read_tokens,omitempty"` + TotalTokens int `json:"total_tokens,omitempty"` + ContextWindow int `json:"context_window,omitempty"` + // CumulativeUsage is the provider's own running total for the session as of + // this record, rather than this record's delta. It is the result figure: + // reading the last one is exact, where summing per-record deltas drifts + // (a real 238-event codex session sums to 29.47M against a reported 29.24M). + // Netted to the disjoint api.Usage contract like the per-record fields. + CumulativeUsage *api.Usage `json:"cumulative_usage,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:"-"` // SourceLine is the 1-based JSONL line the use was extracted from — the // line of the FIRST record when several collapse into one row. It is the // stable identity of the row across re-parses of a growing transcript; diff --git a/pkg/ai/live_catalog.go b/pkg/ai/live_catalog.go index 9fc1a713..29a31b6b 100644 --- a/pkg/ai/live_catalog.go +++ b/pkg/ai/live_catalog.go @@ -18,7 +18,7 @@ func LiveCatalog() ([]Model, error) { if err != nil { return nil, err } - return mergeLiveCatalog(Catalog(), adapters), nil + return mergeLiveCatalog(Catalog(), adapters, liveCatalogOptions{}), nil } // LiveCatalogInfo annotates the live catalog with per-caller selectability, @@ -26,11 +26,15 @@ func LiveCatalog() ([]Model, error) { // provider key is present (configuredProviders), agent/CLI models when their // local backend binary is installed. func LiveCatalogInfo(configuredProviders []string) ([]ModelInfo, error) { - models, err := LiveCatalog() + adapters, err := CachedAdapters(time.Now()) if err != nil { return nil, err } - return catalogInfoFrom(models, configuredProviders), nil + models := mergeLiveCatalog(catalogSnapshot(), adapters, liveCatalogOptions{IncludeDisabled: true}) + return catalogInfoFrom(models, catalogInfoOptions{ + ConfiguredProviders: configuredProviders, + Adapters: adapters, + }), nil } // mergeLiveCatalog upserts each probed model onto the static catalog. Ordering @@ -41,7 +45,11 @@ func LiveCatalogInfo(configuredProviders []string) ([]ModelInfo, error) { // backend it collapses onto: without the menu-backend check, disabling a model // on the claude-agent card would still let the claude-cli probe re-add it under // the same menu id. -func mergeLiveCatalog(static []Model, adapters []AdapterStatus) []Model { +type liveCatalogOptions struct { + IncludeDisabled bool +} + +func mergeLiveCatalog(static []Model, adapters []AdapterStatus, options liveCatalogOptions) []Model { out := append([]Model(nil), static...) pos := make(map[string]int, len(out)) for i, m := range out { @@ -52,11 +60,11 @@ func mergeLiveCatalog(static []Model, adapters []AdapterStatus) []Model { for _, a := range adapters { probed := Backend(a.Backend) menuBackend, hasMenu := menuBackendFor(probed) - if !hasMenu || disabled.Backend(probed) { + if !hasMenu || (!options.IncludeDisabled && disabled.Backend(probed)) { continue } for _, md := range a.ModelDetails { - if disabled.Model(probed, md.ID) || disabled.Model(menuBackend, md.ID) { + if !options.IncludeDisabled && (disabled.Model(probed, md.ID) || disabled.Model(menuBackend, md.ID)) { continue } live := liveModel(menuBackend, md) diff --git a/pkg/ai/live_catalog_test.go b/pkg/ai/live_catalog_test.go index d7b710d6..377a631d 100644 --- a/pkg/ai/live_catalog_test.go +++ b/pkg/ai/live_catalog_test.go @@ -36,7 +36,7 @@ func TestMergeLiveCatalogUpsertsLiveAndPreservesStatic(t *testing.T) { }}, } - merged := mergeLiveCatalog(static, adapters) + merged := mergeLiveCatalog(static, adapters, liveCatalogOptions{}) // A live codex model becomes one codex-agent entry keyed by its exact id. sol := findModel(t, merged, "gpt-5.6-sol") @@ -107,4 +107,10 @@ func TestLiveCatalogInfoAppliesPerProviderConfigured(t *testing.T) { if !ok || openai.Configured { t.Errorf("openai model should be unconfigured (no key) but still listed: %+v", openai) } + if openai.Availability.State != api.AvailabilityMissingCredential || openai.Availability.Remediation == "" { + t.Errorf("openai availability = %+v, want missing credentials with remediation", openai.Availability) + } + if !anthropic.Availability.IsAvailable() { + t.Errorf("anthropic availability = %+v, want available", anthropic.Availability) + } } diff --git a/pkg/ai/loop.go b/pkg/ai/loop.go index 4f550be0..e39ddcf8 100644 --- a/pkg/ai/loop.go +++ b/pkg/ai/loop.go @@ -141,7 +141,7 @@ func runOneIteration(ctx context.Context, opts LoopOptions, req Request, iter *L if model == "" { model = opts.Provider.GetModel() } - iter.CostUSD = PriceResponse(backend, model, &Response{Backend: backend, Model: model, Usage: *ev.Usage}).Total() + iter.CostUSD = PriceUsage(backend, model, *ev.Usage, 0).Total() } if ev.Error != "" && iter.Err == nil { iter.Err = fmt.Errorf("claude returned: %s", ev.Error) diff --git a/pkg/ai/price_usage_ginkgo_test.go b/pkg/ai/price_usage_ginkgo_test.go new file mode 100644 index 00000000..b4a3a1b7 --- /dev/null +++ b/pkg/ai/price_usage_ginkgo_test.go @@ -0,0 +1,85 @@ +package ai + +import ( + "github.com/flanksource/captain/pkg/ai/pricing" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Observed on aichat thread 6f58a8f6: one claude-opus-5 turn on the anthropic +// backend. Opus 5 lists at $5/Mtok input and $25/Mtok output, so the turn's +// cost splits 0.640690 + 0.091700 = 0.732390 — the figure the UI rendered as a +// single opaque total while every per-bucket row showed "-". +const ( + opusModel = "claude-opus-5" + turnInputTokens = 128138 + turnOutputTokens = 3668 + turnInputUSD = 0.640690 + turnOutputUSD = 0.091700 +) + +var _ = Describe("PriceUsage", func() { + BeforeEach(func() { + pricing.EnsureLoaded(pricing.LoadOptions{}) + }) + + It("splits a turn across the per-bucket costs instead of collapsing it into one", func() { + usage := Usage{InputTokens: turnInputTokens, OutputTokens: turnOutputTokens} + + cost := PriceUsage(BackendAnthropic, opusModel, usage, 0) + + Expect(cost.InputCost).To(BeNumerically("~", turnInputUSD, 1e-6)) + Expect(cost.OutputCost).To(BeNumerically("~", turnOutputUSD, 1e-6)) + Expect(cost.Total()).To(BeNumerically("~", turnInputUSD+turnOutputUSD, 1e-6)) + }) + + It("prices cache reads and writes into their own buckets", func() { + // The claude-agent side of the same thread: near-zero input against a + // large cached prefix. Folding cache into input would misprice it by + // 10x, since cache reads list at $0.50/Mtok against input's $5.00. + // Expectations below are computed from those list rates, not from the + // function's own output. + const ( + agentInput = 74 + agentOutput = 31502 + agentCacheRead = 3108284 + agentCacheWrite = 185751 + + expectCacheReadUSD = agentCacheRead * 0.50 / 1e6 // 1.554142 + expectCacheWriteUSD = agentCacheWrite * 6.25 / 1e6 // 1.160944 + ) + usage := Usage{ + InputTokens: agentInput, OutputTokens: agentOutput, + CacheReadTokens: agentCacheRead, CacheWriteTokens: agentCacheWrite, + } + + cost := PriceUsage(BackendClaudeAgent, opusModel, usage, 0) + + Expect(cost.CacheReadCost).To(BeNumerically("~", expectCacheReadUSD, 1e-6)) + Expect(cost.CacheWriteCost).To(BeNumerically("~", expectCacheWriteUSD, 1e-6)) + Expect(cost.TotalTokens).To(Equal(agentInput + agentOutput + agentCacheRead + agentCacheWrite)) + }) + + It("prefers the provider's reported total over the list-price recompute", func() { + // The provider figure covers pricing captain's registry cannot model, + // such as 1-hour cache writes billed above the 5-minute rate. + usage := Usage{InputTokens: turnInputTokens, OutputTokens: turnOutputTokens} + const providerReported = 4.295061 + + cost := PriceUsage(BackendClaudeAgent, opusModel, usage, providerReported) + + Expect(cost.ProviderCostUSD).To(Equal(providerReported)) + Expect(cost.Total()).To(Equal(providerReported)) + Expect(cost.InputCost).To(BeNumerically("~", turnInputUSD, 1e-6), + "the list-price breakdown is still retained for display") + }) + + It("keeps token counts when the model is absent from the pricing registry", func() { + usage := Usage{InputTokens: 10, OutputTokens: 20} + + cost := PriceUsage(BackendAnthropic, "model-that-does-not-exist", usage, 0) + + Expect(cost.TotalTokens).To(Equal(30)) + Expect(cost.Total()).To(BeZero()) + }) +}) diff --git a/pkg/ai/pricing/openrouter.go b/pkg/ai/pricing/openrouter.go index 7151ad0d..0002fde9 100644 --- a/pkg/ai/pricing/openrouter.go +++ b/pkg/ai/pricing/openrouter.go @@ -59,28 +59,30 @@ func (c *PricingCache) IsExpired() bool { return time.Since(c.Timestamp) >= cacheExpiryDuration } -func EnsureLoaded() { +// LoadOptions controls how EnsureLoaded resolves the pricing snapshot. +type LoadOptions struct { + // Refresh skips both the in-memory and the on-disk snapshot and re-queries + // OpenRouter, so `captain whoami --no-cache` reports live prices. A failed + // refresh keeps the snapshot already installed instead of emptying pricing. + Refresh bool +} + +// EnsureLoaded installs the OpenRouter pricing snapshot, fetching it only when +// no usable one is already in memory or on disk (or when opts.Refresh forces it). +func EnsureLoaded(opts LoadOptions) { pricingCacheLock.Lock() defer pricingCacheLock.Unlock() - if pricingCacheErr != nil { - return - } - if pricingCache != nil && !pricingCache.IsExpired() { - return - } - - if pricingCache == nil { - if cache, err := loadFromDisk(); err == nil && cache != nil && !cache.IsExpired() { - log.Debugf("Loaded OpenRouter pricing from cache (age: %s)", time.Since(cache.Timestamp)) - pricingCache = cache - MergeModels(cache.Models) - applyCatalogPrices() + if !opts.Refresh { + if pricingCacheErr != nil { + return + } + if installedPricingIsFresh() { return } } - models, err := fetchOpenRouterPricing() + models, err := fetchPricing() if err != nil { log.Warnf("Failed to fetch OpenRouter pricing: %v", err) pricingCacheErr = err @@ -88,10 +90,33 @@ func EnsureLoaded() { } pricingCache = &PricingCache{Timestamp: time.Now(), Models: models} + pricingCacheErr = nil MergeModels(models) applyCatalogPrices() } +// installedPricingIsFresh reports whether an unexpired snapshot is in play, +// promoting the on-disk one into memory (and into the registry) when it is the +// only copy. Callers must hold pricingCacheLock. +func installedPricingIsFresh() bool { + if pricingCache != nil { + return !pricingCache.IsExpired() + } + cache, err := loadFromDisk() + if err != nil || cache == nil || cache.IsExpired() { + return false + } + log.Debugf("Loaded OpenRouter pricing from cache (age: %s)", time.Since(cache.Timestamp)) + pricingCache = cache + MergeModels(cache.Models) + applyCatalogPrices() + return true +} + +// fetchPricing is the live OpenRouter query, a package var so tests can +// substitute deterministic rows without hitting the network. +var fetchPricing = fetchOpenRouterPricing + func fetchOpenRouterPricing() (map[string]*ModelInfo, error) { resp, err := http.Get(openRouterAPIURL) if err != nil { diff --git a/pkg/ai/pricing/openrouter_test.go b/pkg/ai/pricing/openrouter_test.go new file mode 100644 index 00000000..8053af4a --- /dev/null +++ b/pkg/ai/pricing/openrouter_test.go @@ -0,0 +1,80 @@ +package pricing + +import ( + "errors" + "testing" + "time" +) + +// withStubbedFetch replaces the OpenRouter fetch and installs a fresh in-memory +// snapshot, so EnsureLoaded can be exercised without network or disk access. +func withStubbedFetch(t *testing.T, cached *PricingCache, models map[string]*ModelInfo, fetchErr error) *int { + t.Helper() + + calls := 0 + prevFetch := fetchPricing + fetchPricing = func() (map[string]*ModelInfo, error) { + calls++ + return models, fetchErr + } + + pricingCacheLock.Lock() + savedCache, savedErr := pricingCache, pricingCacheErr + pricingCache, pricingCacheErr = cached, nil + pricingCacheLock.Unlock() + + t.Cleanup(func() { + fetchPricing = prevFetch + pricingCacheLock.Lock() + pricingCache, pricingCacheErr = savedCache, savedErr + pricingCacheLock.Unlock() + }) + return &calls +} + +func freshSnapshot() *PricingCache { + return &PricingCache{ + Timestamp: time.Now().Add(-time.Hour), + Models: map[string]*ModelInfo{"x/cached": {ModelID: "x/cached", InputPrice: 1}}, + } +} + +func TestEnsureLoadedReusesFreshSnapshot(t *testing.T) { + calls := withStubbedFetch(t, freshSnapshot(), nil, nil) + + EnsureLoaded(LoadOptions{}) + + if *calls != 0 { + t.Fatalf("fetch calls = %d, want 0: a fresh snapshot must not re-query OpenRouter", *calls) + } +} + +func TestEnsureLoadedRefreshRefetchesFreshSnapshot(t *testing.T) { + refreshed := map[string]*ModelInfo{"x/refreshed": {ModelID: "x/refreshed", InputPrice: 2}} + calls := withStubbedFetch(t, freshSnapshot(), refreshed, nil) + + EnsureLoaded(LoadOptions{Refresh: true}) + + if *calls != 1 { + t.Fatalf("fetch calls = %d, want 1: refresh must bypass the unexpired snapshot", *calls) + } + if info, ok := GetModelInfo("x/refreshed"); !ok || info.InputPrice != 2 { + t.Fatalf("GetModelInfo(x/refreshed) = %+v, %v; want the refreshed price installed", info, ok) + } +} + +func TestEnsureLoadedRefreshKeepsSnapshotWhenFetchFails(t *testing.T) { + cached := freshSnapshot() + calls := withStubbedFetch(t, cached, nil, errors.New("openrouter unreachable")) + + EnsureLoaded(LoadOptions{Refresh: true}) + + if *calls != 1 { + t.Fatalf("fetch calls = %d, want 1", *calls) + } + pricingCacheLock.Lock() + defer pricingCacheLock.Unlock() + if pricingCache != cached { + t.Fatalf("pricingCache = %+v, want the previous snapshot retained after a failed refresh", pricingCache) + } +} diff --git a/pkg/ai/pricing/registry.go b/pkg/ai/pricing/registry.go index ec10ad49..c1fbe084 100644 --- a/pkg/ai/pricing/registry.go +++ b/pkg/ai/pricing/registry.go @@ -25,7 +25,7 @@ var ( ) func GetModelInfo(model string) (ModelInfo, bool) { - EnsureLoaded() + EnsureLoaded(LoadOptions{}) registryMu.RLock() info, ok := registry[model] registryMu.RUnlock() @@ -42,7 +42,7 @@ func GetModelInfo(model string) (ModelInfo, bool) { // GetModelInfo it has no static-Claude fallback (which prices any claude-ish id, // including typos), so use Contains for membership/validation, not for pricing. func Contains(model string) bool { - EnsureLoaded() + EnsureLoaded(LoadOptions{}) registryMu.RLock() defer registryMu.RUnlock() _, ok := registry[model] @@ -116,7 +116,7 @@ func RegistrySize() int { } func ListModels(filter string) []ModelInfo { - EnsureLoaded() + EnsureLoaded(LoadOptions{}) registryMu.RLock() defer registryMu.RUnlock() diff --git a/pkg/ai/pricing_coverage_test.go b/pkg/ai/pricing_coverage_test.go index 72d7599d..9f305bff 100644 --- a/pkg/ai/pricing_coverage_test.go +++ b/pkg/ai/pricing_coverage_test.go @@ -27,7 +27,7 @@ var knownPricingGaps = map[string]string{} // 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() + pricing.EnsureLoaded(pricing.LoadOptions{}) for _, p := range registry.Providers() { backend, err := p.BackendFor(registry.ModeAPI) @@ -76,7 +76,7 @@ func TestPricingPrefixIsNotCatalogPrefix(t *testing.T) { // 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() + pricing.EnsureLoaded(pricing.LoadOptions{}) for _, model := range []string{"claude-sonnet-5", "claude-opus-4-8"} { t.Run(model, func(t *testing.T) { diff --git a/pkg/ai/provider/caller_tools_ginkgo_test.go b/pkg/ai/provider/caller_tools_ginkgo_test.go new file mode 100644 index 00000000..aac33ad0 --- /dev/null +++ b/pkg/ai/provider/caller_tools_ginkgo_test.go @@ -0,0 +1,79 @@ +package provider + +import ( + "context" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Codex Agent caller tools", func() { + It("injects the same request-scoped MCP endpoint on start and resume", func() { + endpoint := &api.CallerToolEndpoint{ + Name: "captain", URL: "http://127.0.0.1:43210/mcp", + Headers: map[string]string{"Authorization": "Bearer secret"}, + } + request := ai.Request{ + SessionID: "thread-1", + Prompt: api.Prompt{User: "inspect"}, + } + + for _, params := range []map[string]any{ + buildThreadStartParams("gpt-5.4", request, endpoint), + buildResumeParams(request, endpoint), + } { + config, ok := params["config"].(map[string]any) + Expect(ok).To(BeTrue()) + servers, ok := config["mcp_servers"].(map[string]any) + Expect(ok).To(BeTrue()) + serverConfig, ok := servers["captain"].(map[string]any) + Expect(ok).To(BeTrue()) + Expect(serverConfig).To(HaveKeyWithValue("url", endpoint.URL)) + Expect(serverConfig).To(HaveKeyWithValue("http_headers", endpoint.Headers)) + Expect(serverConfig).To(HaveKeyWithValue("required", true)) + } + }) + + It("binds the private capability to the Captain session identity", func() { + provider, err := NewCodexAppServer(ai.Config{ + Model: api.Model{Name: "gpt-5.4"}, + CaptainSessionID: "captain-thread-1", + SessionID: "provider-session-1", + Tools: []api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return "ok", nil }, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(provider.Close) + + Expect(provider.prepareCallerTools(ai.Request{SessionID: "provider-session-1"})).To(Succeed()) + Expect(provider.callerTools).NotTo(BeNil()) + Expect(provider.callerTools.Headers["Authorization"]).To( + HavePrefix("Bearer cap_captain-thread-1."), + ) + Expect(strings.Count(provider.callerTools.Headers["Authorization"], ".")).To(Equal(1)) + }) + + It("does not require MCP when request preferences disable every caller tool", func() { + provider, err := NewCodexAppServer(ai.Config{ + Tools: []api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return "ok", nil }, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(provider.Close) + + request := ai.Request{ + ToolPreferences: api.ToolPreferences{"invoice_get": api.ToolModeOff}, + Permissions: api.Permissions{MCP: api.MCP{Disabled: true}}, + } + Expect(provider.prepareCallerTools(request)).To(Succeed()) + Expect(provider.callerTools).To(BeNil()) + }) +}) diff --git a/pkg/ai/provider/claudeagent/agent.ts b/pkg/ai/provider/claudeagent/agent.ts index 98b7e6d2..c7fd8dc2 100644 --- a/pkg/ai/provider/claudeagent/agent.ts +++ b/pkg/ai/provider/claudeagent/agent.ts @@ -6,7 +6,7 @@ // client -> server requests: // initialize {cwd, model, systemPrompt, appendSystemPrompt, allowedTools, // maxTurns, maxBudgetUsd, permissionMode, resume, approvalMode, -// outputSchema} +// outputSchema, mcpServers} // -> reply {ok:true} // prompt {text, attachments?} -> reply {accepted:true} // interrupt -> reply {} @@ -34,11 +34,22 @@ import { query } from "@anthropic-ai/claude-agent-sdk"; import type { Options, + PreToolUseHookInput, Query, SDKMessage, - SDKUserMessage, } from "@anthropic-ai/claude-agent-sdk"; import { createInterface } from "readline"; +import { + callHost, + diag, + handleResponse, + type JsonRpcId, + notify, + type PromptParams, + reply, + replyError, + TurnQueue, +} from "./protocol.js"; // Strip nested-session markers so the SDK does not refuse to run inside captain // (which may itself have been launched from a Claude Code session). The Go @@ -65,70 +76,11 @@ interface InitializeParams { // 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; - -function send(obj: Record) { - process.stdout.write(JSON.stringify(obj) + "\n"); -} -function notify(method: string, params: Record) { - send({ jsonrpc: "2.0", method, params }); -} -function reply(id: JsonRpcId, result: unknown) { - send({ jsonrpc: "2.0", id, result }); -} -function replyError(id: JsonRpcId, code: number, message: string) { - send({ jsonrpc: "2.0", id, error: { code, message } }); -} -function diag(msg: string) { - process.stderr.write(`[claude-agent] ${msg}\n`); -} - -// callHost issues a server->client request to the Go host and resolves when the -// matching id-bearing response arrives on stdin (routed by handleResponse). Ids -// are string-prefixed so they never collide with the host's numeric Call ids. -interface HostResponse { - result?: unknown; - error?: { code: number; message: string }; -} -let nextHostId = 1; -const pendingHostCalls = new Map void>(); - -function callHost( - method: string, - params: Record, -): Promise { - const id = `agent-${nextHostId++}`; - return new Promise((resolve, reject) => { - pendingHostCalls.set(id, (resp) => { - if (resp.error) { - reject(new Error(resp.error.message)); - } else { - resolve(resp.result); - } - }); - send({ jsonrpc: "2.0", id, method, params }); - }); -} - -// handleResponse resolves a pending callHost when a host response (id, no method) -// arrives. Returns true if the frame was a response we were waiting for. -function handleResponse(frame: { - id?: JsonRpcId; - result?: unknown; - error?: { code: number; message: string }; -}): boolean { - if (frame.id == null || typeof frame.id !== "string") { - return false; - } - const waiter = pendingHostCalls.get(frame.id); - if (!waiter) { - return false; - } - pendingHostCalls.delete(frame.id); - waiter({ result: frame.result, error: frame.error }); - return true; + mcpServers?: Record< + string, + { type: "http"; url: string; headers?: Record } + >; + callerToolUseIDKey?: string; } // HostDecision is the can_use_tool reply shape from the Go host. @@ -138,115 +90,18 @@ interface HostDecision { updatedInput?: Record; } -// TurnQueue is a push async-iterable of SDKUserMessage. Pushing a user message -// resolves the SDK's pending next() so a single query() session processes turns -// as they arrive instead of ending after the first. -class TurnQueue implements AsyncIterable { - private pending: SDKUserMessage[] = []; - private waiters: ((r: IteratorResult) => void)[] = []; - private ended = false; - - 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, - }, - parent_tool_use_id: null, - session_id: "", - }; - const waiter = this.waiters.shift(); - if (waiter) { - waiter({ value: msg, done: false }); - } else { - this.pending.push(msg); - } - } - - end() { - this.ended = true; - const waiter = this.waiters.shift(); - if (waiter) { - waiter({ value: undefined as unknown as SDKUserMessage, done: true }); - } - } - - [Symbol.asyncIterator](): AsyncIterator { - return { - next: (): Promise> => { - const queued = this.pending.shift(); - if (queued) { - return Promise.resolve({ value: queued, done: false }); - } - if (this.ended) { - return Promise.resolve({ - value: undefined as unknown as SDKUserMessage, - done: true, - }); - } - return new Promise((resolve) => this.waiters.push(resolve)); - }, - }; - } -} - let turns: TurnQueue | null = null; let activeQuery: Query | null = null; +let callerToolServers: string[] = []; -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 callerToolName(toolName: string): string | undefined { + for (const server of callerToolServers) { + const prefix = `mcp__${server}__`; + if (toolName.startsWith(prefix) && toolName.length > prefix.length) { + return toolName.slice(prefix.length); + } + } + return undefined; } function buildOptions(params: InitializeParams): Options { @@ -266,6 +121,7 @@ function buildOptions(params: InitializeParams): Options { params.allowedTools && params.allowedTools.length ? params.allowedTools : undefined, + mcpServers: params.mcpServers, stderr: (data: string) => process.stderr.write(data), hooks: { PreToolUse: [ @@ -289,6 +145,49 @@ function buildOptions(params: InitializeParams): Options { }, }; + if (callerToolServers.length > 0) { + const callerToolUseIDKey = params.callerToolUseIDKey; + if (!callerToolUseIDKey) { + throw new Error("caller tools require a provider tool-use ID key"); + } + options.hooks?.PreToolUse?.push({ + hooks: [ + async (input, toolUseID) => { + const hook = input as PreToolUseHookInput; + if (!callerToolName(hook.tool_name)) { + return {}; + } + if (!toolUseID && !hook.tool_use_id) { + return { + decision: "block" as const, + reason: "caller tool has no Claude tool-use ID", + }; + } + if ( + typeof hook.tool_input !== "object" || + hook.tool_input === null || + Array.isArray(hook.tool_input) + ) { + return { + decision: "block" as const, + reason: "caller tool input must be an object", + }; + } + return { + hookSpecificOutput: { + hookEventName: "PreToolUse" as const, + permissionDecision: "allow" as const, + updatedInput: { + ...(hook.tool_input as Record), + [callerToolUseIDKey]: toolUseID || hook.tool_use_id, + }, + }, + }; + }, + ], + }); + } + // 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. @@ -360,14 +259,19 @@ function buildOptions(params: InitializeParams): Options { // PreToolUse git add/commit block above still applies first. if (brokered) { options.canUseTool = async (toolName, input, opts) => { - const toolUseId = - (opts as { toolUseId?: string } | undefined)?.toolUseId ?? ""; + if ( + Object.keys(params.mcpServers ?? {}).some((server) => + toolName.startsWith(`mcp__${server}__`), + ) + ) { + return { behavior: "allow", updatedInput: input }; + } let decision: HostDecision; try { decision = (await callHost("can_use_tool", { tool: toolName, input, - tool_use_id: toolUseId, + tool_use_id: opts.toolUseID, })) as HostDecision; } catch (err) { return { @@ -393,6 +297,7 @@ function handleInitialize(id: JsonRpcId, params: InitializeParams) { return; } try { + callerToolServers = Object.keys(params.mcpServers ?? {}); turns = new TurnQueue(); activeQuery = query({ prompt: turns, options: buildOptions(params) }); reply(id, { ok: true }); @@ -400,6 +305,7 @@ function handleInitialize(id: JsonRpcId, params: InitializeParams) { notify("turn/error", { message: err?.message || String(err) }); }); } catch (err) { + callerToolServers = []; turns = null; activeQuery = null; replyError(id, -32603, `initialize failed: ${(err as Error)?.message || err}`); @@ -487,7 +393,7 @@ function handleMessage(message: SDKMessage) { notify("message/thinking", { text: block.thinking }); } else if (block.type === "tool_use") { notify("message/tool_use", { - tool: block.name, + tool: callerToolName(String(block.name)) ?? block.name, input: block.input, id: block.id, }); diff --git a/pkg/ai/provider/claudeagent/attachments_ginkgo_test.go b/pkg/ai/provider/claudeagent/attachments_ginkgo_test.go index 12c369f7..fa25d0af 100644 --- a/pkg/ai/provider/claudeagent/attachments_ginkgo_test.go +++ b/pkg/ai/provider/claudeagent/attachments_ginkgo_test.go @@ -21,9 +21,12 @@ var _ = Describe("Claude Agent prompt parameters", func() { It("materializes the structured attachment bridge source", func() { directory, err := prepareAgentDir() Expect(err).NotTo(HaveOccurred()) - content, err := os.ReadFile(filepath.Join(directory, "agent.ts")) + content, err := os.ReadFile(filepath.Join(directory, "protocol.ts")) Expect(err).NotTo(HaveOccurred()) Expect(string(content)).To(ContainSubstring("attachments?: PromptAttachment[]")) + agent, err := os.ReadFile(filepath.Join(directory, "agent.ts")) + Expect(err).NotTo(HaveOccurred()) + Expect(string(agent)).To(ContainSubstring(`from "./protocol.js"`)) }) It("encodes prepared image and PDF data as ordered structured inputs", func() { diff --git a/pkg/ai/provider/claudeagent/bridge_params.go b/pkg/ai/provider/claudeagent/bridge_params.go new file mode 100644 index 00000000..0aead56c --- /dev/null +++ b/pkg/ai/provider/claudeagent/bridge_params.go @@ -0,0 +1,20 @@ +package claudeagent + +import "encoding/json" + +type initializeParams struct { + Cwd string `json:"cwd,omitempty"` + Model string `json:"model,omitempty"` + SystemPrompt string `json:"systemPrompt,omitempty"` + AppendSystemPrompt string `json:"appendSystemPrompt,omitempty"` + AllowedTools []string `json:"allowedTools,omitempty"` + MaxTurns int `json:"maxTurns,omitempty"` + MaxBudgetUsd float64 `json:"maxBudgetUsd,omitempty"` + PermissionMode string `json:"permissionMode,omitempty"` + Resume string `json:"resume,omitempty"` + ApprovalMode string `json:"approvalMode,omitempty"` + OutputSchema json.RawMessage `json:"outputSchema,omitempty"` + MonitorURL string `json:"monitorUrl,omitempty"` + MCPServers map[string]callerToolServer `json:"mcpServers,omitempty"` + CallerToolUseIDKey string `json:"callerToolUseIDKey,omitempty"` +} diff --git a/pkg/ai/provider/claudeagent/caller_tools.go b/pkg/ai/provider/claudeagent/caller_tools.go new file mode 100644 index 00000000..e7c07599 --- /dev/null +++ b/pkg/ai/provider/claudeagent/caller_tools.go @@ -0,0 +1,78 @@ +package claudeagent + +import ( + "fmt" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/callertools" + aitools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" +) + +type callerToolServer struct { + Type string `json:"type"` + URL string `json:"url"` + Headers map[string]string `json:"headers,omitempty"` +} + +func (p *Provider) prepareCallerTools(req ai.Request) error { + p.callerToolsMu.Lock() + defer p.callerToolsMu.Unlock() + if p.callerTools != nil { + if req.Permissions.MCP.Disabled { + return fmt.Errorf("claude-agent: caller tools require MCP but MCP is disabled") + } + return p.callerTools.Validate() + } + if len(p.cfg.Tools) == 0 { + return nil + } + definitions, err := aitools.ResolveDefinitions(p.cfg.Tools, req.ToolPreferences) + if err != nil { + return fmt.Errorf("claude-agent caller tools: %w", err) + } + if len(definitions) == 0 { + return nil + } + if req.Permissions.MCP.Disabled { + return fmt.Errorf("claude-agent: caller tools require MCP but MCP is disabled") + } + runtime, err := callertools.New(callertools.Options{ + Definitions: definitions, CanUseTool: p.cfg.CanUseTool, + SessionID: firstNonEmpty(p.cfg.CaptainSessionID, req.SessionID, p.cfg.SessionID), + }) + if err != nil { + return fmt.Errorf("start claude-agent caller tools: %w", err) + } + endpoint := runtime.Endpoint() + p.callerToolsRuntime = runtime + p.callerTools = &endpoint + return nil +} + +func callerToolServers(endpoint *api.CallerToolEndpoint) map[string]callerToolServer { + if endpoint == nil { + return nil + } + return map[string]callerToolServer{ + endpoint.Name: {Type: "http", URL: endpoint.URL, Headers: cloneHeaders(endpoint.Headers)}, + } +} + +func callerToolUseIDKey(endpoint *api.CallerToolEndpoint) string { + if endpoint == nil { + return "" + } + return callertools.ToolUseIDInputKey +} + +func cloneHeaders(headers map[string]string) map[string]string { + if headers == nil { + return nil + } + cloned := make(map[string]string, len(headers)) + for key, value := range headers { + cloned[key] = value + } + return cloned +} diff --git a/pkg/ai/provider/claudeagent/caller_tools_ginkgo_test.go b/pkg/ai/provider/claudeagent/caller_tools_ginkgo_test.go new file mode 100644 index 00000000..83640334 --- /dev/null +++ b/pkg/ai/provider/claudeagent/caller_tools_ginkgo_test.go @@ -0,0 +1,131 @@ +package claudeagent + +import ( + "context" + "encoding/json" + "os" + "strings" + "sync/atomic" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/clicky/exec" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Claude Agent caller tools", func() { + It("injects a request-scoped HTTP MCP endpoint", func() { + provider := &Provider{ + model: "claude-sonnet-5", + callerTools: &api.CallerToolEndpoint{ + Name: "captain", URL: "http://127.0.0.1:43210/mcp", + Headers: map[string]string{"Authorization": "Bearer secret"}, + }, + } + + params := provider.initializeParams(ai.Request{Prompt: api.Prompt{User: "inspect"}}) + + Expect(params.MCPServers).To(HaveKey("captain")) + Expect(params.MCPServers["captain"].Type).To(Equal("http")) + Expect(params.MCPServers["captain"].URL).To(Equal("http://127.0.0.1:43210/mcp")) + Expect(params.MCPServers["captain"].Headers).To(HaveKeyWithValue("Authorization", "Bearer secret")) + raw, err := json.Marshal(params) + Expect(err).NotTo(HaveOccurred()) + Expect(string(raw)).To(ContainSubstring(`"callerToolUseIDKey":"__captain_tool_use_id"`)) + }) + + It("executes an allowed caller tool through the fake Claude runtime", func(ctx SpecContext) { + self, err := os.Executable() + Expect(err).NotTo(HaveOccurred()) + original := newAgentProcess + newAgentProcess = func(*Provider) (*exec.Process, error) { + return exec.NewExec(self).WithStdioPipe().WithEnv(map[string]string{ + fakeServerEnv: "1", fakeModeEnv: "caller-tools", + }), nil + } + DeferCleanup(func() { newAgentProcess = original }) + + var calls atomic.Int32 + permissions := make(chan api.PermissionRequest, 1) + provider, err := New(ai.Config{ + Model: api.Model{Name: "claude-sonnet-5"}, + CaptainSessionID: "captain-thread-1", + CanUseTool: func(_ context.Context, request api.PermissionRequest) (api.PermissionDecision, error) { + permissions <- request + return api.PermissionDecision{Allow: true}, nil + }, + Tools: []api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeAsk, + Handler: func(_ context.Context, input map[string]any) (any, error) { + calls.Add(1) + return map[string]any{"id": input["id"], "status": "draft"}, nil + }, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(provider.Close) + + events, err := provider.ExecuteStream(ctx, ai.Request{Prompt: api.Prompt{User: "inspect invoice"}}) + Expect(err).NotTo(HaveOccurred()) + var toolUse, toolResult api.Event + for event := range events { + switch event.Kind { + case api.EventToolUse: + toolUse = event + case api.EventToolResult: + toolResult = event + } + } + Expect(calls.Load()).To(Equal(int32(1))) + Expect(toolUse.Tool).To(Equal("invoice_get")) + Expect(toolUse.ToolCallID).To(Equal("claude-tool-use-1")) + var permission api.PermissionRequest + Eventually(permissions).Should(Receive(&permission)) + Expect(permission.Tool).To(Equal(toolUse.Tool)) + Expect(permission.ToolUseID).To(Equal(toolUse.ToolCallID)) + Expect(permission.Input).To(Equal(map[string]any{"id": "inv-1"})) + Expect(toolResult.ToolCallID).To(Equal(toolUse.ToolCallID)) + Expect(toolResult.Success).To(BeTrue()) + }) + + It("binds the private capability to the Captain session identity", func() { + provider, err := New(ai.Config{ + Model: api.Model{Name: "claude-sonnet-5"}, + CaptainSessionID: "captain-thread-1", + SessionID: "provider-session-1", + Tools: []api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return "ok", nil }, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(provider.Close) + + Expect(provider.prepareCallerTools(ai.Request{SessionID: "provider-session-1"})).To(Succeed()) + Expect(provider.callerTools).NotTo(BeNil()) + Expect(provider.callerTools.Headers["Authorization"]).To( + HavePrefix("Bearer cap_captain-thread-1."), + ) + Expect(strings.Count(provider.callerTools.Headers["Authorization"], ".")).To(Equal(1)) + }) + + It("does not require MCP when request preferences disable every caller tool", func() { + provider, err := New(ai.Config{ + Tools: []api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return "ok", nil }, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(provider.Close) + + request := ai.Request{ + ToolPreferences: api.ToolPreferences{"invoice_get": api.ToolModeOff}, + Permissions: api.Permissions{MCP: api.MCP{Disabled: true}}, + } + Expect(provider.prepareCallerTools(request)).To(Succeed()) + Expect(provider.callerTools).To(BeNil()) + }) +}) diff --git a/pkg/ai/provider/claudeagent/fake_agent_test.go b/pkg/ai/provider/claudeagent/fake_agent_test.go new file mode 100644 index 00000000..792f23f5 --- /dev/null +++ b/pkg/ai/provider/claudeagent/fake_agent_test.go @@ -0,0 +1,227 @@ +package claudeagent + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "testing" + "time" + + "github.com/flanksource/clicky/exec" + mcpclient "github.com/mark3labs/mcp-go/client" + "github.com/mark3labs/mcp-go/client/transport" + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/require" +) + +const ( + fakeServerEnv = "CLAUDEAGENT_FAKE_SERVER" + fakeModeEnv = "CLAUDEAGENT_FAKE_MODE" + fakeMarkerEnv = "CLAUDEAGENT_FAKE_MARKER" +) + +type fakeCallerToolServer struct { + URL string `json:"url"` + Headers map[string]string `json:"headers"` +} + +type fakeInitializeParams struct { + OutputSchema json.RawMessage `json:"outputSchema"` + MCPServers map[string]fakeCallerToolServer `json:"mcpServers"` + CallerToolUseIDKey string `json:"callerToolUseIDKey"` +} + +func TestMain(m *testing.M) { + if os.Getenv(fakeServerEnv) == "1" { + runFakeServer() + os.Exit(0) + } + os.Exit(m.Run()) +} + +func runFakeServer() { + mode := os.Getenv(fakeModeEnv) + marker := os.Getenv(fakeMarkerEnv) + var initialization fakeInitializeParams + initHadSchema := false + promptCount := 0 + + enc := func(obj map[string]any) { + encoded, _ := json.Marshal(obj) + _, _ = os.Stdout.Write(append(encoded, '\n')) + } + id := func(raw json.RawMessage) any { + if len(raw) == 0 { + return nil + } + return raw + } + completed := func(resultText string) { + enc(map[string]any{"jsonrpc": "2.0", "method": "turn/completed", "params": map[string]any{ + "success": true, "session_id": "fake-sess", "cost_usd": 0.01, + "result_text": resultText, + "usage": map[string]any{"input_tokens": 10, "output_tokens": 5}, + }}) + } + + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 64*1024), 4*1024*1024) + for scanner.Scan() { + var frame struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Result json.RawMessage `json:"result"` + Params fakeInitializeParams `json:"params"` + } + if json.Unmarshal(scanner.Bytes(), &frame) != nil { + continue + } + if frame.Method == "" && len(frame.Result) > 0 { + completed("decision=" + string(frame.Result)) + continue + } + switch frame.Method { + case "initialize": + initialization = frame.Params + initHadSchema = len(frame.Params.OutputSchema) > 0 + enc(map[string]any{"jsonrpc": "2.0", "id": id(frame.ID), "result": map[string]any{"ok": true}}) + enc(map[string]any{"jsonrpc": "2.0", "method": "session/init", "params": map[string]any{ + "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"}}) + runFakeTurn(mode, promptCount, initHadSchema, initialization, enc, completed) + case "interrupt": + if marker != "" { + _ = os.WriteFile(marker, []byte("interrupted"), 0o644) + } + enc(map[string]any{"jsonrpc": "2.0", "id": id(frame.ID), "result": map[string]any{}}) + case "shutdown": + enc(map[string]any{"jsonrpc": "2.0", "id": id(frame.ID), "result": map[string]any{}}) + os.Exit(0) + } + } +} + +func runFakeTurn( + mode string, + promptCount int, + initHadSchema bool, + initialization fakeInitializeParams, + enc func(map[string]any), + completed func(string), +) { + switch mode { + case "error-output": + _, _ = os.Stderr.WriteString("claude subprocess authentication detail\n") + enc(map[string]any{"jsonrpc": "2.0", "method": "turn/error", "params": map[string]any{ + "message": "Claude Code process exited with code 1", + }}) + case "steer": + if promptCount == 2 { + completed("first prompt complete") + completed("steered prompt complete") + } + case "approval": + 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": + 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": + case "structured": + enc(map[string]any{"jsonrpc": "2.0", "method": "turn/completed", "params": map[string]any{ + "success": true, "session_id": "fake-sess", "cost_usd": 0.01, "subtype": "success", + "usage": map[string]any{"input_tokens": 10, "output_tokens": 5}, + "structured_output": map[string]any{ + "company_name": "Anthropic", "founded_year": 2021, "received_schema": initHadSchema, + }, + }}) + case "caller-tools": + if err := runFakeCallerTool(initialization, enc, completed); err != nil { + enc(map[string]any{"jsonrpc": "2.0", "method": "turn/error", "params": map[string]any{"message": err.Error()}}) + } + default: + enc(map[string]any{"jsonrpc": "2.0", "method": "message/tool_use", "params": map[string]any{ + "tool": "Read", "id": "t1", "input": map[string]any{"file_path": "/x"}, + }}) + completed("hi from fake") + } +} + +func runFakeCallerTool( + initialization fakeInitializeParams, + enc func(map[string]any), + completed func(string), +) error { + server, ok := initialization.MCPServers["captain"] + if !ok || initialization.CallerToolUseIDKey == "" { + return fmt.Errorf("fake caller tools were not initialized") + } + const toolUseID = "claude-tool-use-1" + enc(map[string]any{"jsonrpc": "2.0", "method": "message/tool_use", "params": map[string]any{ + "tool": "invoice_get", "id": toolUseID, "input": map[string]any{"id": "inv-1"}, + }}) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + channel, err := transport.NewStreamableHTTP(server.URL, transport.WithHTTPHeaders(server.Headers)) + if err != nil { + return err + } + client := mcpclient.NewClient(channel) + if err := client.Start(ctx); err != nil { + return err + } + defer client.Close() + request := mcp.InitializeRequest{} + request.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION + request.Params.ClientInfo = mcp.Implementation{Name: "claude-agent-fake", Version: "1.0.0"} + if _, err := client.Initialize(ctx, request); err != nil { + return err + } + call := mcp.CallToolRequest{} + call.Params.Name = "invoice_get" + call.Params.Arguments = map[string]any{ + "id": "inv-1", initialization.CallerToolUseIDKey: toolUseID, + } + result, err := client.CallTool(ctx, call) + if err != nil { + return err + } + content, err := json.Marshal(result.StructuredContent) + if err != nil { + return err + } + enc(map[string]any{"jsonrpc": "2.0", "method": "message/tool_result", "params": map[string]any{ + "id": toolUseID, "content": string(content), "is_error": result.IsError, + }}) + completed("caller tool complete") + return nil +} + +func withFakeAgentProcess(t *testing.T) { + t.Helper() + withFakeAgentProcessEnv(t, map[string]string{fakeServerEnv: "1"}) +} + +func withFakeAgentProcessEnv(t *testing.T, env map[string]string) { + t.Helper() + self, err := os.Executable() + require.NoError(t, err) + original := newAgentProcess + newAgentProcess = func(*Provider) (*exec.Process, error) { + return exec.NewExec(self).WithStdioPipe().WithEnv(env), nil + } + t.Cleanup(func() { newAgentProcess = original }) +} diff --git a/pkg/ai/provider/claudeagent/interrupt_ginkgo_test.go b/pkg/ai/provider/claudeagent/interrupt_ginkgo_test.go new file mode 100644 index 00000000..429ff7df --- /dev/null +++ b/pkg/ai/provider/claudeagent/interrupt_ginkgo_test.go @@ -0,0 +1,33 @@ +package claudeagent + +import ( + "encoding/json" + "time" + + "github.com/flanksource/captain/pkg/ai" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Claude Agent interruption", func() { + It("does not queue a terminal result while an interrupt is in progress", func() { + turn := &turnState{ + inbox: make(chan ai.Event, 1), + term: make(chan struct{}), + quit: make(chan struct{}), + pending: 1, + } + turn.interrupting.Store(true) + provider := &Provider{model: testModel} + provider.setActive(turn) + + provider.onNotification(notifyTurnDone, json.RawMessage(`{ + "success":false, + "subtype":"error_during_execution", + "session_id":"session-interrupted" + }`)) + + Consistently(turn.inbox, 50*time.Millisecond).ShouldNot(Receive()) + Eventually(turn.term).Should(BeClosed()) + }) +}) diff --git a/pkg/ai/provider/claudeagent/process_env.go b/pkg/ai/provider/claudeagent/process_env.go new file mode 100644 index 00000000..05aae259 --- /dev/null +++ b/pkg/ai/provider/claudeagent/process_env.go @@ -0,0 +1,21 @@ +package claudeagent + +import "github.com/flanksource/captain/pkg/ai" + +func agentProcessEnv(cfg ai.Config, environ []string) map[string]string { + env := nestingEnvOverrides(environ) + if cfg.APIURL != "" { + env["ANTHROPIC_BASE_URL"] = cfg.APIURL + env["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"] = "1" + env["DISABLE_NON_ESSENTIAL_MODEL_CALLS"] = "1" + env["DISABLE_TELEMETRY"] = "1" + env["DISABLE_ERROR_REPORTING"] = "1" + env["DISABLE_AUTOUPDATER"] = "1" + env["DISABLE_BUG_COMMAND"] = "1" + } + if cfg.APIKey != "" { + env["ANTHROPIC_API_KEY"] = cfg.APIKey + env["ANTHROPIC_AUTH_TOKEN"] = cfg.APIKey + } + return env +} diff --git a/pkg/ai/provider/claudeagent/protocol.ts b/pkg/ai/provider/claudeagent/protocol.ts new file mode 100644 index 00000000..8aabc28b --- /dev/null +++ b/pkg/ai/provider/claudeagent/protocol.ts @@ -0,0 +1,167 @@ +import type { SDKUserMessage } from "@anthropic-ai/claude-agent-sdk"; + +export type JsonRpcId = number | string | null; + +export interface PromptAttachment { + mediaType: string; + data: string; + filename?: string; +} + +export interface PromptParams { + text?: string; + attachments?: PromptAttachment[]; +} + +export function send(obj: Record) { + process.stdout.write(JSON.stringify(obj) + "\n"); +} + +export function notify(method: string, params: Record) { + send({ jsonrpc: "2.0", method, params }); +} + +export function reply(id: JsonRpcId, result: unknown) { + send({ jsonrpc: "2.0", id, result }); +} + +export function replyError(id: JsonRpcId, code: number, message: string) { + send({ jsonrpc: "2.0", id, error: { code, message } }); +} + +export function diag(msg: string) { + process.stderr.write(`[claude-agent] ${msg}\n`); +} + +interface HostResponse { + result?: unknown; + error?: { code: number; message: string }; +} + +let nextHostId = 1; +const pendingHostCalls = new Map void>(); + +export function callHost( + method: string, + params: Record, +): Promise { + const id = `agent-${nextHostId++}`; + return new Promise((resolve, reject) => { + pendingHostCalls.set(id, (resp) => { + if (resp.error) { + reject(new Error(resp.error.message)); + } else { + resolve(resp.result); + } + }); + send({ jsonrpc: "2.0", id, method, params }); + }); +} + +export function handleResponse(frame: { + id?: JsonRpcId; + result?: unknown; + error?: { code: number; message: string }; +}): boolean { + if (frame.id == null || typeof frame.id !== "string") { + return false; + } + const waiter = pendingHostCalls.get(frame.id); + if (!waiter) { + return false; + } + pendingHostCalls.delete(frame.id); + waiter({ result: frame.result, error: frame.error }); + return true; +} + +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, + ); +} + +export class TurnQueue implements AsyncIterable { + private pending: SDKUserMessage[] = []; + private waiters: ((result: IteratorResult) => void)[] = []; + private ended = false; + + 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 message: SDKUserMessage = { + type: "user", + message: { role: "user", content }, + parent_tool_use_id: null, + session_id: "", + }; + const waiter = this.waiters.shift(); + if (waiter) { + waiter({ value: message, done: false }); + } else { + this.pending.push(message); + } + } + + end() { + this.ended = true; + const waiter = this.waiters.shift(); + if (waiter) { + waiter({ value: undefined as unknown as SDKUserMessage, done: true }); + } + } + + [Symbol.asyncIterator](): AsyncIterator { + return { + next: (): Promise> => { + const queued = this.pending.shift(); + if (queued) { + return Promise.resolve({ value: queued, done: false }); + } + if (this.ended) { + return Promise.resolve({ + value: undefined as unknown as SDKUserMessage, + done: true, + }); + } + return new Promise((resolve) => this.waiters.push(resolve)); + }, + }; + } +} diff --git a/pkg/ai/provider/claudeagent/provider.go b/pkg/ai/provider/claudeagent/provider.go index 22a47a8d..8ef7c2d7 100644 --- a/pkg/ai/provider/claudeagent/provider.go +++ b/pkg/ai/provider/claudeagent/provider.go @@ -24,6 +24,7 @@ import ( "time" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/callertools" "github.com/flanksource/captain/pkg/ai/provider/jsonrpc" "github.com/flanksource/captain/pkg/api" "github.com/flanksource/clicky/exec" @@ -66,7 +67,7 @@ var _ ai.StreamingProvider = (*Provider)(nil) // newAgentProcess builds the supervised child command. It is a package var so // tests can substitute a fake JSON-RPC server without npm or a claude binary. -var newAgentProcess = func(*Provider) (*exec.Process, error) { +var newAgentProcess = func(provider *Provider) (*exec.Process, error) { agentDir, err := prepareAgentDir() if err != nil { return nil, err @@ -83,7 +84,7 @@ var newAgentProcess = func(*Provider) (*exec.Process, error) { return exec.NewExec(tsxPath, agentTSPath). WithCwd(agentDir). WithStdioPipe(). - WithEnv(nestingEnvOverrides(os.Environ())), nil + WithEnv(agentProcessEnv(provider.cfg, os.Environ())), nil } // Provider drives a supervised Claude Agent SDK process over JSON-RPC. @@ -120,6 +121,10 @@ type Provider struct { // so it is pinned from the first turn and every later turn must match it. sessionSchemaOnce sync.Once sessionSchema json.RawMessage + + callerToolsMu sync.Mutex + callerToolsRuntime *callertools.Runtime + callerTools *api.CallerToolEndpoint } // New builds a claude-agent provider. The supervised process is started lazily @@ -131,18 +136,27 @@ func New(cfg ai.Config) (*Provider, error) { } model = ai.NormalizeModelForBackend(ai.BackendClaudeAgent, model) ctx, cancel := context.WithCancel(context.Background()) - return &Provider{ + provider := &Provider{ model: model, cfg: cfg, baseCtx: ctx, baseCancel: cancel, initDone: make(chan struct{}), procExited: make(chan struct{}), - }, nil + } + if cfg.CallerTools != nil { + endpoint := *cfg.CallerTools + endpoint.Headers = cloneHeaders(cfg.CallerTools.Headers) + provider.callerTools = &endpoint + } + return provider, nil } -func (p *Provider) GetModel() string { return p.model } -func (p *Provider) GetBackend() ai.Backend { return ai.BackendClaudeAgent } +func (p *Provider) GetModel() string { return p.model } +func (p *Provider) GetBackend() ai.Backend { return ai.BackendClaudeAgent } +func (p *Provider) SupportsCallerTools() bool { return true } + +var _ api.ToolCapableProvider = (*Provider)(nil) // Execute drains its own ExecuteStream into a buffered ai.Response. When the // request carries a structured-output schema, the validated JSON the SDK @@ -269,6 +283,9 @@ func (p *Provider) ExecuteStream(ctx context.Context, req ai.Request) (<-chan ai // structured session, or a differing schema, cannot be honoured). p.sessionSchemaOnce.Do(func() { p.sessionSchema = schema }) + if err := p.prepareCallerTools(req); err != nil { + return nil, err + } if err := p.ensureStarted(req); err != nil { return nil, err } @@ -307,6 +324,9 @@ func (p *Provider) Close() error { if p.baseCancel != nil { p.baseCancel() } + if p.callerToolsRuntime != nil { + return p.callerToolsRuntime.Close() + } return nil } @@ -438,6 +458,8 @@ func (p *Provider) initializeParams(req ai.Request) initializeParams { ApprovalMode: approvalMode, OutputSchema: p.sessionSchema, MonitorURL: monitorHooksURL(req), + MCPServers: callerToolServers(p.callerTools), + CallerToolUseIDKey: callerToolUseIDKey(p.callerTools), } } @@ -451,21 +473,6 @@ func monitorHooksURL(req ai.Request) string { return api.ServeBaseURL() } -type initializeParams struct { - Cwd string `json:"cwd,omitempty"` - Model string `json:"model,omitempty"` - SystemPrompt string `json:"systemPrompt,omitempty"` - AppendSystemPrompt string `json:"appendSystemPrompt,omitempty"` - AllowedTools []string `json:"allowedTools,omitempty"` - MaxTurns int `json:"maxTurns,omitempty"` - MaxBudgetUsd float64 `json:"maxBudgetUsd,omitempty"` - PermissionMode string `json:"permissionMode,omitempty"` - 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) { p.initMu.Lock() defer p.initMu.Unlock() diff --git a/pkg/ai/provider/claudeagent/provider_test.go b/pkg/ai/provider/claudeagent/provider_test.go index 094b0443..201dd238 100644 --- a/pkg/ai/provider/claudeagent/provider_test.go +++ b/pkg/ai/provider/claudeagent/provider_test.go @@ -1,188 +1,18 @@ package claudeagent import ( - "bufio" "context" "encoding/json" - "os" "path/filepath" "testing" "time" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" - "github.com/flanksource/clicky/exec" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// fakeServerEnv switches the test binary into a fake JSON-RPC agent server when -// re-exec'd by the supervised process, so the provider lifecycle is exercised -// without npm, tsx or a real claude binary. -const fakeServerEnv = "CLAUDEAGENT_FAKE_SERVER" - -// fakeModeEnv selects the fake server's turn behaviour: "" runs the default -// happy-path turn; "approval" emits a can_use_tool request and finishes only -// after the host replies; "hang" emits text then waits for an interrupt; -// "error-output" emits both process streams before a terminal error. -const fakeModeEnv = "CLAUDEAGENT_FAKE_MODE" - -// fakeMarkerEnv, when set in "hang" mode, names a file the fake creates when it -// receives an interrupt — proof the graceful control request arrived (no kill). -const fakeMarkerEnv = "CLAUDEAGENT_FAKE_MARKER" - -func TestMain(m *testing.M) { - if os.Getenv(fakeServerEnv) == "1" { - runFakeServer() - os.Exit(0) - } - os.Exit(m.Run()) -} - -// runFakeServer speaks the agent.ts JSON-RPC protocol over stdio: it answers -// initialize/prompt/interrupt/shutdown and pushes a scripted set of turn -// notifications so the Go provider has a realistic stream to map. The turn shape -// is selected by fakeModeEnv so a single binary covers the happy path, the -// can_use_tool round-trip, and the interrupt-without-kill path. -func runFakeServer() { - mode := os.Getenv(fakeModeEnv) - marker := os.Getenv(fakeMarkerEnv) - - enc := func(obj map[string]any) { - b, _ := json.Marshal(obj) - _, _ = os.Stdout.Write(append(b, '\n')) - } - id := func(raw json.RawMessage) any { - if len(raw) == 0 { - return nil - } - return raw - } - completed := func(resultText string) { - enc(map[string]any{"jsonrpc": "2.0", "method": "turn/completed", "params": map[string]any{ - "success": true, "session_id": "fake-sess", "cost_usd": 0.01, - "result_text": resultText, - "usage": map[string]any{"input_tokens": 10, "output_tokens": 5}, - }}) - } - - // 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) - for scanner.Scan() { - var frame struct { - ID json.RawMessage `json:"id"` - Method string `json:"method"` - Result json.RawMessage `json:"result"` - Params struct { - OutputSchema json.RawMessage `json:"outputSchema"` - } `json:"params"` - } - if err := json.Unmarshal(scanner.Bytes(), &frame); err != nil { - continue - } - // The host's reply to our can_use_tool request (id, result, no method) - // finishes the approval turn, echoing the decision so the test can verify - // the round-trip reached the agent. - if frame.Method == "" && len(frame.Result) > 0 { - completed("decision=" + string(frame.Result)) - continue - } - switch frame.Method { - case "initialize": - initHadSchema = len(frame.Params.OutputSchema) > 0 - enc(map[string]any{"jsonrpc": "2.0", "id": id(frame.ID), "result": map[string]any{"ok": true}}) - enc(map[string]any{"jsonrpc": "2.0", "method": "session/init", "params": map[string]any{ - "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 "error-output": - _, _ = os.Stderr.WriteString("claude subprocess authentication detail\n") - enc(map[string]any{"jsonrpc": "2.0", "method": "turn/error", "params": map[string]any{ - "message": "Claude Code process exited with code 1", - }}) - 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": - // Complete with a structured_output payload, echoing whether the - // host actually transmitted the schema on initialize. - enc(map[string]any{"jsonrpc": "2.0", "method": "turn/completed", "params": map[string]any{ - "success": true, "session_id": "fake-sess", "cost_usd": 0.01, "subtype": "success", - "usage": map[string]any{"input_tokens": 10, "output_tokens": 5}, - "structured_output": map[string]any{ - "company_name": "Anthropic", - "founded_year": 2021, - "received_schema": initHadSchema, - }, - }}) - default: - enc(map[string]any{"jsonrpc": "2.0", "method": "message/tool_use", "params": map[string]any{ - "tool": "Read", "id": "t1", "input": map[string]any{"file_path": "/x"}, - }}) - completed("hi from fake") - } - case "interrupt": - if marker != "" { - _ = os.WriteFile(marker, []byte("interrupted"), 0o644) - } - enc(map[string]any{"jsonrpc": "2.0", "id": id(frame.ID), "result": map[string]any{}}) - case "shutdown": - enc(map[string]any{"jsonrpc": "2.0", "id": id(frame.ID), "result": map[string]any{}}) - os.Exit(0) - } - } -} - -func withFakeAgentProcess(t *testing.T) { - t.Helper() - withFakeAgentProcessEnv(t, map[string]string{fakeServerEnv: "1"}) -} - -func withFakeAgentProcessEnv(t *testing.T, env map[string]string) { - t.Helper() - self, err := os.Executable() - require.NoError(t, err) - - orig := newAgentProcess - newAgentProcess = func(*Provider) (*exec.Process, error) { - return exec.NewExec(self). - WithStdioPipe(). - WithEnv(env), nil - } - t.Cleanup(func() { newAgentProcess = orig }) -} - func TestProvider_StreamLifecycle(t *testing.T) { withFakeAgentProcess(t) @@ -219,6 +49,18 @@ func TestProvider_StreamLifecycle(t *testing.T) { assert.Equal(t, 5, result.Usage.OutputTokens) } +func TestAgentProcessEnvHonoursAPIURL(t *testing.T) { + got := agentProcessEnv(ai.Config{ + APIURL: "http://127.0.0.1:4010", APIKey: "captain-mock", + }, []string{"CLAUDECODE=1", "CLAUDE_CODE_ENTRYPOINT=cli"}) + assert.Equal(t, "http://127.0.0.1:4010", got["ANTHROPIC_BASE_URL"]) + assert.Equal(t, "captain-mock", got["ANTHROPIC_API_KEY"]) + assert.Equal(t, "captain-mock", got["ANTHROPIC_AUTH_TOKEN"]) + assert.Equal(t, "1", got["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"]) + assert.Empty(t, got["CLAUDECODE"]) + assert.Empty(t, got["CLAUDE_CODE_ENTRYPOINT"]) +} + func TestProvider_ExecuteCoalesce(t *testing.T) { withFakeAgentProcess(t) diff --git a/pkg/ai/provider/claudeagent/runner.go b/pkg/ai/provider/claudeagent/runner.go index 0578d250..aee6ffd2 100644 --- a/pkg/ai/provider/claudeagent/runner.go +++ b/pkg/ai/provider/claudeagent/runner.go @@ -17,6 +17,9 @@ import ( //go:embed agent.ts var agentTS string +//go:embed protocol.ts +var protocolTS string + //go:embed package.json var agentPackageJSON string @@ -38,6 +41,9 @@ func prepareAgentDir() (string, error) { if err := writeIfChanged(filepath.Join(agentDir, "agent.ts"), agentTS); err != nil { return "", err } + if err := writeIfChanged(filepath.Join(agentDir, "protocol.ts"), protocolTS); err != nil { + return "", err + } if err := writeIfChanged(filepath.Join(agentDir, "package.json"), agentPackageJSON); err != nil { return "", err } diff --git a/pkg/ai/provider/claudeagent/runtime_status.go b/pkg/ai/provider/claudeagent/runtime_status.go new file mode 100644 index 00000000..89ab75f8 --- /dev/null +++ b/pkg/ai/provider/claudeagent/runtime_status.go @@ -0,0 +1,72 @@ +package claudeagent + +import ( + "encoding/json" + "os" + "os/exec" + "path/filepath" + + "github.com/flanksource/captain/pkg/ai" +) + +type runtimeProbeOptions struct { + cacheDir string + lookPath func(string) (string, error) + readFile func(string) ([]byte, error) + stat func(string) (os.FileInfo, error) +} + +// ProbeRuntime mirrors newAgentProcess without materialising files or +// installing dependencies during a readiness check. +func ProbeRuntime() ai.RuntimeStatus { + cacheDir, err := os.UserCacheDir() + if err != nil { + return ai.RuntimeStatus{Error: err.Error()} + } + return probeRuntimeStatus(runtimeProbeOptions{ + cacheDir: cacheDir, + lookPath: exec.LookPath, + readFile: os.ReadFile, + stat: os.Stat, + }) +} + +func probeRuntimeStatus(options runtimeProbeOptions) ai.RuntimeStatus { + agentDir := filepath.Join(options.cacheDir, "captain", "claude-agent") + current, err := dependenciesCurrent(options.readFile, agentDir) + if err != nil { + return ai.RuntimeStatus{Error: err.Error()} + } + if !current { + if npm, pathErr := options.lookPath("npm"); pathErr == nil { + return ai.RuntimeStatus{Provisioner: npm} + } + return ai.RuntimeStatus{DependencyMissing: "npm"} + } + localTsx := filepath.Join(agentDir, "node_modules", ".bin", "tsx") + if _, statErr := options.stat(localTsx); statErr == nil { + return ai.RuntimeStatus{Binary: localTsx} + } + if tsx, pathErr := options.lookPath("tsx"); pathErr == nil { + return ai.RuntimeStatus{Binary: tsx} + } + return ai.RuntimeStatus{DependencyMissing: "tsx"} +} + +func dependenciesCurrent(readFile func(string) ([]byte, error), agentDir string) (bool, error) { + required, err := requiredSDKVersion() + if err != nil { + return false, err + } + data, err := readFile(filepath.Join(agentDir, "node_modules", "@anthropic-ai", "claude-agent-sdk", "package.json")) + if err != nil { + return false, nil + } + var installed struct { + Version string `json:"version"` + } + if json.Unmarshal(data, &installed) != nil { + return false, nil + } + return installed.Version == required, nil +} diff --git a/pkg/ai/provider/claudeagent/runtime_status_test.go b/pkg/ai/provider/claudeagent/runtime_status_test.go new file mode 100644 index 00000000..a44230bb --- /dev/null +++ b/pkg/ai/provider/claudeagent/runtime_status_test.go @@ -0,0 +1,81 @@ +package claudeagent + +import ( + "fmt" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/ai" +) + +type runtimeStatusCase struct { + installed bool + localTsx bool + paths map[string]string + wantBinary string + wantProvisioner string + wantDependency string +} + +var _ = Describe("Claude Agent runtime prerequisites", func() { + DescribeTable("reports how the runtime can become ready", + func(test runtimeStatusCase) { + cacheDir := GinkgoT().TempDir() + agentDir := filepath.Join(cacheDir, "captain", "claude-agent") + if test.installed { + writeInstalledSDK(agentDir) + } + localTsx := filepath.Join(agentDir, "node_modules", ".bin", "tsx") + if test.localTsx { + Expect(os.MkdirAll(filepath.Dir(localTsx), 0o755)).To(Succeed()) + Expect(os.WriteFile(localTsx, []byte("tsx"), 0o755)).To(Succeed()) + } + got := probeRuntimeStatus(runtimeProbeOptions{ + cacheDir: cacheDir, + lookPath: func(binary string) (string, error) { + if path := test.paths[binary]; path != "" { + return path, nil + } + return "", os.ErrNotExist + }, + readFile: os.ReadFile, + stat: os.Stat, + }) + wantBinary := test.wantBinary + if wantBinary == "local" { + wantBinary = localTsx + } + Expect(got).To(Equal(ai.RuntimeStatus{ + Binary: wantBinary, + Provisioner: test.wantProvisioner, + DependencyMissing: test.wantDependency, + })) + }, + Entry("uses the provisioner for a cold cache", runtimeStatusCase{ + paths: map[string]string{"npm": "/bin/npm"}, wantProvisioner: "/bin/npm", + }), + Entry("reports npm for a cold cache without a provisioner", runtimeStatusCase{ + wantDependency: "npm", + }), + Entry("uses the provisioned tsx", runtimeStatusCase{ + installed: true, localTsx: true, wantBinary: "local", + }), + Entry("uses a global tsx when the provisioned executable is missing", runtimeStatusCase{ + installed: true, paths: map[string]string{"tsx": "/bin/tsx"}, wantBinary: "/bin/tsx", + }), + Entry("reports tsx when installed dependencies are incomplete", runtimeStatusCase{ + installed: true, wantDependency: "tsx", + }), + ) +}) + +func writeInstalledSDK(agentDir string) { + version, err := requiredSDKVersion() + Expect(err).NotTo(HaveOccurred()) + manifest := filepath.Join(agentDir, "node_modules", "@anthropic-ai", "claude-agent-sdk", "package.json") + Expect(os.MkdirAll(filepath.Dir(manifest), 0o755)).To(Succeed()) + Expect(os.WriteFile(manifest, []byte(fmt.Sprintf(`{"version":%q}`, version)), 0o644)).To(Succeed()) +} diff --git a/pkg/ai/provider/claudeagent/turn.go b/pkg/ai/provider/claudeagent/turn.go index 515c3d57..8afcf0e1 100644 --- a/pkg/ai/provider/claudeagent/turn.go +++ b/pkg/ai/provider/claudeagent/turn.go @@ -8,6 +8,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "time" "github.com/flanksource/captain/pkg/ai" @@ -37,6 +38,8 @@ type turnState struct { promptMu sync.Mutex pending int ended bool + + interrupting atomic.Bool } type promptParams struct { @@ -158,6 +161,9 @@ func (p *Provider) onNotification(method string, params json.RawMessage) { return } + if ts.interrupting.Load() && (method == notifyTurnDone || method == notifyTurnError) { + ok = false + } if ok { select { case ts.inbox <- ev: @@ -321,7 +327,15 @@ func (p *Provider) Interrupt(ctx context.Context) error { if p.rpc == nil { return fmt.Errorf("claude-agent: provider not started") } + p.activeMu.Lock() + ts := p.active + p.activeMu.Unlock() + if ts == nil { + return fmt.Errorf("claude-agent: no active turn to interrupt") + } + ts.interrupting.Store(true) if _, err := p.rpc.Call(ctx, methodInterrupt, nil); err != nil { + ts.interrupting.Store(false) return fmt.Errorf("claude-agent interrupt failed: %w", err) } return nil diff --git a/pkg/ai/provider/cmux/sessionstats.go b/pkg/ai/provider/cmux/sessionstats.go index 047d5ef6..50f50612 100644 --- a/pkg/ai/provider/cmux/sessionstats.go +++ b/pkg/ai/provider/cmux/sessionstats.go @@ -12,6 +12,7 @@ import ( "github.com/flanksource/captain/pkg/ai/history" "github.com/flanksource/captain/pkg/ai/pricing" + "github.com/flanksource/captain/pkg/api" ) // High-level agent states surfaced to the dashboard, derived from the last @@ -135,6 +136,11 @@ type SessionStats struct { // Error is the API/network failure reason when State == "error" — the synthetic // "API Error: …" message Claude Code records when a request fails after retries. Error string `json:"error,omitempty"` + + // responses deduplicates the per-content-block lines one API response is + // written across. A cmux session log carries no result record, so the totals + // have to be reconstructed from each response; see api.ResponseSet. + responses api.ResponseSet } // sessionLogLine is the subset of a Claude session-log entry needed for stats: @@ -145,6 +151,9 @@ type sessionLogLine struct { Subtype string `json:"subtype"` Timestamp string `json:"timestamp"` Message struct { + // ID identifies the API response. Several lines share it when a response + // spans multiple content blocks, each repeating the same usage. + ID string `json:"id"` Model string `json:"model"` Usage struct { InputTokens int `json:"input_tokens"` @@ -173,7 +182,14 @@ func (l sessionLogLine) isCompaction() bool { // applyUsage folds one assistant entry's usage into the running totals and snaps // the live context window to this turn's prompt size (input + cache). +// +// Lines repeating a response already counted are skipped: they carry the same +// usage object, so folding each one in would multiply the session's tokens (and +// its cost) by the number of content blocks the response happened to contain. func (s *SessionStats) applyUsage(l sessionLogLine) { + if !s.responses.First(l.Message.ID) { + return + } u := l.Message.Usage s.InputTokens += u.InputTokens s.OutputTokens += u.OutputTokens diff --git a/pkg/ai/provider/cmux/sessionstats_test.go b/pkg/ai/provider/cmux/sessionstats_test.go index 544f0b3e..f9af0618 100644 --- a/pkg/ai/provider/cmux/sessionstats_test.go +++ b/pkg/ai/provider/cmux/sessionstats_test.go @@ -19,6 +19,62 @@ func assistantLine(ts, model string, in, out, cacheRead, cacheCreate int) string ) } +// assistantLineWithID is assistantLine plus the response id several +// content-block lines of one API response share. +func assistantLineWithID(ts, id, model string, in, out, cacheRead, cacheCreate int) string { + return fmt.Sprintf( + `{"type":"assistant","timestamp":%q,"message":{"id":%q,"model":%q,"usage":{"input_tokens":%d,"output_tokens":%d,"cache_read_input_tokens":%d,"cache_creation_input_tokens":%d},"content":[{"type":"text","text":"hi"}]}}`, + ts, id, model, in, out, cacheRead, cacheCreate, + ) +} + +// One API response spanning thinking, text and tool_use is written as three +// lines that each repeat the whole usage object. Counting them separately +// triples the session's tokens and its cost. +func TestComputeSessionStatsCountsOneResponseOnce(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + writeSessionLog(t, path, + assistantLineWithID("2026-06-23T10:00:00Z", "msg_a", "claude-opus-4-8", 100, 20, 5, 50), + assistantLineWithID("2026-06-23T10:00:00Z", "msg_a", "claude-opus-4-8", 100, 20, 5, 50), + assistantLineWithID("2026-06-23T10:00:00Z", "msg_a", "claude-opus-4-8", 100, 20, 5, 50), + assistantLineWithID("2026-06-23T10:00:30Z", "msg_b", "claude-opus-4-8", 200, 40, 7, 0), + ) + + stats, err := computeSessionStats(path) + if err != nil { + t.Fatalf("computeSessionStats() error = %v", err) + } + if stats.InputTokens != 300 || stats.OutputTokens != 60 { + t.Fatalf("tokens = in:%d out:%d, want in:300 out:60 for two responses", + stats.InputTokens, stats.OutputTokens) + } + if stats.CacheReadTokens != 12 || stats.CacheCreationTokens != 50 { + t.Fatalf("cache tokens = read:%d create:%d, want read:12 create:50", + stats.CacheReadTokens, stats.CacheCreationTokens) + } + if stats.Turns != 2 { + t.Fatalf("Turns = %d, want 2 responses rather than 4 lines", stats.Turns) + } +} + +// Lines without a response id cannot be correlated, so each must still count — +// dropping them would under-report every older session log. +func TestComputeSessionStatsCountsEveryUnidentifiedLine(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + writeSessionLog(t, path, + assistantLine("2026-06-23T10:00:00Z", "claude-opus-4-8", 100, 20, 0, 0), + assistantLine("2026-06-23T10:00:30Z", "claude-opus-4-8", 100, 20, 0, 0), + ) + + stats, err := computeSessionStats(path) + if err != nil { + t.Fatalf("computeSessionStats() error = %v", err) + } + if stats.InputTokens != 200 || stats.Turns != 2 { + t.Fatalf("tokens = in:%d turns:%d, want in:200 turns:2", stats.InputTokens, stats.Turns) + } +} + func TestComputeSessionStatsAggregatesUsage(t *testing.T) { path := filepath.Join(t.TempDir(), "s.jsonl") // Two assistant turns 30s apart plus a non-assistant line that must be ignored diff --git a/pkg/ai/provider/codex_appserver.go b/pkg/ai/provider/codex_appserver.go index f474c822..e474adaa 100644 --- a/pkg/ai/provider/codex_appserver.go +++ b/pkg/ai/provider/codex_appserver.go @@ -5,12 +5,14 @@ import ( "encoding/json" "fmt" osexec "os/exec" - "strings" "sync" "time" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/callertools" "github.com/flanksource/captain/pkg/ai/provider/jsonrpc" + aitools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" "github.com/flanksource/clicky/exec" "github.com/flanksource/commons/logger" ) @@ -27,6 +29,7 @@ var log = logger.GetLogger("ai") // unit-testable mapAppServerNotification. type CodexAppServer struct { model string + cfg ai.Config turnMu sync.Mutex // serializes turns; held by ExecuteStream, freed by its driver @@ -36,6 +39,10 @@ type CodexAppServer struct { rpcDone chan struct{} // closed by the rpc Run goroutine when the child exits active *turnState threadID string + + callerToolsMu sync.Mutex + callerToolsRuntime *callertools.Runtime + callerTools *api.CallerToolEndpoint } const ( @@ -45,16 +52,26 @@ const ( // NewCodexAppServer builds a codex app-server provider. The supervised process // is started lazily on the first ExecuteStream. -func NewCodexAppServer(model string) (*CodexAppServer, error) { +func NewCodexAppServer(cfg ai.Config) (*CodexAppServer, error) { + model := cfg.Model.Name if model == "" { model = CodexCLIDefaultModel } model = ai.NormalizeModelForBackend(ai.BackendCodexAgent, model) - return &CodexAppServer{model: model}, nil + provider := &CodexAppServer{model: model, cfg: cfg} + if cfg.CallerTools != nil { + endpoint := *cfg.CallerTools + endpoint.Headers = cloneStringMap(cfg.CallerTools.Headers) + provider.callerTools = &endpoint + } + return provider, nil } -func (c *CodexAppServer) GetModel() string { return c.model } -func (c *CodexAppServer) GetBackend() ai.Backend { return ai.BackendCodexAgent } +func (c *CodexAppServer) GetModel() string { return c.model } +func (c *CodexAppServer) GetBackend() ai.Backend { return ai.BackendCodexAgent } +func (c *CodexAppServer) SupportsCallerTools() bool { return true } + +var _ api.ToolCapableProvider = (*CodexAppServer)(nil) // 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 @@ -107,6 +124,9 @@ func (c *CodexAppServer) ExecuteStream(ctx context.Context, req ai.Request) (<-c return nil, err } + if err := c.prepareCallerTools(req); err != nil { + return nil, err + } c.turnMu.Lock() if err := c.ensureStarted(ctx); err != nil { c.turnMu.Unlock() @@ -120,7 +140,7 @@ func (c *CodexAppServer) ExecuteStream(ctx context.Context, req ai.Request) (<-c ch: make(chan ai.Event, 16), usage: &ai.Usage{}, model: c.model, - streamed: map[string]bool{}, + streamed: map[string]string{}, toolOutput: map[string]string{}, terminal: make(chan struct{}), started: make(chan struct{}), @@ -193,7 +213,7 @@ 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{ + sup := newCodexAppServerProcess(c.cfg).WithStdioPipe().Supervise(exec.SuperviseOptions{ // No restart: a crash surfaces as EventError, never a silent retry. RestartPolicy: exec.RestartNo, OnStarted: func(p *exec.Process) { @@ -236,13 +256,6 @@ 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 { @@ -280,7 +293,7 @@ func (c *CodexAppServer) startThread(ctx context.Context, req ai.Request) (strin return threadID, nil } if req.SessionID != "" { - raw, err := rpc.Call(ctx, "thread/resume", buildResumeParams(req)) + raw, err := rpc.Call(ctx, "thread/resume", buildResumeParams(req, c.callerTools)) if err != nil { return "", err } @@ -288,7 +301,7 @@ func (c *CodexAppServer) startThread(ctx context.Context, req ai.Request) (strin c.rememberThread(threadID) return threadID, nil } - raw, err := rpc.Call(ctx, "thread/start", buildThreadStartParams(c.model, req)) + raw, err := rpc.Call(ctx, "thread/start", buildThreadStartParams(c.model, req, c.callerTools)) if err != nil { return "", err } @@ -358,9 +371,58 @@ func (c *CodexAppServer) Interrupt(ctx context.Context) error { func (c *CodexAppServer) Close() error { c.teardown(true) + if c.callerToolsRuntime != nil { + return c.callerToolsRuntime.Close() + } return nil } +func (c *CodexAppServer) prepareCallerTools(req ai.Request) error { + c.callerToolsMu.Lock() + defer c.callerToolsMu.Unlock() + if c.callerTools != nil { + if req.Permissions.MCP.Disabled { + return fmt.Errorf("codex app-server: caller tools require MCP but MCP is disabled") + } + return c.callerTools.Validate() + } + if len(c.cfg.Tools) == 0 { + return nil + } + definitions, err := aitools.ResolveDefinitions(c.cfg.Tools, req.ToolPreferences) + if err != nil { + return fmt.Errorf("codex app-server caller tools: %w", err) + } + if len(definitions) == 0 { + return nil + } + if req.Permissions.MCP.Disabled { + return fmt.Errorf("codex app-server: caller tools require MCP but MCP is disabled") + } + runtime, err := callertools.New(callertools.Options{ + Definitions: definitions, CanUseTool: c.cfg.CanUseTool, + SessionID: firstNonEmpty(c.cfg.CaptainSessionID, req.SessionID, c.cfg.SessionID), + }) + if err != nil { + return fmt.Errorf("start codex app-server caller tools: %w", err) + } + endpoint := runtime.Endpoint() + c.callerToolsRuntime = runtime + c.callerTools = &endpoint + return nil +} + +func cloneStringMap(values map[string]string) map[string]string { + if values == nil { + return nil + } + cloned := make(map[string]string, len(values)) + for key, value := range values { + cloned[key] = value + } + return cloned +} + // 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. @@ -372,8 +434,9 @@ func (c *CodexAppServer) handleNotification(method string, params json.RawMessag ctx := appServerEventContext{Model: ts.model, Usage: ts.usage} switch method { case "item/agentMessage/delta": - if id := parseAppServerNotif(params).ItemID; id != "" { - ts.streamed[id] = true + notification := parseAppServerNotif(params) + if notification.ItemID != "" { + ts.streamed[notification.ItemID] += notification.Delta } if len(ts.outputSchema) > 0 { return @@ -394,7 +457,15 @@ func (c *CodexAppServer) handleNotification(method string, params json.RawMessag return } } - if appServerStreamedAgentMessage(params, ts.streamed) { + remainder, streamed, err := appServerAgentMessageRemainder(params, ts.streamed) + if err != nil { + ts.send(ai.Event{Kind: ai.EventError, Error: err.Error(), Model: ts.model}) + return + } + if streamed { + if remainder != "" { + ts.send(ai.Event{Kind: ai.EventText, Text: remainder, Model: ts.model}) + } return } if it := parseAppServerNotif(params).Item; it != nil { @@ -414,22 +485,6 @@ func (c *CodexAppServer) handleNotification(method string, params json.RawMessag } } -// handleApproval auto-approves server→client approval requests, mirroring the -// `--dangerously-bypass-approvals` default of the exec path. Decision shapes -// differ per method (see the *ApprovalResponse schemas). -func (c *CodexAppServer) handleApproval(method string, _ json.RawMessage) (any, *jsonrpc.RPCError) { - switch method { - case "item/commandExecution/requestApproval", "item/fileChange/requestApproval": - return map[string]string{"decision": "accept"}, nil - case "item/permissions/requestApproval": - return map[string]any{"permissions": map[string]any{}, "scope": "turn"}, nil - case "item/tool/requestUserInput": - return map[string]any{}, nil - default: // execCommandApproval, applyPatchApproval, unknown - return map[string]string{"decision": "approved"}, nil - } -} - // --- turn state ------------------------------------------------------------ // turnState is the routing target for one turn's server notifications. send and diff --git a/pkg/ai/provider/codex_appserver_approval.go b/pkg/ai/provider/codex_appserver_approval.go new file mode 100644 index 00000000..ce4b293d --- /dev/null +++ b/pkg/ai/provider/codex_appserver_approval.go @@ -0,0 +1,22 @@ +package provider + +import ( + "encoding/json" + + "github.com/flanksource/captain/pkg/ai/provider/jsonrpc" +) + +// handleApproval auto-approves server-to-client approval requests, mirroring +// the bypass-permissions default of the exec path. +func (c *CodexAppServer) handleApproval(method string, _ json.RawMessage) (any, *jsonrpc.RPCError) { + switch method { + case "item/commandExecution/requestApproval", "item/fileChange/requestApproval": + return map[string]string{"decision": "accept"}, nil + case "item/permissions/requestApproval": + return map[string]any{"permissions": map[string]any{}, "scope": "turn"}, nil + case "item/tool/requestUserInput": + return map[string]any{}, nil + default: + return map[string]string{"decision": "approved"}, nil + } +} diff --git a/pkg/ai/provider/codex_appserver_lifecycle_ginkgo_test.go b/pkg/ai/provider/codex_appserver_lifecycle_ginkgo_test.go index a5c0146c..326e309e 100644 --- a/pkg/ai/provider/codex_appserver_lifecycle_ginkgo_test.go +++ b/pkg/ai/provider/codex_appserver_lifecycle_ginkgo_test.go @@ -77,6 +77,18 @@ var _ = Describe("Codex app-server tool lifecycle", func() { }) var _ = Describe("Codex app-server turn control", func() { + It("does not emit a successful result for an interrupted turn", func() { + client, turn := activeGinkgoTurn() + + client.handleNotification("turn/completed", json.RawMessage(`{ + "threadId":"thread-1", + "turn":{"id":"turn-1","status":"interrupted"} + }`)) + + Expect(drainEvents(turn)).To(BeEmpty()) + Eventually(turn.terminal).Should(BeClosed()) + }) + It("waits for thread and turn identifiers before interrupting", func() { turn := &turnState{terminal: make(chan struct{}), started: make(chan struct{})} go func() { @@ -122,13 +134,13 @@ var _ = Describe("Codex CLI attachments", func() { }) func activeGinkgoTurn() (*CodexAppServer, *turnState) { - client, err := NewCodexAppServer("gpt-5") + client, err := NewCodexAppServer(ai.Config{Model: api.Model{Name: "gpt-5"}}) Expect(err).NotTo(HaveOccurred()) turn := &turnState{ ch: make(chan ai.Event, 16), usage: &ai.Usage{}, model: "gpt-5", - streamed: map[string]bool{}, + streamed: map[string]string{}, toolOutput: map[string]string{}, terminal: make(chan struct{}), } diff --git a/pkg/ai/provider/codex_appserver_params_test.go b/pkg/ai/provider/codex_appserver_params_test.go new file mode 100644 index 00000000..73db9254 --- /dev/null +++ b/pkg/ai/provider/codex_appserver_params_test.go @@ -0,0 +1,124 @@ +package provider + +import ( + "testing" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/commons-db/shell" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestComposePrompt(t *testing.T) { + assert.Equal(t, "task", composePrompt(req(api.Prompt{User: "task"}))) + assert.Equal(t, "be brief\n\ntask", composePrompt(req(api.Prompt{User: "task", System: "be brief"}))) + assert.Equal(t, "task\n\ntail", composePrompt(req(api.Prompt{User: "task", AppendSystem: "tail"}))) + assert.Equal(t, "sys\n\ntask\n\ntail", + composePrompt(req(api.Prompt{User: "task", System: "sys", AppendSystem: "tail"}))) +} + +func TestBuildThreadStartParams_Safety(t *testing.T) { + tests := []struct { + name string + req ai.Request + wantSandbox string + wantApproval string + wantEphem bool + wantNoMCP bool + }{ + {name: "default is read-only on-request", req: req(api.Prompt{User: "p"}), wantSandbox: "read-only", wantApproval: "on-request"}, + { + name: "edit maps to workspace-write", + req: ai.Request{Prompt: api.Prompt{User: "p"}, Permissions: api.Permissions{Presets: []api.Preset{api.PresetEdit}}}, + wantSandbox: "workspace-write", wantApproval: "on-request", + }, + { + name: "explicit permission mode skips workspace-write default", + req: ai.Request{Prompt: api.Prompt{User: "p"}, Permissions: api.Permissions{Presets: []api.Preset{api.PresetEdit}, Mode: api.PermissionDefault}}, + wantSandbox: "read-only", wantApproval: "on-request", + }, + { + name: "bypass permissions maps to danger-full-access never", + req: ai.Request{Prompt: api.Prompt{User: "p"}, Permissions: api.Permissions{Mode: api.PermissionBypass}}, + wantSandbox: "danger-full-access", wantApproval: "never", + }, + { + name: "no-memory sets ephemeral", + req: ai.Request{Prompt: api.Prompt{User: "p"}, Memory: api.Memory{SkipMemory: true}}, + wantSandbox: "read-only", wantApproval: "on-request", wantEphem: true, + }, + { + name: "bare sets ephemeral", + req: ai.Request{Prompt: api.Prompt{User: "p"}, Memory: api.Memory{Bare: true}}, + wantSandbox: "read-only", wantApproval: "on-request", wantEphem: true, + }, + { + name: "no-mcp sets empty mcp_servers override", + req: ai.Request{Prompt: api.Prompt{User: "p"}, Permissions: api.Permissions{MCP: api.MCP{Disabled: true}}}, + wantSandbox: "read-only", wantApproval: "on-request", wantNoMCP: true, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + p := buildThreadStartParams("gpt-5", tc.req, nil) + assert.Equal(t, tc.wantSandbox, p["sandbox"]) + assert.Equal(t, tc.wantApproval, p["approvalPolicy"]) + _, hasEphem := p["ephemeral"] + assert.Equal(t, tc.wantEphem, hasEphem) + cfg, hasCfg := p["config"].(map[string]any) + assert.Equal(t, tc.wantNoMCP, hasCfg) + if tc.wantNoMCP { + assert.Equal(t, map[string]any{}, cfg["mcp_servers"]) + } + }) + } +} + +func TestBuildThreadStartParams_CwdAndModel(t *testing.T) { + p := buildThreadStartParams("gpt-5", ai.Request{ + Prompt: api.Prompt{User: "p"}, Setup: &shell.Setup{Cwd: "/repo"}, + }, nil) + assert.Equal(t, "/repo", p["cwd"]) + assert.Equal(t, "gpt-5", p["model"]) + noModel := buildThreadStartParams("", req(api.Prompt{User: "p"}), nil) + _, hasModel := noModel["model"] + assert.False(t, hasModel, "empty model must be omitted") +} + +func TestBuildResumeParams(t *testing.T) { + p := buildResumeParams(ai.Request{SessionID: "thread-9", Setup: &shell.Setup{Cwd: "/repo"}}, nil) + assert.Equal(t, "thread-9", p["threadId"]) + assert.Equal(t, "/repo", p["cwd"]) +} + +func TestHandleApproval_AutoApproves(t *testing.T) { + c, err := NewCodexAppServer(ai.Config{Model: api.Model{Name: "m"}}) + require.NoError(t, err) + tests := []struct { + method string + key string + want any + }{ + {"execCommandApproval", "decision", "approved"}, + {"applyPatchApproval", "decision", "approved"}, + {"item/commandExecution/requestApproval", "decision", "accept"}, + {"item/fileChange/requestApproval", "decision", "accept"}, + {"some/unknown/approval", "decision", "approved"}, + } + for _, tc := range tests { + t.Run(tc.method, func(t *testing.T) { + res, rpcErr := c.handleApproval(tc.method, nil) + assert.Nil(t, rpcErr) + m, ok := res.(map[string]string) + require.True(t, ok, "decision approvals return a string map") + assert.Equal(t, tc.want, m[tc.key]) + }) + } + res, rpcErr := c.handleApproval("item/permissions/requestApproval", nil) + assert.Nil(t, rpcErr) + perm, ok := res.(map[string]any) + require.True(t, ok) + assert.Equal(t, "turn", perm["scope"]) + assert.NotNil(t, perm["permissions"]) +} diff --git a/pkg/ai/provider/codex_appserver_process.go b/pkg/ai/provider/codex_appserver_process.go new file mode 100644 index 00000000..5b489137 --- /dev/null +++ b/pkg/ai/provider/codex_appserver_process.go @@ -0,0 +1,28 @@ +package provider + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/clicky/exec" +) + +func newCodexAppServerProcess(cfg ai.Config) *exec.Process { + args := []string{"app-server"} + if cfg.APIURL != "" { + args = append(args, codexProviderOverride(cfg.APIURL)...) + } + process := exec.NewExec("codex", args...) + if cfg.APIKey != "" { + process.WithEnv(map[string]string{"OPENAI_API_KEY": cfg.APIKey}) + } + return process +} + +func appServerProcessError(err error, stderr string) error { + if detail := strings.TrimSpace(stderr); detail != "" { + return fmt.Errorf("%w: %s", err, detail) + } + return err +} diff --git a/pkg/ai/provider/codex_appserver_protocol.go b/pkg/ai/provider/codex_appserver_protocol.go index b706d242..edb190b1 100644 --- a/pkg/ai/provider/codex_appserver_protocol.go +++ b/pkg/ai/provider/codex_appserver_protocol.go @@ -3,6 +3,7 @@ package provider import ( "encoding/json" "fmt" + "strings" "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/ai/history" @@ -64,7 +65,9 @@ type appServerErrorBody struct { } type appServerRef struct { - ID string `json:"id"` + ID string `json:"id"` + Status string `json:"status"` + Error *appServerErrorBody `json:"error"` } type appServerTokenUsage struct { @@ -159,6 +162,18 @@ func mapAppServerNotification(method string, params json.RawMessage, ctx appServ return ai.Event{}, false case "turn/completed": + if n.Turn != nil { + switch n.Turn.Status { + case "interrupted": + return ai.Event{}, false + case "failed": + message := "codex turn failed" + if n.Turn.Error != nil { + message = firstNonEmpty(n.Turn.Error.Message, n.Turn.Error.AdditionalDetails, message) + } + return ai.Event{Kind: ai.EventError, Error: extractCodexErrorText(message), SessionID: n.threadID(), Model: ctx.Model}, true + } + } 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 @@ -278,11 +293,19 @@ func appServerErrorIsFatal(method string, params json.RawMessage) bool { return !parseAppServerNotif(params).WillRetry } -// appServerStreamedAgentMessage reports whether an item/completed notification is -// an agent message whose text was already streamed via item/agentMessage/delta. -func appServerStreamedAgentMessage(params json.RawMessage, streamed map[string]bool) bool { +func appServerAgentMessageRemainder(params json.RawMessage, streamed map[string]string) (string, bool, error) { it := parseAppServerNotif(params).Item - return it != nil && it.Type == "agentMessage" && streamed[it.ID] + if it == nil || it.Type != "agentMessage" { + return "", false, nil + } + prefix, ok := streamed[it.ID] + if !ok { + return "", false, nil + } + if !strings.HasPrefix(it.Text, prefix) { + return "", true, fmt.Errorf("codex app-server completed agent message %q does not extend its streamed text", it.ID) + } + return strings.TrimPrefix(it.Text, prefix), true, nil } // --- request params -------------------------------------------------------- @@ -303,7 +326,7 @@ func composePrompt(req ai.Request) string { // (req.Memory.SkipUser/SkipProject/SkipHooks) have no first-class equivalent in // the versioned thread/start schema, so only ephemeral + an empty mcp_servers // override (the knobs the protocol exposes) are emitted. -func buildThreadStartParams(model string, req ai.Request) map[string]any { +func buildThreadStartParams(model string, req ai.Request, callerTools *api.CallerToolEndpoint) map[string]any { p := map[string]any{} if cwd := req.Cwd(); cwd != "" { p["cwd"] = cwd @@ -316,8 +339,8 @@ func buildThreadStartParams(model string, req ai.Request) map[string]any { if req.Memory.SkipMemory || req.Memory.Bare { p["ephemeral"] = true } - if req.Permissions.MCP.Disabled { - p["config"] = map[string]any{"mcp_servers": map[string]any{}} + if config := codexThreadConfig(req, callerTools); config != nil { + p["config"] = config } return p } @@ -367,10 +390,28 @@ func buildTurnStartParams(model string, req ai.Request, threadID string, outputS return p, nil } -func buildResumeParams(req ai.Request) map[string]any { +func buildResumeParams(req ai.Request, callerTools *api.CallerToolEndpoint) map[string]any { p := map[string]any{"threadId": req.SessionID} if cwd := req.Cwd(); cwd != "" { p["cwd"] = cwd } + if config := codexThreadConfig(req, callerTools); config != nil { + p["config"] = config + } return p } + +func codexThreadConfig(req ai.Request, callerTools *api.CallerToolEndpoint) map[string]any { + if req.Permissions.MCP.Disabled { + return map[string]any{"mcp_servers": map[string]any{}} + } + if callerTools == nil { + return nil + } + return map[string]any{"mcp_servers": map[string]any{ + callerTools.Name: map[string]any{ + "url": callerTools.URL, "http_headers": cloneStringMap(callerTools.Headers), + "required": true, "enabled": true, "default_tools_approval_mode": "approve", + }, + }} +} diff --git a/pkg/ai/provider/codex_appserver_test.go b/pkg/ai/provider/codex_appserver_test.go index 63056f81..5aa0a165 100644 --- a/pkg/ai/provider/codex_appserver_test.go +++ b/pkg/ai/provider/codex_appserver_test.go @@ -8,22 +8,37 @@ import ( "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/claude" - "github.com/flanksource/commons-db/shell" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestNewCodexAppServer_Defaults(t *testing.T) { - c, err := NewCodexAppServer("") + c, err := NewCodexAppServer(ai.Config{}) require.NoError(t, err) assert.Equal(t, CodexCLIDefaultModel, c.GetModel()) assert.Equal(t, ai.BackendCodexAgent, c.GetBackend()) - c2, err := NewCodexAppServer("gpt-5.4") + c2, err := NewCodexAppServer(ai.Config{Model: api.Model{Name: "gpt-5.4"}}) require.NoError(t, err) assert.Equal(t, "gpt-5.4", c2.GetModel()) } +func TestCodexAppServerProcessHonoursAPIURL(t *testing.T) { + process := newCodexAppServerProcess(ai.Config{ + APIURL: "http://127.0.0.1:4020/v1", APIKey: "captain-mock", + }) + assert.Equal(t, "codex", process.Cmd) + assert.Equal(t, []string{ + "app-server", + "-c", "model_provider=captain", + "-c", "model_providers.captain.name=captain", + "-c", "model_providers.captain.base_url=http://127.0.0.1:4020/v1", + "-c", "model_providers.captain.env_key=OPENAI_API_KEY", + "-c", "model_providers.captain.wire_api=responses", + }, process.Args) + assert.Equal(t, "captain-mock", process.Env["OPENAI_API_KEY"]) +} + 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") @@ -197,13 +212,13 @@ func drainEvents(ts *turnState) []ai.Event { // route to it, mirroring what ExecuteStream sets up. func activeTurn(t *testing.T, schema json.RawMessage) (*CodexAppServer, *turnState) { t.Helper() - c, err := NewCodexAppServer("gpt-5") + c, err := NewCodexAppServer(ai.Config{Model: api.Model{Name: "gpt-5"}}) require.NoError(t, err) ts := &turnState{ ch: make(chan ai.Event, 16), usage: &ai.Usage{}, model: "gpt-5", - streamed: map[string]bool{}, + streamed: map[string]string{}, toolOutput: map[string]string{}, terminal: make(chan struct{}), outputSchema: schema, @@ -292,6 +307,23 @@ func TestHandleNotification_TextModeStreamsDeltasAndDeduplicatesCompletedMessage assert.Empty(t, resultEvent(t, events).StructuredData) } +func TestHandleNotification_TextModeBackfillsUnstreamedCompletedSuffix(t *testing.T) { + c, ts := activeTurn(t, nil) + c.handleNotification("item/agentMessage/delta", + json.RawMessage(`{"itemId":"a1","delta":"plain "}`)) + c.handleNotification("item/completed", + json.RawMessage(`{"item":{"id":"a1","type":"agentMessage","text":"plain answer"}}`)) + + events := drainEvents(ts) + var text []string + for _, event := range events { + if event.Kind == ai.EventText { + text = append(text, event.Text) + } + } + assert.Equal(t, []string{"plain ", "answer"}, text) +} + // A text-mode turn (no schema) leaves the result's StructuredData empty. func TestHandleNotification_NoStructuredWithoutSchema(t *testing.T) { c, ts := activeTurn(t, nil) @@ -372,20 +404,24 @@ func TestAppServerErrorIsFatal(t *testing.T) { } } -func TestAppServerStreamedAgentMessage(t *testing.T) { - streamed := map[string]bool{"i1": true} +func TestAppServerAgentMessageRemainder(t *testing.T) { + streamed := map[string]string{"i1": "partial"} - assert.True(t, appServerStreamedAgentMessage( - json.RawMessage(`{"item":{"id":"i1","type":"agentMessage","text":"x"}}`), streamed), - "completed agent message whose deltas streamed should be deduped") + remainder, handled, err := appServerAgentMessageRemainder( + json.RawMessage(`{"item":{"id":"i1","type":"agentMessage","text":"partial result"}}`), streamed) + require.NoError(t, err) + assert.True(t, handled) + assert.Equal(t, " result", remainder) - assert.False(t, appServerStreamedAgentMessage( - json.RawMessage(`{"item":{"id":"i2","type":"agentMessage","text":"x"}}`), streamed), - "a different item id was not streamed") + _, handled, err = appServerAgentMessageRemainder( + json.RawMessage(`{"item":{"id":"i2","type":"agentMessage","text":"x"}}`), streamed) + require.NoError(t, err) + assert.False(t, handled) - assert.False(t, appServerStreamedAgentMessage( - json.RawMessage(`{"item":{"id":"i1","type":"commandExecution"}}`), streamed), - "non agent-message items are never deduped") + _, handled, err = appServerAgentMessageRemainder( + json.RawMessage(`{"item":{"id":"i1","type":"commandExecution"}}`), streamed) + require.NoError(t, err) + assert.False(t, handled) } func TestThreadID(t *testing.T) { @@ -397,14 +433,6 @@ func TestThreadID(t *testing.T) { assert.Equal(t, "", tid(`{}`)) } -func TestComposePrompt(t *testing.T) { - assert.Equal(t, "task", composePrompt(req(api.Prompt{User: "task"}))) - assert.Equal(t, "be brief\n\ntask", composePrompt(req(api.Prompt{User: "task", System: "be brief"}))) - assert.Equal(t, "task\n\ntail", composePrompt(req(api.Prompt{User: "task", AppendSystem: "tail"}))) - assert.Equal(t, "sys\n\ntask\n\ntail", - composePrompt(req(api.Prompt{User: "task", System: "sys", AppendSystem: "tail"}))) -} - // req builds an ai.Request carrying only the given prompt, keeping the nested // api.Spec literal out of the table tests above. func req(p api.Prompt) ai.Request { @@ -464,149 +492,3 @@ func TestBuildTurnStartParams_OutputSchema(t *testing.T) { assert.Equal(t, []any{"answer", "detail"}, decoded["required"]) assert.Equal(t, false, decoded["additionalProperties"]) } - -func TestBuildThreadStartParams_Safety(t *testing.T) { - tests := []struct { - name string - req ai.Request - wantSandbox string - wantApproval string - wantEphem bool - wantNoMCP bool - }{ - { - name: "default is read-only on-request", - req: req(api.Prompt{User: "p"}), - wantSandbox: "read-only", - wantApproval: "on-request", - }, - { - name: "edit maps to workspace-write", - req: ai.Request{ - Prompt: api.Prompt{User: "p"}, - Permissions: api.Permissions{Presets: []api.Preset{api.PresetEdit}}, - }, - wantSandbox: "workspace-write", - wantApproval: "on-request", - }, - { - name: "explicit permission mode skips workspace-write default", - req: ai.Request{ - Prompt: api.Prompt{User: "p"}, - Permissions: api.Permissions{Presets: []api.Preset{api.PresetEdit}, Mode: api.PermissionDefault}, - }, - wantSandbox: "read-only", - wantApproval: "on-request", - }, - { - name: "bypass permissions maps to danger-full-access never", - req: ai.Request{ - Prompt: api.Prompt{User: "p"}, - Permissions: api.Permissions{Mode: api.PermissionBypass}, - }, - wantSandbox: "danger-full-access", - wantApproval: "never", - }, - { - name: "no-memory sets ephemeral", - req: ai.Request{ - Prompt: api.Prompt{User: "p"}, - Memory: api.Memory{SkipMemory: true}, - }, - wantSandbox: "read-only", - wantApproval: "on-request", - wantEphem: true, - }, - { - name: "bare sets ephemeral", - req: ai.Request{ - Prompt: api.Prompt{User: "p"}, - Memory: api.Memory{Bare: true}, - }, - wantSandbox: "read-only", - wantApproval: "on-request", - wantEphem: true, - }, - { - name: "no-mcp sets empty mcp_servers override", - req: ai.Request{ - Prompt: api.Prompt{User: "p"}, - Permissions: api.Permissions{MCP: api.MCP{Disabled: true}}, - }, - wantSandbox: "read-only", - wantApproval: "on-request", - wantNoMCP: true, - }, - } - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - p := buildThreadStartParams("gpt-5", tc.req) - assert.Equal(t, tc.wantSandbox, p["sandbox"]) - assert.Equal(t, tc.wantApproval, p["approvalPolicy"]) - - _, hasEphem := p["ephemeral"] - assert.Equal(t, tc.wantEphem, hasEphem) - - cfg, hasCfg := p["config"].(map[string]any) - assert.Equal(t, tc.wantNoMCP, hasCfg) - if tc.wantNoMCP { - assert.Equal(t, map[string]any{}, cfg["mcp_servers"]) - } - }) - } -} - -func TestBuildThreadStartParams_CwdAndModel(t *testing.T) { - p := buildThreadStartParams("gpt-5", ai.Request{ - Prompt: api.Prompt{User: "p"}, - Setup: &shell.Setup{Cwd: "/repo"}, - }) - assert.Equal(t, "/repo", p["cwd"]) - assert.Equal(t, "gpt-5", p["model"]) - - noModel := buildThreadStartParams("", req(api.Prompt{User: "p"})) - _, hasModel := noModel["model"] - assert.False(t, hasModel, "empty model must be omitted") -} - -func TestBuildResumeParams(t *testing.T) { - p := buildResumeParams(ai.Request{ - SessionID: "thread-9", - Setup: &shell.Setup{Cwd: "/repo"}, - }) - assert.Equal(t, "thread-9", p["threadId"]) - assert.Equal(t, "/repo", p["cwd"]) -} - -func TestHandleApproval_AutoApproves(t *testing.T) { - c, err := NewCodexAppServer("m") - require.NoError(t, err) - - tests := []struct { - method string - key string - want any - }{ - {"execCommandApproval", "decision", "approved"}, - {"applyPatchApproval", "decision", "approved"}, - {"item/commandExecution/requestApproval", "decision", "accept"}, - {"item/fileChange/requestApproval", "decision", "accept"}, - {"some/unknown/approval", "decision", "approved"}, - } - for _, tc := range tests { - t.Run(tc.method, func(t *testing.T) { - res, rpcErr := c.handleApproval(tc.method, nil) - assert.Nil(t, rpcErr) - m, ok := res.(map[string]string) - require.True(t, ok, "decision approvals return a string map") - assert.Equal(t, tc.want, m[tc.key]) - }) - } - - res, rpcErr := c.handleApproval("item/permissions/requestApproval", nil) - assert.Nil(t, rpcErr) - perm, ok := res.(map[string]any) - require.True(t, ok) - assert.Equal(t, "turn", perm["scope"]) - assert.NotNil(t, perm["permissions"]) -} diff --git a/pkg/ai/provider/codex_appserver_turn.go b/pkg/ai/provider/codex_appserver_turn.go index 0f5e0ce2..75ec20c6 100644 --- a/pkg/ai/provider/codex_appserver_turn.go +++ b/pkg/ai/provider/codex_appserver_turn.go @@ -13,7 +13,7 @@ type turnState struct { ch chan ai.Event usage *ai.Usage model string - streamed map[string]bool + streamed map[string]string toolOutput map[string]string outputSchema json.RawMessage diff --git a/pkg/ai/provider/genkit/approval.go b/pkg/ai/provider/genkit/approval.go index 3c3ffdb1..2575201c 100644 --- a/pkg/ai/provider/genkit/approval.go +++ b/pkg/ai/provider/genkit/approval.go @@ -22,7 +22,13 @@ func toolApprovalState(req ai.Request, response *gkai.ModelResponse) (*api.ToolA if err != nil { return nil, err } - state := &api.ToolApprovalState{Messages: append(messages, assistant), Calls: calls} + checkpoint, err := encodeToolApprovalCheckpoint(response) + if err != nil { + return nil, err + } + state := &api.ToolApprovalState{ + Messages: append(messages, assistant), Calls: calls, ProviderCheckpoint: checkpoint, + } if err := state.Validate(); err != nil { return nil, fmt.Errorf("genkit approval state: %w", err) } @@ -116,7 +122,7 @@ func prepareToolApprovalResume(resume *api.ToolApprovalResume) ([]*gkai.Message, if err := resume.Validate(); err != nil { return nil, nil, nil, err } - messages, err := conversationMessages(resume.State.Messages) + messages, err := decodeToolApprovalCheckpoint(resume.State.ProviderCheckpoint) if err != nil { return nil, nil, nil, err } @@ -152,8 +158,13 @@ func prepareToolApprovalResume(resume *api.ToolApprovalResume) ([]*gkai.Message, if reason == "" { reason = "tool call denied" } + output := map[string]any{"denied": true, "reason": reason} + if part.Metadata == nil { + part.Metadata = map[string]any{} + } + part.Metadata["pendingOutput"] = output responses = append(responses, gkai.NewToolResponsePart(&gkai.ToolResponse{ - Name: call.Request.Tool, Ref: call.Request.ToolCallID, Output: map[string]any{"denied": true, "reason": reason}, + Name: call.Request.Tool, Ref: call.Request.ToolCallID, Output: output, })) case api.ToolApprovalRespond: output, err := approvalResultOutput(decision.Result) diff --git a/pkg/ai/provider/genkit/approval_checkpoint.go b/pkg/ai/provider/genkit/approval_checkpoint.go new file mode 100644 index 00000000..e31e8cbb --- /dev/null +++ b/pkg/ai/provider/genkit/approval_checkpoint.go @@ -0,0 +1,153 @@ +package genkit + +import ( + "encoding/base64" + "encoding/json" + "fmt" + + "github.com/flanksource/captain/pkg/api" + + gkai "github.com/firebase/genkit/go/ai" +) + +const ( + genkitApprovalCheckpointCodec = "genkit-messages-json" + genkitApprovalCheckpointVersion = 1 + checkpointBytesKey = "$captainBytes" +) + +func encodeToolApprovalCheckpoint(response *gkai.ModelResponse) (*api.ProviderCheckpoint, error) { + if response == nil || response.Request == nil || response.Message == nil { + return nil, fmt.Errorf("genkit approval checkpoint requires the model request and response message") + } + messages := cloneCheckpointMessages(response.Request.Messages) + messages = append(messages, response.Message.Clone()) + encodeCheckpointMetadata(messages) + payload, err := json.Marshal(messages) + if err != nil { + return nil, fmt.Errorf("encode genkit approval checkpoint: %w", err) + } + return &api.ProviderCheckpoint{ + Codec: genkitApprovalCheckpointCodec, Version: genkitApprovalCheckpointVersion, Payload: payload, + }, nil +} + +func decodeToolApprovalCheckpoint(checkpoint *api.ProviderCheckpoint) ([]*gkai.Message, error) { + if checkpoint == nil { + return nil, fmt.Errorf("genkit approval checkpoint is missing") + } + if checkpoint.Codec != genkitApprovalCheckpointCodec || checkpoint.Version != genkitApprovalCheckpointVersion { + return nil, fmt.Errorf("unsupported genkit approval checkpoint %q version %d", checkpoint.Codec, checkpoint.Version) + } + var messages []*gkai.Message + if err := json.Unmarshal(checkpoint.Payload, &messages); err != nil { + return nil, fmt.Errorf("decode genkit approval checkpoint: %w", err) + } + if len(messages) == 0 { + return nil, fmt.Errorf("genkit approval checkpoint has no messages") + } + if err := decodeCheckpointMetadata(messages); err != nil { + return nil, err + } + return messages, nil +} + +func cloneCheckpointMessages(messages []*gkai.Message) []*gkai.Message { + cloned := make([]*gkai.Message, len(messages)) + for i, message := range messages { + cloned[i] = message.Clone() + } + return cloned +} + +func encodeCheckpointMetadata(messages []*gkai.Message) { + for _, message := range messages { + message.Metadata = encodeCheckpointMap(message.Metadata) + for _, part := range message.Content { + part.Metadata = encodeCheckpointMap(part.Metadata) + } + } +} + +func encodeCheckpointMap(values map[string]any) map[string]any { + for key, value := range values { + values[key] = encodeCheckpointValue(value) + } + return values +} + +func encodeCheckpointValue(value any) any { + switch typed := value.(type) { + case []byte: + return map[string]any{checkpointBytesKey: base64.StdEncoding.EncodeToString(typed)} + case map[string]any: + return encodeCheckpointMap(typed) + case []any: + for i := range typed { + typed[i] = encodeCheckpointValue(typed[i]) + } + return typed + default: + return value + } +} + +func decodeCheckpointMetadata(messages []*gkai.Message) error { + for _, message := range messages { + if err := decodeCheckpointMap(message.Metadata); err != nil { + return err + } + for _, part := range message.Content { + if err := decodeCheckpointMap(part.Metadata); err != nil { + return err + } + } + } + return nil +} + +func decodeCheckpointMap(values map[string]any) error { + for key, value := range values { + decoded, err := decodeCheckpointValue(value) + if err != nil { + return fmt.Errorf("decode genkit checkpoint metadata %q: %w", key, err) + } + values[key] = decoded + } + return nil +} + +func decodeCheckpointValue(value any) (any, error) { + switch typed := value.(type) { + case map[string]any: + if encoded, ok := typed[checkpointBytesKey]; ok { + if len(typed) != 1 { + return nil, fmt.Errorf("byte envelope has unexpected fields") + } + text, ok := encoded.(string) + if !ok { + return nil, fmt.Errorf("byte envelope payload is %T", encoded) + } + decoded, err := base64.StdEncoding.DecodeString(text) + if err != nil { + return nil, fmt.Errorf("invalid byte envelope: %w", err) + } + return decoded, nil + } + if err := decodeCheckpointMap(typed); err != nil { + return nil, err + } + return typed, nil + case []any: + for i := range typed { + decoded, err := decodeCheckpointValue(typed[i]) + if err != nil { + return nil, err + } + typed[i] = decoded + } + return typed, nil + default: + return value, nil + } +} diff --git a/pkg/ai/provider/genkit/mapping.go b/pkg/ai/provider/genkit/mapping.go index a7f017bf..33ca5b56 100644 --- a/pkg/ai/provider/genkit/mapping.go +++ b/pkg/ai/provider/genkit/mapping.go @@ -1,9 +1,7 @@ package genkit import ( - "encoding/json" "fmt" - "reflect" "sync" "time" @@ -75,51 +73,31 @@ func (c *toolEventCorrelation) observeRequest(request *gkai.ToolRequest) { c.pending = append(c.pending, request) } -func (c *toolEventCorrelation) begin(name string, input map[string]any) (string, error) { +func (c *toolEventCorrelation) begin(request *gkai.ToolRequest) (string, error) { + if request == nil || request.Name == "" { + return "", fmt.Errorf("genkit tool execution has no provider request name") + } + if request.Ref == "" { + return "", fmt.Errorf("genkit tool %q provider request has no call reference", request.Name) + } 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 { + for i, pending := range c.pending { + if pending.Ref != request.Ref { continue } - matchingName = append(matchingName, i) - if requestInput, ok := toolRequestInput(request.Input); ok && reflect.DeepEqual(requestInput, input) { - exact = append(exact, i) + if pending.Name != request.Name { + return "", fmt.Errorf("genkit tool request %q names %q, expected %q", request.Ref, request.Name, pending.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[:i], c.pending[i+1:]...) + c.started[request.Ref] = request.Name + return request.Ref, nil } - - 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) - } + return "", fmt.Errorf("genkit tool %q execution reference %q has no correlated provider request", request.Name, request.Ref) } func (c *toolEventCorrelation) finish(ref string) error { @@ -163,38 +141,6 @@ func (c *toolEventCorrelation) observeResponse(response *gkai.ToolResponse) erro 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 -// map[string]any shape captain's Event.Input expects. -func toInputMap(v any) map[string]any { - if v == nil { - return nil - } - if m, ok := v.(map[string]any); ok { - return m - } - b, err := json.Marshal(v) - if err != nil { - return nil - } - var m map[string]any - if err := json.Unmarshal(b, &m); err != nil { - return nil - } - return m -} - // 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 diff --git a/pkg/ai/provider/genkit/options.go b/pkg/ai/provider/genkit/options.go index 26564af4..13f583f8 100644 --- a/pkg/ai/provider/genkit/options.go +++ b/pkg/ai/provider/genkit/options.go @@ -1,6 +1,7 @@ package genkit import ( + "context" "encoding/base64" "encoding/json" "fmt" @@ -46,7 +47,10 @@ func generateOptions(p *Provider, req ai.Request, stream gkai.ModelStreamCallbac if err := req.ValidateRequestMode(); err != nil { return nil, err } - opts := []gkai.GenerateOption{gkai.WithModelName(p.modelRef)} + opts := []gkai.GenerateOption{ + gkai.WithModelName(p.modelRef), + gkai.WithUse(gkai.MiddlewareFunc(captureGenkitRequests)), + } toolOptions, err := p.toolOptions(req.ToolPreferences, emit) if err != nil { return nil, err @@ -132,6 +136,26 @@ func generateOptions(p *Provider, req ai.Request, stream gkai.ModelStreamCallbac return opts, nil } +type genkitToolRequestContextKey struct{} + +func captureGenkitRequests(context.Context) (*gkai.Hooks, error) { + return &gkai.Hooks{ + WrapModel: func(ctx context.Context, params *gkai.ModelParams, next gkai.ModelNext) (*gkai.ModelResponse, error) { + response, err := next(ctx, params) + if response != nil { + response.Request = &gkai.ModelRequest{Messages: cloneCheckpointMessages(params.Request.Messages)} + } + return response, err + }, + WrapTool: func(ctx context.Context, params *gkai.ToolParams, next gkai.ToolNext) (*gkai.MultipartToolResponse, error) { + if params == nil || params.Request == nil { + return nil, fmt.Errorf("genkit tool middleware received no provider request") + } + return next(context.WithValue(ctx, genkitToolRequestContextKey{}, params.Request), params) + }, + }, nil +} + func promptParts(req ai.Request) ([]*gkai.Part, error) { parts := make([]*gkai.Part, 0, len(req.Prompt.Attachments)) if req.Prompt.User != "" { diff --git a/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go b/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go index 16f545b5..65518a53 100644 --- a/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go +++ b/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go @@ -22,7 +22,9 @@ var _ = Describe("Genkit resumable tool approval", func() { events := make([]ai.Event, 0, 2) ran := false - _, err := provider.runTool(context.Background(), api.ToolDefinition{ + _, err := provider.runTool(context.WithValue(context.Background(), genkitToolRequestContextKey{}, &gkai.ToolRequest{ + Name: "invoice_update", Ref: "call-update", Input: map[string]any{"amount": 10}, + }), api.ToolDefinition{ Name: "invoice_update", DefaultPermission: api.ToolModeAsk, Handler: func(context.Context, map[string]any) (any, error) { @@ -47,7 +49,10 @@ var _ = Describe("Genkit resumable tool approval", func() { 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} + response := &gkai.ModelResponse{ + Request: &gkai.ModelRequest{Messages: []*gkai.Message{gkai.NewUserTextMessage("Update then inspect.")}}, + Message: gkai.NewModelMessage(pending, completed), FinishReason: gkai.FinishReasonInterrupted, + } state, err := toolApprovalState(request, response) Expect(err).NotTo(HaveOccurred()) @@ -60,6 +65,35 @@ var _ = Describe("Genkit resumable tool approval", func() { Expect(state.Calls[1].Result.Output).To(MatchJSON(`{"amount":10}`)) }) + It("round trips Gemini thought signatures through the private approval checkpoint", func() { + signature := []byte("gemini-thought-signature") + pending := gkai.NewToolRequestPart(&gkai.ToolRequest{ + Name: "accounts_edit", Ref: "call-signed", Input: map[string]any{"id": "acc-1"}, + }) + pending.Metadata = map[string]any{ + "interrupt": map[string]any{"approvalRequired": true}, + "signature": signature, + } + response := &gkai.ModelResponse{ + Request: &gkai.ModelRequest{Messages: []*gkai.Message{gkai.NewUserTextMessage("Edit the account")}}, + Message: gkai.NewModelMessage(pending), FinishReason: gkai.FinishReasonInterrupted, + } + + state, err := toolApprovalState(api.Spec{Prompt: api.Prompt{User: "Edit the account"}}, response) + Expect(err).NotTo(HaveOccurred()) + Expect(state.ProviderCheckpoint).NotTo(BeNil()) + + messages, _, _, err := prepareToolApprovalResume(&api.ToolApprovalResume{ + State: *state, + Decisions: []api.ToolApprovalDecision{{ + ToolCallID: "call-signed", Tool: "accounts_edit", Action: api.ToolApprovalApprove, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + requests := genkitApprovalRequests(messages[len(messages)-1]) + Expect(requests["call-signed"].Metadata["signature"]).To(Equal(signature)) + }) + It("restarts an approved call with edited input and never replays completed siblings", func(ctx SpecContext) { var updateRuns atomic.Int32 var readRuns atomic.Int32 @@ -144,20 +178,20 @@ var _ = Describe("Genkit resumable tool approval", func() { }) 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{ + deny := gkai.NewToolRequestPart(&gkai.ToolRequest{ + Name: "invoice_delete", Ref: "call-deny", Input: map[string]any{"id": "inv-1"}, + }) + deny.Metadata = map[string]any{"interrupt": map[string]any{"approvalRequired": true}} + respond := gkai.NewToolRequestPart(&gkai.ToolRequest{ + Name: "invoice_update", Ref: "call-respond", Input: map[string]any{"amount": 10}, + }) + respond.Metadata = map[string]any{"interrupt": map[string]any{"approvalRequired": true}} + state, err := toolApprovalState(api.Spec{Prompt: api.Prompt{User: "Change it."}}, &gkai.ModelResponse{ + Request: &gkai.ModelRequest{Messages: []*gkai.Message{gkai.NewUserTextMessage("Change it.")}}, + Message: gkai.NewModelMessage(deny, respond), FinishReason: gkai.FinishReasonInterrupted, + }) + Expect(err).NotTo(HaveOccurred()) + 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}`), diff --git a/pkg/ai/provider/genkit/tool_lifecycle_ginkgo_test.go b/pkg/ai/provider/genkit/tool_lifecycle_ginkgo_test.go index 7d77327c..0f77d9fd 100644 --- a/pkg/ai/provider/genkit/tool_lifecycle_ginkgo_test.go +++ b/pkg/ai/provider/genkit/tool_lifecycle_ginkgo_test.go @@ -2,6 +2,8 @@ package genkit import ( "context" + "encoding/json" + "sync/atomic" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -10,6 +12,7 @@ import ( "github.com/flanksource/captain/pkg/api" gkai "github.com/firebase/genkit/go/ai" + gk "github.com/firebase/genkit/go/genkit" ) var _ = Describe("Genkit tool event correlation", func() { @@ -44,7 +47,7 @@ var _ = Describe("Genkit tool event correlation", func() { Expect(err).NotTo(HaveOccurred()) Expect(mapped).To(BeEmpty()) - _, err = provider.runTool(context.Background(), tool, request.Input, emit, correlation) + _, err = provider.runTool(context.WithValue(context.Background(), genkitToolRequestContextKey{}, request), tool, request.Input, emit, correlation) Expect(err).NotTo(HaveOccurred()) Expect(events).To(HaveLen(2)) Expect(events[0].Kind).To(Equal(ai.EventToolUse)) @@ -67,7 +70,7 @@ var _ = Describe("Genkit tool event correlation", func() { _, 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) + _, err = provider.runTool(context.WithValue(context.Background(), genkitToolRequestContextKey{}, first), 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)) @@ -82,7 +85,7 @@ var _ = Describe("Genkit tool event correlation", func() { Expect(err).NotTo(HaveOccurred()) request.Ref = "genkit-assigned-ref" - _, err = provider.runTool(context.Background(), tool, request.Input, emit, correlation) + _, err = provider.runTool(context.WithValue(context.Background(), genkitToolRequestContextKey{}, request), tool, request.Input, emit, correlation) Expect(err).NotTo(HaveOccurred()) Expect(events).To(HaveLen(2)) Expect(events[0].ToolCallID).To(Equal(request.Ref)) @@ -100,13 +103,65 @@ var _ = Describe("Genkit tool event correlation", func() { _, err := chunkToEvents(toolRequestChunk(request), provider.GetModel(), correlation) Expect(err).NotTo(HaveOccurred()) - _, err = provider.runTool(context.Background(), tool, request.Input, emit, correlation) + _, err = provider.runTool(context.WithValue(context.Background(), genkitToolRequestContextKey{}, request), 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("correlates concurrent same-name calls after schema normalization", func(ctx SpecContext) { + var runs atomic.Int32 + genkit := gk.Init(ctx) + modelRef := "test/parallel-journals" + 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) { + if request.Messages[len(request.Messages)-1].Role == gkai.RoleTool { + return &gkai.ModelResponse{Message: gkai.NewModelTextMessage("done"), FinishReason: gkai.FinishReasonStop}, nil + } + message := gkai.NewModelMessage( + gkai.NewToolRequestPart(&gkai.ToolRequest{Name: "journals", Ref: "call-10", Input: json.RawMessage(`{"limit":10}`)}), + gkai.NewToolRequestPart(&gkai.ToolRequest{Name: "journals", Ref: "call-20", Input: json.RawMessage(`{"limit":20}`)}), + ) + 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: "parallel-journals", Backend: api.BackendAnthropic}, + Tools: []api.ToolDefinition{{ + Name: "journals", + DefaultPermission: api.ToolModeOn, + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{"limit": map[string]any{"type": "integer"}}, + "required": []string{"limit"}, + }, + Handler: func(_ context.Context, input map[string]any) (any, error) { + Expect(input["limit"]).To(BeAssignableToTypeOf(int64(0))) + runs.Add(1) + return input, nil + }, + }}, + }, + backend: api.BackendAnthropic, + g: genkit, modelRef: modelRef, + } + + stream, err := provider.ExecuteStream(ctx, api.Spec{Prompt: api.Prompt{User: "Inspect both journal sets."}}) + Expect(err).NotTo(HaveOccurred()) + var callIDs []string + for event := range stream { + Expect(event.Kind).NotTo(Equal(ai.EventError), event.Error) + if event.Kind == ai.EventToolUse { + callIDs = append(callIDs, event.ToolCallID) + } + } + Expect(runs.Load()).To(Equal(int32(2))) + Expect(callIDs).To(ConsistOf("call-10", "call-20")) + }) + It("fails loudly when a tool response has no correlated request", func() { correlation := newToolEventCorrelation() _, err := chunkToEvents(toolResponseChunk(tool.Name, "missing"), provider.GetModel(), correlation) @@ -116,7 +171,7 @@ var _ = Describe("Genkit tool event correlation", func() { 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(err).To(MatchError(ContainSubstring(`no provider request context`))) Expect(events).To(BeEmpty()) }) }) diff --git a/pkg/ai/provider/genkit/tools.go b/pkg/ai/provider/genkit/tools.go index 4eb5f84d..4bc0ef7a 100644 --- a/pkg/ai/provider/genkit/tools.go +++ b/pkg/ai/provider/genkit/tools.go @@ -50,44 +50,7 @@ func (p *Provider) toolOptions(preferences api.ToolPreferences, emit func(ai.Eve } 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 + return captools.ResolveDefinitions(definitions, preferences) } func anthropicStrictToolDefinitions(definitions []api.ToolDefinition) []api.ToolDefinition { @@ -180,8 +143,15 @@ func (p *Provider) runTool( if correlation == nil { return nil, fmt.Errorf("genkit tool %q execution has no correlation state", def.Name) } + request, ok := ctx.Value(genkitToolRequestContextKey{}).(*gkai.ToolRequest) + if !ok || request == nil { + return nil, fmt.Errorf("genkit tool %q execution has no provider request context", def.Name) + } + if request.Name != def.Name { + return nil, fmt.Errorf("genkit tool execution names %q, expected %q", request.Name, def.Name) + } var err error - callID, err = correlation.begin(def.Name, args) + callID, err = correlation.begin(request) if err != nil { return nil, err } diff --git a/pkg/ai/provider/genkit/tools_test.go b/pkg/ai/provider/genkit/tools_test.go index 87c87cf9..fe0bb6cf 100644 --- a/pkg/ai/provider/genkit/tools_test.go +++ b/pkg/ai/provider/genkit/tools_test.go @@ -155,6 +155,8 @@ func eventKinds(events []ai.Event) []ai.EventKind { 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) + request := &gkai.ToolRequest{Name: def.Name, Ref: "provider-call-1", Input: input} + correlation.observeRequest(request) + ctx := context.WithValue(context.Background(), genkitToolRequestContextKey{}, request) + return p.runTool(ctx, def, input, emit, correlation) } diff --git a/pkg/ai/provider/init.go b/pkg/ai/provider/init.go index 87fd077d..12236799 100644 --- a/pkg/ai/provider/init.go +++ b/pkg/ai/provider/init.go @@ -12,6 +12,7 @@ import ( ) func init() { + ai.RegisterRuntimeProbe(ai.BackendClaudeAgent, claudeagent.ProbeRuntime) // API backends are served by Firebase Genkit (replaces the per-SDK providers). ai.RegisterProvider(ai.BackendAnthropic, func(cfg ai.Config) (ai.Provider, error) { return genkit.New(cfg) }) ai.RegisterProvider(ai.BackendOpenAI, func(cfg ai.Config) (ai.Provider, error) { return genkit.New(cfg) }) @@ -28,7 +29,7 @@ func init() { }) ai.RegisterProvider(ai.BackendCodexCLI, func(cfg ai.Config) (ai.Provider, error) { return NewCodexCLI(cfg), nil }) - ai.RegisterProvider(ai.BackendCodexAgent, func(cfg ai.Config) (ai.Provider, error) { return NewCodexAppServer(cfg.Model.Name) }) + ai.RegisterProvider(ai.BackendCodexAgent, func(cfg ai.Config) (ai.Provider, error) { return NewCodexAppServer(cfg) }) // cmux drives an interactive claude/codex TUI inside a tmux/cmux surface, // tailing the session JSONL; the same provider serves both agents (it reads diff --git a/pkg/ai/runtime_probe.go b/pkg/ai/runtime_probe.go new file mode 100644 index 00000000..48e7fd92 --- /dev/null +++ b/pkg/ai/runtime_probe.go @@ -0,0 +1,38 @@ +package ai + +import "sync" + +// RuntimeStatus describes provider-owned local prerequisites independently of +// authentication. A provisioner may make a runtime ready on first use. +type RuntimeStatus struct { + Binary string + BinaryMissing string + DependencyMissing string + Provisioner string + Error string +} + +type RuntimeProbe func() RuntimeStatus + +var ( + runtimeProbeMu sync.RWMutex + runtimeProbes = map[Backend]RuntimeProbe{} +) + +// RegisterRuntimeProbe lets a provider report the prerequisites its real +// launch path uses instead of relying on a generic PATH check. +func RegisterRuntimeProbe(backend Backend, probe RuntimeProbe) { + runtimeProbeMu.Lock() + defer runtimeProbeMu.Unlock() + runtimeProbes[backend] = probe +} + +func probeRuntime(backend Backend) (RuntimeStatus, bool) { + runtimeProbeMu.RLock() + probe, ok := runtimeProbes[backend] + runtimeProbeMu.RUnlock() + if !ok { + return RuntimeStatus{}, false + } + return probe(), true +} diff --git a/pkg/ai/runtime_selector_test.go b/pkg/ai/runtime_selector_test.go index b4ce75d9..efc781de 100644 --- a/pkg/ai/runtime_selector_test.go +++ b/pkg/ai/runtime_selector_test.go @@ -242,12 +242,13 @@ func TestResolvedModelCarriesCapabilities(t *testing.T) { wantResume bool wantIntr bool wantSteer bool + wantTools 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/*"}}, + {"agent:sonnet", registry.ModeAgent, true, true, true, true, []string{"image/png", "image/jpeg", "image/gif", "image/webp", "application/pdf"}}, + {"cli:sonnet", registry.ModeCLI, true, false, false, false, []string{}}, + {"api:sonnet", registry.ModeAPI, false, false, false, true, []string{"image/*"}}, + {"agent:sol", registry.ModeAgent, true, true, false, true, []string{"image/*"}}, } for _, tc := range cases { t.Run(tc.selector, func(t *testing.T) { @@ -265,6 +266,9 @@ func TestResolvedModelCarriesCapabilities(t *testing.T) { 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 got.CallerTools != tc.wantTools { + t.Errorf("CallerTools = %v, want %v", got.CallerTools, tc.wantTools) + } if !reflect.DeepEqual(got.MediaTypes, tc.wantMedia) { t.Errorf("MediaTypes = %v, want %v", got.MediaTypes, tc.wantMedia) } diff --git a/pkg/ai/tools/definitions_ginkgo_test.go b/pkg/ai/tools/definitions_ginkgo_test.go new file mode 100644 index 00000000..d53d9f1e --- /dev/null +++ b/pkg/ai/tools/definitions_ginkgo_test.go @@ -0,0 +1,71 @@ +package tools_test + +import ( + "context" + + "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Caller tool definitions", func() { + noop := func(context.Context, map[string]any) (any, error) { return "ok", nil } + + It("resolves exact preferences before groups and omits disabled tools", func() { + definitions, err := tools.ResolveDefinitions([]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}, + }, api.ToolPreferences{ + "billing": api.ToolModeOff, + "invoice_list": api.ToolModeOn, + "search": api.ToolModeAsk, + }) + + Expect(err).NotTo(HaveOccurred()) + Expect(definitions).To(HaveLen(2)) + Expect(definitions[0].Name).To(Equal("invoice_list")) + Expect(definitions[0].DefaultPermission).To(Equal(api.ToolModeOn)) + Expect(definitions[1].Name).To(Equal("search")) + Expect(definitions[1].DefaultPermission).To(Equal(api.ToolModeAsk)) + }) + + It("validates definitions even when a preference disables them", func() { + _, err := tools.ResolveDefinitions([]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("allows auto only for explicitly read-only non-destructive tools", func() { + readOnly, nonDestructive := true, false + definitions, err := tools.ResolveDefinitions([]api.ToolDefinition{ + { + Name: "invoice_get", ReadOnlyHint: &readOnly, DestructiveHint: &nonDestructive, + DefaultPermission: api.ToolModeAuto, Handler: noop, + }, + {Name: "invoice_update", DefaultPermission: api.ToolModeAuto, Handler: noop}, + }, nil) + + Expect(err).NotTo(HaveOccurred()) + Expect(definitions).To(HaveLen(2)) + Expect(definitions[0].DefaultPermission).To(Equal(api.ToolModeOn)) + Expect(definitions[1].DefaultPermission).To(Equal(api.ToolModeAsk)) + }) + + It("rejects duplicate and provider-unsafe tool names", func() { + _, err := tools.ResolveDefinitions([]api.ToolDefinition{ + {Name: "invoice_get", Handler: noop}, + {Name: "invoice_get", Handler: noop}, + }, nil) + Expect(err).To(MatchError(ContainSubstring(`duplicate caller tool "invoice_get"`))) + + _, err = tools.ResolveDefinitions([]api.ToolDefinition{{ + Name: "invoice/get", Handler: noop, + }}, nil) + Expect(err).To(MatchError(ContainSubstring(`caller tool name "invoice/get"`))) + }) +}) diff --git a/pkg/ai/tools/tools.go b/pkg/ai/tools/tools.go index 530ca8b1..3dea1cf6 100644 --- a/pkg/ai/tools/tools.go +++ b/pkg/ai/tools/tools.go @@ -9,6 +9,7 @@ package tools import ( "context" + "fmt" "sort" "github.com/flanksource/captain/pkg/api" @@ -219,6 +220,74 @@ func NormalizedPreference(prefs ToolPreferences, name string) (ToolMode, bool) { return NormalizeToolMode(mode) } +// ResolveDefinitions validates caller tools, applies exact/group preferences, +// omits disabled tools, and writes the effective permission onto a copy of each +// selected definition. Every provider uses this function so API and agent +// runtimes cannot disagree about the visible tool set. +func ResolveDefinitions(definitions []api.ToolDefinition, preferences ToolPreferences) ([]api.ToolDefinition, error) { + if err := preferences.Validate(); err != nil { + return nil, err + } + selected := make([]api.ToolDefinition, 0, len(definitions)) + seen := make(map[string]struct{}, len(definitions)) + for _, definition := range definitions { + if definition.Name == "" { + return nil, fmt.Errorf("caller tool name cannot be empty") + } + if !validCallerToolName(definition.Name) { + return nil, fmt.Errorf("caller tool name %q contains unsupported characters", definition.Name) + } + if _, ok := seen[definition.Name]; ok { + return nil, fmt.Errorf("duplicate caller tool %q", definition.Name) + } + seen[definition.Name] = struct{}{} + if definition.Handler == nil { + return nil, fmt.Errorf("caller tool %q has no handler", definition.Name) + } + mode := ToolModeAuto + if definition.DefaultPermission != "" { + var ok bool + mode, ok = NormalizeToolMode(definition.DefaultPermission) + if !ok { + return nil, fmt.Errorf("tool %q has invalid default permission %q", definition.Name, definition.DefaultPermission) + } + } + if preferred, ok := EffectivePreference(preferences, ToolInfo{ + Name: definition.Name, Group: definition.Group, + }); ok && preferred != ToolModeAuto { + mode = preferred + } + if mode == ToolModeOff { + continue + } + if mode == ToolModeAuto { + if definition.ReadOnlyHint != nil && *definition.ReadOnlyHint && + definition.DestructiveHint != nil && !*definition.DestructiveHint { + mode = ToolModeOn + } else { + mode = ToolModeAsk + } + } + definition.DefaultPermission = mode + selected = append(selected, definition) + } + return selected, nil +} + +func validCallerToolName(name string) bool { + for _, value := range name { + if value >= 'a' && value <= 'z' || + value >= 'A' && value <= 'Z' || + value >= '0' && value <= '9' || + value == '-' || + value == '_' { + continue + } + return false + } + return true +} + // ToolEntry is one row in the tool-preferences UI: a single ungrouped tool, or a // collapsed group listing its member names. type ToolEntry struct { diff --git a/pkg/aichat/agent_prompt.go b/pkg/aichat/agent_prompt.go new file mode 100644 index 00000000..a9232142 --- /dev/null +++ b/pkg/aichat/agent_prompt.go @@ -0,0 +1,81 @@ +package aichat + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/api" +) + +func agentPrompt(messages []api.Message, resumed bool) (string, []api.AttachmentRef, error) { + selected := messages + if resumed { + index := lastUserMessage(messages) + if index < 0 { + return "", nil, fmt.Errorf("resumed agent chat requires a user message") + } + selected = messages[index : index+1] + } + blocks := make([]string, 0, len(selected)) + attachments := make([]api.AttachmentRef, 0) + for _, message := range selected { + text, refs, err := agentMessageText(message) + if err != nil { + return "", nil, err + } + attachments = append(attachments, refs...) + if strings.TrimSpace(text) != "" { + blocks = append(blocks, fmt.Sprintf("%s:\n%s", message.Role, text)) + } + } + if len(blocks) == 1 && len(selected) == 1 && selected[0].Role == api.RoleUser { + return strings.TrimPrefix(blocks[0], string(api.RoleUser)+":\n"), attachments, nil + } + return strings.Join(blocks, "\n\n"), attachments, nil +} + +func lastUserMessage(messages []api.Message) int { + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == api.RoleUser { + return i + } + } + return -1 +} + +func agentMessageText(message api.Message) (string, []api.AttachmentRef, error) { + lines := make([]string, 0, len(message.Parts)) + attachments := make([]api.AttachmentRef, 0) + for _, part := range message.Parts { + switch part.Type { + case api.PartText: + lines = append(lines, part.Text) + case api.PartReasoning: + continue + case api.PartAttachment: + attachments = append(attachments, *part.Attachment) + lines = append(lines, "[Attachment: "+part.Attachment.Filename+"]") + case api.PartToolRequest: + lines = append(lines, fmt.Sprintf("Tool request %s (%s): %s", + part.ToolRequest.Name, part.ToolRequest.ToolCallID, jsonText(part.ToolRequest.Input))) + case api.PartToolResult: + if part.ToolResult.Error != "" { + lines = append(lines, fmt.Sprintf("Tool result %s failed: %s", part.ToolResult.ToolCallID, part.ToolResult.Error)) + } else { + lines = append(lines, fmt.Sprintf("Tool result %s: %s", + part.ToolResult.ToolCallID, jsonText(part.ToolResult.Output))) + } + default: + return "", nil, fmt.Errorf("unsupported agent prompt part %q", part.Type) + } + } + return strings.Join(lines, "\n"), attachments, nil +} + +func jsonText(raw json.RawMessage) string { + if len(raw) == 0 { + return "{}" + } + return string(raw) +} diff --git a/pkg/aichat/aimock_lifecycle_integration_test.go b/pkg/aichat/aimock_lifecycle_integration_test.go new file mode 100644 index 00000000..94f1b7f7 --- /dev/null +++ b/pkg/aichat/aimock_lifecycle_integration_test.go @@ -0,0 +1,504 @@ +package aichat_test + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/flanksource/captain/pkg/ai" + _ "github.com/flanksource/captain/pkg/ai/provider" + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/aimock" + "github.com/flanksource/captain/pkg/aimock/anthropicmock" + "github.com/flanksource/captain/pkg/aimock/openaimock" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/flanksource/commons-db/dbtest" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type lifecycleRuntime struct { + name string + model api.Model + protocol string + agent bool + binaries []string + scenario string + configEnv string +} + +type lifecycleMock struct { + server aimock.Server + apiURL string + remaining func() []string +} + +type lifecycleRoundTripFunc func(*http.Request) (*http.Response, error) + +func (f lifecycleRoundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { + return f(request) +} + +type lifecycleSignalReadCloser struct { + io.ReadCloser + match string + signal chan struct{} + once sync.Once + read strings.Builder +} + +func (r *lifecycleSignalReadCloser) Read(buffer []byte) (int, error) { + read, err := r.ReadCloser.Read(buffer) + if read > 0 { + r.read.Write(buffer[:read]) + if strings.Contains(r.read.String(), r.match) { + r.once.Do(func() { close(r.signal) }) + } + } + return read, err +} + +type realChatResolver struct{} + +func (realChatResolver) Models(context.Context) (aichat.ModelCatalogResponse, error) { + return nil, nil +} + +func (realChatResolver) Runtimes(context.Context) ([]api.RuntimeFamily, error) { + return api.RuntimeCatalog(), nil +} + +func (realChatResolver) 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 { + return nil, fmt.Errorf("backend %q is not streaming", provider.GetBackend()) + } + return streaming, nil +} + +type httpResult struct { + status int + body []byte + err error +} + +var _ = Describe("Mocked Captain chat lifecycle", func() { + DescribeTable("persists request, response, approval, interruption, and resume", + func(ctx SpecContext, runtime lifecycleRuntime) { + if runtime.agent { + if os.Getenv("CAPTAIN_AIMOCK_AGENT_E2E") != "1" { + Skip("set CAPTAIN_AIMOCK_AGENT_E2E=1 to run real agent-process lifecycle tests") + } + for _, binary := range runtime.binaries { + _, err := exec.LookPath(binary) + Expect(err).NotTo(HaveOccurred(), "%s is required when the agent E2E gate is enabled", binary) + } + GinkgoT().Setenv(runtime.configEnv, GinkgoT().TempDir()) + } + GinkgoT().Setenv(api.MonitorHooksEnv, "off") + + mock := startLifecycleMock(runtime) + DeferCleanup(mock.server.Close) + dbName := "captain_aichat_mock_" + strings.NewReplacer(" ", "_", "-", "_").Replace(runtime.name) + testDB := dbtest.ForGinkgo(dbtest.Options{Name: dbName}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + store, err := aichat.NewDatabaseThreadStore(db) + Expect(err).NotTo(HaveOccurred()) + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + + var toolCalls atomic.Int32 + var inputMu sync.Mutex + var approvedInput map[string]any + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: realChatResolver{}, Threads: aichat.FixedThreadStore(store), Authority: authority, + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{ProviderConfig: api.Config{ + APIURL: mock.apiURL, APIKey: aimock.DummyKey, + }}, nil + }), + Tools: aichat.StaticToolProvider([]api.ToolDefinition{{ + Name: "accounts_edit", DefaultPermission: api.ToolModeAsk, + Handler: func(_ context.Context, input map[string]any) (any, error) { + toolCalls.Add(1) + inputMu.Lock() + approvedInput = input + inputMu.Unlock() + return map[string]any{"id": input["id"], "name": input["name"], "updated": true}, nil + }, + }}), + }) + server := httptest.NewServer(service.Handler()) + DeferCleanup(server.Close) + client := server.Client() + + responseSession := createLifecycleSession(ctx, client, server.URL, "Response") + response := sendLifecycleChat(ctx, client, server.URL, responseSession.ID, runtime.model, nil, + "user-response", "Return the lifecycle greeting") + Expect(response.err).NotTo(HaveOccurred()) + Expect(response.status).To(Equal(http.StatusOK), string(response.body)) + Expect(lifecycleSSEText(response.body)).To(Equal("Lifecycle response complete."), string(response.body)) + assertCompletedSession(ctx, client, server.URL, responseSession.ID, runtime, "Lifecycle response complete.") + + approved := runApprovalFlow(ctx, client, server.URL, mock.server, runtime.model, "Approve", true, "approved by test") + Expect(approved.Requests).To(HaveLen(1)) + Expect(approved.Requests[0].State).To(Equal(string(database.TurnRequestStateApproved))) + Expect(toolCalls.Load()).To(Equal(int32(1))) + inputMu.Lock() + Expect(approvedInput).To(Equal(map[string]any{"id": "acc-1", "name": "Approved Account"})) + inputMu.Unlock() + + rejected := runApprovalFlow(ctx, client, server.URL, mock.server, runtime.model, "Reject", false, "rejected by test") + Expect(rejected.Requests).To(HaveLen(1)) + Expect(rejected.Requests[0].State).To(Equal(string(database.TurnRequestStateDenied))) + Expect(rejected.Requests[0].Reason).To(Equal("rejected by test")) + Expect(toolCalls.Load()).To(Equal(int32(1)), "a rejected tool must not execute") + + interruptSession := createLifecycleSession(ctx, client, server.URL, "Interrupt") + chatResult := make(chan httpResult, 1) + responseStarted := make(chan struct{}) + streamClient := *client + transport := client.Transport + if transport == nil { + transport = http.DefaultTransport + } + streamClient.Transport = lifecycleRoundTripFunc(func(request *http.Request) (*http.Response, error) { + response, roundTripErr := transport.RoundTrip(request) + if roundTripErr == nil { + response.Body = &lifecycleSignalReadCloser{ + ReadCloser: response.Body, match: "Partial", signal: responseStarted, + } + } + return response, roundTripErr + }) + go func() { + chatResult <- sendLifecycleChat(ctx, &streamClient, server.URL, interruptSession.ID, runtime.model, nil, + "user-interrupt", "Wait for the lifecycle interrupt") + }() + Eventually(responseStarted).WithTimeout(30 * time.Second).Should(BeClosed()) + Eventually(func(g Gomega) { + aggregate := getLifecycleSession(ctx, client, server.URL, interruptSession.ID) + g.Expect(aggregate.LifecycleStatus).To(Equal(string(database.SessionLifecycleRunning))) + if runtime.agent { + g.Expect(aggregate.ProviderSessionID).NotTo(BeEmpty()) + } + }).WithTimeout(30 * time.Second).Should(Succeed()) + interruptedProviderID := getLifecycleSession(ctx, client, server.URL, interruptSession.ID).ProviderSessionID + interruptContext, cancelInterrupt := context.WithTimeout(ctx, 5*time.Second) + defer cancelInterrupt() + interrupt := postLifecycleJSON(interruptContext, client, http.MethodPost, + server.URL+"/api/chat/sessions/"+interruptSession.ID+"/interrupt", nil) + Expect(interrupt.err).NotTo(HaveOccurred()) + Expect(interrupt.status).To(Equal(http.StatusOK), string(interrupt.body)) + var interrupted httpResult + Eventually(chatResult).WithTimeout(30 * time.Second).Should(Receive(&interrupted)) + Expect(interrupted.err).NotTo(HaveOccurred()) + Expect(string(interrupted.body)).To(ContainSubstring(`"interrupted":true`)) + Expect(string(interrupted.body)).NotTo(ContainSubstring(`"type":"error"`)) + + interruptedSession := getLifecycleSession(ctx, client, server.URL, interruptSession.ID) + Expect(interruptedSession.LifecycleStatus).To(Equal(string(database.SessionLifecycleInterrupted))) + Expect(interruptedSession.ActivityState).To(Equal(string(database.SessionActivityIdle))) + Expect(interruptedSession.Turns).To(HaveLen(1)) + Expect(interruptedSession.Turns[0].Status).To(Equal(string(database.TurnStatusInterrupted))) + Expect(interruptedSession.Turns[0].StopReason).To(Equal("interrupt")) + + resumedMessages := lifecycleMessages(interruptedSession.Messages) + resumed := sendLifecycleChat(ctx, client, server.URL, interruptSession.ID, runtime.model, resumedMessages, + "user-resume", "Resume after the lifecycle interrupt") + Expect(resumed.err).NotTo(HaveOccurred()) + Expect(resumed.status).To(Equal(http.StatusOK), string(resumed.body)) + finalSession := getLifecycleSession(ctx, client, server.URL, interruptSession.ID) + Expect(finalSession.LifecycleStatus).To(Equal(string(database.SessionLifecycleSucceeded))) + Expect(finalSession.ActivityState).To(Equal(string(database.SessionActivityIdle))) + Expect(finalSession.Turns).To(HaveLen(2)) + Expect(finalSession.Turns[0].Status).To(Equal(string(database.TurnStatusInterrupted))) + Expect(finalSession.Turns[1].Status).To(Equal(string(database.TurnStatusEnded))) + if runtime.agent { + Expect(interruptedProviderID).NotTo(BeEmpty()) + Expect(finalSession.ProviderSessionID).To(Equal(interruptedProviderID)) + } else { + Expect(finalSession.ProviderSessionID).To(BeEmpty()) + } + + Eventually(func() bool { + for _, request := range mock.server.Requests() { + if strings.Contains(request.Request.LastUserText(), "Wait for the lifecycle interrupt") { + return request.Cancelled && request.Miss == "" + } + } + return false + }).Should(BeTrue()) + Expect(mock.remaining()).To(BeEmpty()) + for _, request := range mock.server.Requests() { + Expect(request.Miss).To(BeEmpty(), "%s %s", request.Method, request.Path) + } + }, + Entry("Anthropic API", lifecycleRuntime{ + name: "anthropic_api", model: api.Model{Name: "claude-sonnet-4-6", Backend: api.BackendAnthropic, Mode: api.ModeAPI}, + protocol: aimock.SectionAnthropic, scenario: "chat-api-flows.yaml", + }), + Entry("OpenAI API", lifecycleRuntime{ + name: "openai_api", model: api.Model{Name: "gpt-5", Backend: api.BackendOpenAI, Mode: api.ModeAPI}, + protocol: aimock.SectionOpenAI, scenario: "chat-api-flows.yaml", + }), + Entry("Claude Agent", lifecycleRuntime{ + name: "claude_agent", model: api.Model{Name: "claude-sonnet-5", Backend: api.BackendClaudeAgent, Mode: api.ModeAgent}, + protocol: aimock.SectionAnthropic, agent: true, binaries: []string{"npm", "claude"}, + scenario: "chat-agent-flows.yaml", configEnv: "CLAUDE_CONFIG_DIR", + }), + Entry("Codex Agent", lifecycleRuntime{ + name: "codex_agent", model: api.Model{Name: "gpt-5.6-sol", Backend: api.BackendCodexAgent, Mode: api.ModeAgent}, + protocol: aimock.SectionOpenAI, agent: true, binaries: []string{"codex"}, + scenario: "chat-agent-flows.yaml", configEnv: "CODEX_HOME", + }), + ) +}) + +func startLifecycleMock(runtime lifecycleRuntime) lifecycleMock { + scenario, err := aimock.Load(filepath.Join("..", "aimock", "testdata", "scenarios", runtime.scenario)) + Expect(err).NotTo(HaveOccurred()) + if runtime.protocol == aimock.SectionAnthropic { + server, err := anthropicmock.Start(anthropicmock.Options{Scenario: scenario}) + Expect(err).NotTo(HaveOccurred()) + return lifecycleMock{server: server, apiURL: server.APIURL(), remaining: server.Remaining} + } + server, err := openaimock.Start(openaimock.Options{Scenario: scenario}) + Expect(err).NotTo(HaveOccurred()) + return lifecycleMock{server: server, apiURL: server.APIURL(), remaining: server.Remaining} +} + +func createLifecycleSession(ctx context.Context, client *http.Client, baseURL, title string) aichat.Thread { + result := postLifecycleJSON(ctx, client, http.MethodPost, baseURL+"/api/chat/sessions", map[string]string{"title": title}) + Expect(result.err).NotTo(HaveOccurred()) + Expect(result.status).To(Equal(http.StatusCreated), string(result.body)) + var thread aichat.Thread + Expect(json.Unmarshal(result.body, &thread)).To(Succeed()) + return thread +} + +func sendLifecycleChat( + ctx context.Context, + client *http.Client, + baseURL, sessionID string, + model api.Model, + messages []aichat.UIMessage, + messageID, prompt string, +) httpResult { + messages = append(messages, aichat.UIMessage{ + ID: messageID, Role: string(api.RoleUser), Parts: []aichat.UIPart{{Type: "text", Text: prompt}}, + }) + return postLifecycleJSON(ctx, client, http.MethodPost, baseURL+"/api/chat", aichat.ChatRequest{ + ID: sessionID, ThreadID: sessionID, Trigger: "submit-message", Runtime: &model, Messages: messages, + }) +} + +func runApprovalFlow( + ctx context.Context, + client *http.Client, + baseURL string, + mockServer aimock.Server, + model api.Model, + verb string, + approved bool, + reason string, +) session.Session { + thread := createLifecycleSession(ctx, client, baseURL, verb) + chatResult := make(chan httpResult, 1) + go func() { + chatResult <- sendLifecycleChat(ctx, client, baseURL, thread.ID, model, nil, + "user-"+strings.ToLower(verb), verb+" the account update") + }() + var pending session.Session + Eventually(func(g Gomega) { + pending = getLifecycleSession(ctx, client, baseURL, thread.ID) + if len(pending.Requests) == 0 { + select { + case chat := <-chatResult: + g.Expect(chat.err).NotTo(HaveOccurred()) + g.Expect(chat.status).To(Equal(http.StatusOK), string(chat.body)) + g.Expect(pending.Requests).To(HaveLen(1), string(chat.body)) + default: + g.Expect(pending.Requests).To(HaveLen(1), lifecycleRequestsJSON(mockServer.Requests())) + } + return + } + g.Expect(pending.Requests).To(HaveLen(1)) + g.Expect(pending.Requests[0].State).To(Equal(string(database.TurnRequestStatePending))) + }).WithTimeout(30 * time.Second).Should(Succeed()) + body := map[string]any{"approved": approved, "reason": reason} + if approved { + body["updatedInput"] = map[string]any{"id": "acc-1", "name": "Approved Account"} + } + decision := postLifecycleJSON(ctx, client, http.MethodPost, + baseURL+"/api/chat/sessions/"+thread.ID+"/approvals/"+pending.Requests[0].ID, body) + Expect(decision.err).NotTo(HaveOccurred()) + Expect(decision.status).To(Equal(http.StatusOK), string(decision.body)) + var completed session.Session + Eventually(func(g Gomega) { + completed = getLifecycleSession(ctx, client, baseURL, thread.ID) + g.Expect(completed.Turns).To(HaveLen(1)) + g.Expect(completed.Turns[0].Status).To(Equal(string(database.TurnStatusEnded)), + lifecycleRequestsJSON(mockServer.Requests())) + }).WithTimeout(30 * time.Second).Should(Succeed()) + var chat httpResult + Eventually(chatResult).WithTimeout(30 * time.Second).Should(Receive(&chat)) + Expect(chat.err).NotTo(HaveOccurred()) + Expect(chat.status).To(Equal(http.StatusOK), string(chat.body)) + conflict := postLifecycleJSON(ctx, client, http.MethodPost, + baseURL+"/api/chat/sessions/"+thread.ID+"/approvals/"+pending.Requests[0].ID, + map[string]any{"approved": !approved, "reason": "conflicting replay"}) + Expect(conflict.err).NotTo(HaveOccurred()) + Expect(conflict.status).To(Equal(http.StatusConflict), string(conflict.body)) + return completed +} + +func assertCompletedSession( + ctx context.Context, + client *http.Client, + baseURL, sessionID string, + runtime lifecycleRuntime, + text string, +) { + aggregate := getLifecycleSession(ctx, client, baseURL, sessionID) + Expect(aggregate.LifecycleStatus).To(Equal(string(database.SessionLifecycleSucceeded))) + Expect(aggregate.ActivityState).To(Equal(string(database.SessionActivityIdle))) + Expect(aggregate.ExecutionMode).To(Equal(runtime.model.Mode)) + Expect(aggregate.Backend).To(Equal(string(runtime.model.Backend))) + Expect(aggregate.Model).To(Equal(runtime.model.Name)) + Expect(aggregate.Turns).To(HaveLen(1)) + Expect(aggregate.Turns[0].Status).To(Equal(string(database.TurnStatusEnded))) + Expect(aggregate.Turns[0].StopReason).To(Equal("stop")) + encoded, err := json.Marshal(aggregate.Messages) + Expect(err).NotTo(HaveOccurred()) + Expect(string(encoded)).To(ContainSubstring(text)) + Expect(aggregate.Usage.TotalTokens()).To(BeNumerically(">", 0)) +} + +func getLifecycleSession(ctx context.Context, client *http.Client, baseURL, id string) session.Session { + request, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+"/api/chat/sessions/"+id, nil) + Expect(err).NotTo(HaveOccurred()) + response, err := client.Do(request) + Expect(err).NotTo(HaveOccurred()) + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + Expect(err).NotTo(HaveOccurred()) + Expect(response.StatusCode).To(Equal(http.StatusOK), string(body)) + var aggregate session.Session + Expect(json.Unmarshal(body, &aggregate)).To(Succeed()) + return aggregate +} + +func lifecycleMessages(messages []session.Message) []aichat.UIMessage { + encoded, err := json.Marshal(messages) + Expect(err).NotTo(HaveOccurred()) + var output []aichat.UIMessage + Expect(json.Unmarshal(encoded, &output)).To(Succeed()) + return output +} + +func lifecycleSSEText(body []byte) string { + var textValue strings.Builder + for _, line := range strings.Split(string(body), "\n") { + payload := strings.TrimPrefix(line, "data: ") + if payload == line || payload == "[DONE]" { + continue + } + part := struct { + Type string `json:"type"` + Delta string `json:"delta"` + }{} + if json.Unmarshal([]byte(payload), &part) == nil && part.Type == "text-delta" { + textValue.WriteString(part.Delta) + } + } + return textValue.String() +} + +func lifecycleRequestsJSON(requests []aimock.Recorded) string { + type requestDiagnostic struct { + Path string `json:"path"` + LastUserText string `json:"lastUserText,omitempty"` + ToolResults []string `json:"toolResults,omitempty"` + ToolNames []string `json:"toolNames,omitempty"` + MCPTools map[string]json.RawMessage `json:"mcpTools,omitempty"` + MCPDefinitions map[string]json.RawMessage `json:"mcpDefinitions,omitempty"` + Miss string `json:"miss,omitempty"` + Cancelled bool `json:"cancelled,omitempty"` + } + diagnostics := make([]requestDiagnostic, 0, len(requests)) + for _, request := range requests { + diagnostic := requestDiagnostic{ + Path: request.Path, LastUserText: request.Request.LastUserText(), + ToolResults: request.Request.ToolResultNames(), + ToolNames: request.Request.ToolNames, Miss: request.Miss, Cancelled: request.Cancelled, + } + for name, schema := range request.Request.ToolSchemas { + if strings.HasPrefix(name, "mcp__") { + if diagnostic.MCPTools == nil { + diagnostic.MCPTools = map[string]json.RawMessage{} + } + diagnostic.MCPTools[name] = schema + } + } + for name, definition := range request.Request.ToolDefinitions { + if strings.HasPrefix(name, "mcp__") { + if diagnostic.MCPDefinitions == nil { + diagnostic.MCPDefinitions = map[string]json.RawMessage{} + } + diagnostic.MCPDefinitions[name] = definition + } + } + diagnostics = append(diagnostics, diagnostic) + } + encoded, err := json.MarshalIndent(diagnostics, "", " ") + if err != nil { + return err.Error() + } + return string(encoded) +} + +func postLifecycleJSON(ctx context.Context, client *http.Client, method, url string, body any) httpResult { + var payload io.Reader + if body != nil { + encoded, err := json.Marshal(body) + if err != nil { + return httpResult{err: err} + } + payload = bytes.NewReader(encoded) + } + request, err := http.NewRequestWithContext(ctx, method, url, payload) + if err != nil { + return httpResult{err: err} + } + if body != nil { + request.Header.Set("Content-Type", "application/json") + } + response, err := client.Do(request) + if err != nil { + return httpResult{err: err} + } + defer response.Body.Close() + responseBody, err := io.ReadAll(response.Body) + return httpResult{status: response.StatusCode, body: responseBody, err: err} +} diff --git a/pkg/aichat/approval_execution.go b/pkg/aichat/approval_execution.go new file mode 100644 index 00000000..fdc50326 --- /dev/null +++ b/pkg/aichat/approval_execution.go @@ -0,0 +1,99 @@ +package aichat + +import ( + "context" + "fmt" + "strings" + + aitools "github.com/flanksource/captain/pkg/ai/tools" + "github.com/flanksource/captain/pkg/api" +) + +func (s *Service) resumeToolApproval(ctx context.Context, threadID string, continuation *ApprovalContinuation) error { + if continuation == nil || continuation.Execution == nil || continuation.Spec.ToolApproval == nil { + return fmt.Errorf("tool approval continuation is incomplete") + } + execution := continuation.Execution + defer closeExecution(execution) + settings, err := s.runtimeSettings(ctx) + if err != nil { + return fmt.Errorf("load chat runtime settings: %w", err) + } + set, err := s.loadTools(ctx) + if err != nil { + return err + } + definitions, err := aitools.ResolveDefinitions(set.Definitions, continuation.Spec.ToolPreferences) + if err != nil { + return err + } + config := settings.ProviderConfig + config.Model = continuation.Spec.Model + config.Budget = continuation.Spec.Budget + config.SessionID = continuation.Spec.SessionID + config.CaptainSessionID = execution.CaptainSessionID() + config.Tools = definitions + config, err = s.prepareProviderConfig(ctx, config) + if err != nil { + return err + } + continuation.Spec.Model = config.Model + provider, err := s.resolver.Provider(ctx, config) + if err != nil { + return err + } + defer func() { + if closeErr := closeProvider(provider); closeErr != nil { + serviceLog.Errorf("close approval continuation provider: %v", closeErr) + } + }() + if len(definitions) > 0 { + capability, ok := api.ProviderAs[api.ToolCapableProvider](provider) + if !ok || !capability.SupportsCallerTools() { + return fmt.Errorf("backend %q does not support caller tools", provider.GetBackend()) + } + } + store, err := s.threads(ctx) + if err != nil { + return err + } + thread, err := store.Get(ctx, threadID) + if err != nil { + return err + } + if len(thread.Messages) == 0 { + return fmt.Errorf("captain chat session %s has no suspended assistant message", threadID) + } + seed := thread.Messages[len(thread.Messages)-1] + if !strings.EqualFold(seed.Role, string(api.RoleAssistant)) || seed.TurnID != execution.TurnID() { + return fmt.Errorf("captain chat session %s does not end with the suspended turn %s", threadID, execution.TurnID()) + } + request := ChatRequest{ + ID: threadID, ThreadID: threadID, Trigger: "submit-message", MessageID: seed.ID, + Messages: []UIMessage{seed}, ToolApproval: continuation.Spec.ToolApproval, + } + streamContext, cancel := context.WithCancel(ctx) + defer cancel() + events, err := provider.ExecuteStream(streamContext, continuation.Spec) + if err != nil { + return err + } + active := newActiveTurn(streamContext, provider, execution, cancel) + if err := s.registerActiveTurn(threadID, active); err != nil { + return err + } + defer s.unregisterActiveTurn(threadID, active) + events = active.stream(events) + events = observeExecutionEvents(streamContext, execution, events) + // A resumed approval writes no SSE stream, so no TurnCosts sink is needed; + // the thread total is recomputed from the database on the next read. + resumed := s.persistedEvents(streamContext, persistedEventOptions{ + Request: request, TurnID: execution.TurnID(), Model: continuation.Spec.Model, + }, events) + for event := range resumed { + if event.Kind == api.EventError { + return fmt.Errorf("resume provider approval: %s", event.Error) + } + } + return nil +} diff --git a/pkg/aichat/approval_http.go b/pkg/aichat/approval_http.go new file mode 100644 index 00000000..1ea5b882 --- /dev/null +++ b/pkg/aichat/approval_http.go @@ -0,0 +1,75 @@ +package aichat + +import ( + "encoding/json" + "fmt" + "net/http" +) + +func (s *Service) handleResolveToolApproval(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w, request) + if store == nil { + return + } + if s.options.Authority == nil { + http.Error(w, "execution authority is not configured", http.StatusNotImplemented) + return + } + threadID := request.PathValue("id") + if _, err := store.Get(request.Context(), threadID); err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + body := struct { + Approved *bool `json:"approved"` + UpdatedInput map[string]any `json:"updatedInput,omitempty"` + Reason string `json:"reason,omitempty"` + }{} + decoder := json.NewDecoder(request.Body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&body); err != nil { + http.Error(w, fmt.Sprintf("invalid tool approval decision: %v", err), http.StatusBadRequest) + return + } + if body.Approved == nil { + http.Error(w, "tool approval decision requires approved", http.StatusBadRequest) + return + } + if !*body.Approved && body.UpdatedInput != nil { + http.Error(w, "denied tool approval cannot replace input", http.StatusBadRequest) + return + } + continuation, err := s.options.Authority.ResolveToolApproval(request.Context(), ToolApprovalResolution{ + ThreadID: threadID, ApprovalID: request.PathValue("approvalID"), + Approved: *body.Approved, UpdatedInput: body.UpdatedInput, Reason: body.Reason, + }) + if err != nil { + http.Error(w, err.Error(), http.StatusConflict) + return + } + if continuation != nil { + if err := s.resumeToolApproval(request.Context(), threadID, continuation); err != nil { + http.Error(w, err.Error(), http.StatusBadGateway) + return + } + } + if sessions, ok := store.(SessionReader); ok { + aggregate, err := sessions.GetSession(request.Context(), threadID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := writeJSON(w, http.StatusOK, aggregate); err != nil { + serviceLog.Errorf("write approved chat session %q: %v", threadID, err) + } + return + } + thread, err := store.Get(request.Context(), threadID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + if err := writeJSON(w, http.StatusOK, thread); err != nil { + serviceLog.Errorf("write approved chat thread %q: %v", threadID, err) + } +} diff --git a/pkg/aichat/approval_resume.go b/pkg/aichat/approval_resume.go deleted file mode 100644 index 0d46dc2e..00000000 --- a/pkg/aichat/approval_resume.go +++ /dev/null @@ -1,114 +0,0 @@ -package aichat - -import ( - "encoding/json" - "fmt" - "reflect" - - "github.com/flanksource/captain/pkg/api" -) - -func resolveToolApproval(request *ChatRequest) error { - if request.ToolApproval != nil || len(request.Messages) == 0 { - return nil - } - message := request.Messages[len(request.Messages)-1] - if message.Role != string(api.RoleAssistant) { - return nil - } - responded := make(map[string]UIPart) - var stateData json.RawMessage - for _, part := range message.Parts { - if part.Type == "data-tool-approval" { - stateData = part.Data - } - if !part.IsTool() || part.State != "approval-responded" { - continue - } - if part.ToolCallID == "" { - return fmt.Errorf("approval response has no tool call ID") - } - if _, exists := responded[part.ToolCallID]; exists { - return fmt.Errorf("duplicate approval response for tool call %q", part.ToolCallID) - } - responded[part.ToolCallID] = part - } - if len(responded) == 0 { - return nil - } - if len(stateData) == 0 { - return fmt.Errorf("approval response is missing durable tool approval state") - } - var state api.ToolApprovalState - if err := json.Unmarshal(stateData, &state); err != nil { - return fmt.Errorf("decode durable tool approval state: %w", err) - } - if err := state.Validate(); err != nil { - return fmt.Errorf("validate durable tool approval state: %w", err) - } - decisions := make([]api.ToolApprovalDecision, 0, len(responded)) - matched := make(map[string]bool, len(responded)) - for _, pending := range state.Pending() { - part, ok := responded[pending.ToolCallID] - if !ok { - return fmt.Errorf("pending tool call %q has no approval response", pending.ToolCallID) - } - if err := validateApprovalPart(pending, part); err != nil { - return err - } - action := api.ToolApprovalDeny - if *part.Approval.Approved { - action = api.ToolApprovalApprove - } - decisions = append(decisions, api.ToolApprovalDecision{ - ToolCallID: pending.ToolCallID, - Tool: pending.Tool, - Action: action, - Message: part.Approval.Reason, - }) - if action == api.ToolApprovalApprove { - decisions[len(decisions)-1].Message = "" - } - matched[pending.ToolCallID] = true - } - for id := range responded { - if !matched[id] { - return fmt.Errorf("approval response references non-pending tool call %q", id) - } - } - resume := &api.ToolApprovalResume{State: state, Decisions: decisions} - if err := resume.Validate(); err != nil { - return fmt.Errorf("validate tool approval resume: %w", err) - } - request.ToolApproval = resume - return nil -} - -func validateApprovalPart(pending api.ToolApprovalRequest, part UIPart) error { - if part.Approval == nil || part.Approval.Approved == nil { - return fmt.Errorf("tool call %q has no completed approval response", pending.ToolCallID) - } - if part.Approval.ID != pending.ToolCallID { - return fmt.Errorf("tool call %q approval ID is %q", pending.ToolCallID, part.Approval.ID) - } - if part.EffectiveToolName() != pending.Tool { - return fmt.Errorf( - "tool call %q approval names %q, want %q", - pending.ToolCallID, part.EffectiveToolName(), pending.Tool, - ) - } - if !equalPartJSON(part.Input, pending.Input) { - return fmt.Errorf("tool call %q approval input does not match durable state", pending.ToolCallID) - } - return nil -} - -func equalPartJSON(left, right json.RawMessage) bool { - if len(left) == 0 || len(right) == 0 { - return len(left) == len(right) - } - var leftValue, rightValue any - return json.Unmarshal(left, &leftValue) == nil && - json.Unmarshal(right, &rightValue) == nil && - reflect.DeepEqual(leftValue, rightValue) -} diff --git a/pkg/aichat/approval_resume_ginkgo_test.go b/pkg/aichat/approval_resume_ginkgo_test.go deleted file mode 100644 index 7c545db1..00000000 --- a/pkg/aichat/approval_resume_ginkgo_test.go +++ /dev/null @@ -1,249 +0,0 @@ -package aichat_test - -import ( - "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 approvalFixture struct { - state api.ToolApprovalState - user aichat.UIMessage - assistant aichat.UIMessage -} - -func newApprovalFixture(approved ...bool) approvalFixture { - requests := make([]api.Part, len(approved)) - calls := make([]api.ToolApprovalCall, len(approved)) - parts := make([]aichat.UIPart, len(approved)) - for i, allow := range approved { - callID := fmt.Sprintf("call-%02d", i+1) - tool := fmt.Sprintf("example_tool_%02d", i+1) - input := json.RawMessage(fmt.Sprintf(`{"index":%d}`, i+1)) - requests[i] = api.Part{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ - ToolCallID: callID, Name: tool, Input: input, - }} - calls[i] = api.ToolApprovalCall{Request: api.ToolApprovalRequest{ - ToolCallID: callID, Tool: tool, Input: input, - }} - allowCopy := allow - parts[i] = aichat.UIPart{ - Type: "dynamic-tool", ToolName: tool, ToolCallID: callID, - State: "approval-responded", Input: input, - Approval: &aichat.Approval{ID: callID, Approved: &allowCopy}, - } - } - user := aichat.UIMessage{ID: "message-user", Role: "user", Parts: []aichat.UIPart{{ - Type: "text", Text: "Run the example tools.", - }}} - state := api.ToolApprovalState{ - Messages: []api.Message{ - {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Run the example tools."}}}, - {Role: api.RoleAssistant, Parts: requests}, - }, - Calls: calls, - } - raw, err := json.Marshal(state) - Expect(err).NotTo(HaveOccurred()) - parts = append(parts, aichat.UIPart{Type: "data-tool-approval", Data: raw}) - return approvalFixture{ - state: state, - user: user, - assistant: aichat.UIMessage{ID: "message-assistant", Role: "assistant", Parts: parts}, - } -} - -func suspendedAssistant(fixture approvalFixture) aichat.UIMessage { - message := fixture.assistant - message.Parts = append([]aichat.UIPart(nil), fixture.assistant.Parts...) - for i := range fixture.state.Calls { - message.Parts[i].State = "approval-requested" - message.Parts[i].Approval = &aichat.Approval{ID: fixture.state.Calls[i].Request.ToolCallID} - } - return message -} - -var _ = Describe("AI SDK approval resume", func() { - It("reconstructs all batched approval decisions from DefaultChatTransport messages", func() { - const reportedCallCount = 13 - approved := make([]bool, reportedCallCount) - for i := range approved { - approved[i] = true - } - fixture := newApprovalFixture(approved...) - 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: "anthropic/claude-opus-5", Messages: []aichat.UIMessage{fixture.user, fixture.assistant}, - })) - - Expect(response.Code).To(Equal(http.StatusOK), response.Body.String()) - Expect(provider.specs).To(HaveLen(1)) - resume := provider.specs[0].ToolApproval - Expect(resume).NotTo(BeNil()) - Expect(provider.specs[0].Messages).To(BeNil()) - Expect(resume.State).To(Equal(fixture.state)) - Expect(resume.Decisions).To(HaveLen(reportedCallCount)) - for i, decision := range resume.Decisions { - Expect(decision).To(Equal(api.ToolApprovalDecision{ - ToolCallID: fmt.Sprintf("call-%02d", i+1), - Tool: fmt.Sprintf("example_tool_%02d", i+1), - Action: api.ToolApprovalApprove, - })) - } - }) - - It("rejects an approval response without its durable state before provider execution", func() { - fixture := newApprovalFixture(true) - fixture.assistant.Parts = fixture.assistant.Parts[:1] - provider := &fakeStreamingProvider{} - service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}}) - - response := httptest.NewRecorder() - service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ - Model: "anthropic/claude-opus-5", Messages: []aichat.UIMessage{fixture.user, fixture.assistant}, - })) - - Expect(response.Code).To(Equal(http.StatusBadRequest)) - Expect(response.Body.String()).To(ContainSubstring("durable tool approval state")) - Expect(provider.specs).To(BeEmpty()) - }) - - DescribeTable("rejects approval responses that do not match durable state", - func(mutate func(*approvalFixture), want string) { - fixture := newApprovalFixture(true, true) - mutate(&fixture) - provider := &fakeStreamingProvider{} - service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}}) - - response := httptest.NewRecorder() - service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ - Model: "anthropic/claude-opus-5", Messages: []aichat.UIMessage{fixture.user, fixture.assistant}, - })) - - Expect(response.Code).To(Equal(http.StatusBadRequest)) - Expect(response.Body.String()).To(ContainSubstring(want)) - Expect(provider.specs).To(BeEmpty()) - }, - Entry("approval id", func(fixture *approvalFixture) { - fixture.assistant.Parts[0].Approval.ID = "approval-other" - }, `tool call "call-01" approval ID is "approval-other"`), - Entry("tool name", func(fixture *approvalFixture) { - fixture.assistant.Parts[0].ToolName = "example_tool_other" - }, `approval names "example_tool_other"`), - Entry("tool input", func(fixture *approvalFixture) { - fixture.assistant.Parts[0].Input = json.RawMessage(`{"index":99}`) - }, `approval input does not match durable state`), - Entry("incomplete batch", func(fixture *approvalFixture) { - fixture.assistant.Parts[1].State = "approval-requested" - fixture.assistant.Parts[1].Approval.Approved = nil - }, `pending tool call "call-02" has no approval response`), - ) - - It("replaces a suspended thread message with mixed approved and denied results", func() { - fixture := newApprovalFixture(true, false) - fixture.assistant.Parts[1].Approval.Reason = "Keep the existing value." - store := aichat.NewMemoryThreadStore() - thread, err := store.Create(context.Background(), "Approval") - Expect(err).NotTo(HaveOccurred()) - Expect(store.AppendMessage(context.Background(), thread.ID, fixture.user)).To(Succeed()) - Expect(store.AppendMessage(context.Background(), thread.ID, suspendedAssistant(fixture))).To(Succeed()) - - provider := &fakeStreamingProvider{events: []api.Event{ - { - Kind: api.EventToolUse, ToolCallID: "call-01", Tool: "example_tool_01", - Input: map[string]any{"index": float64(1)}, - }, - { - Kind: api.EventToolResult, ToolCallID: "call-01", Tool: "example_tool_01", - Text: `{"updated":true}`, Success: true, - }, - {Kind: api.EventText, Text: "Finished."}, - {Kind: api.EventResult, Success: true, Model: "claude-opus-5"}, - }} - 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-approval", ThreadID: thread.ID, Model: "anthropic/claude-opus-5", - Messages: []aichat.UIMessage{fixture.user, fixture.assistant}, - })) - - Expect(response.Code).To(Equal(http.StatusOK), response.Body.String()) - Expect(response.Body.String()).To(ContainSubstring(`"type":"tool-output-denied","toolCallId":"call-02"`)) - Expect(provider.specs).To(HaveLen(1)) - Expect(provider.specs[0].ToolApproval.Decisions).To(Equal([]api.ToolApprovalDecision{ - {ToolCallID: "call-01", Tool: "example_tool_01", Action: api.ToolApprovalApprove}, - { - ToolCallID: "call-02", Tool: "example_tool_02", - Action: api.ToolApprovalDeny, Message: "Keep the existing value.", - }, - })) - - stored, err := store.Get(context.Background(), thread.ID) - Expect(err).NotTo(HaveOccurred()) - Expect(stored.Messages).To(HaveLen(2)) - assistant := stored.Messages[1] - Expect(assistant.ID).To(Equal("message-assistant")) - Expect(assistant.Parts[0].State).To(Equal("output-available")) - Expect(assistant.Parts[0].Output).To(MatchJSON(`{"updated":true}`)) - Expect(assistant.Parts[1].State).To(Equal("output-denied")) - Expect(assistant.Parts).To(ContainElement(SatisfyAll( - HaveField("Type", "text"), - HaveField("Text", "Finished."), - ))) - - nextProvider := &fakeStreamingProvider{events: []api.Event{ - {Kind: api.EventText, Text: "Ready."}, - {Kind: api.EventResult, Success: true, Model: "claude-opus-5"}, - }} - nextService := aichat.NewService(aichat.ServiceOptions{ - Resolver: &fakeResolver{provider: nextProvider}, Threads: store, - }) - nextMessages := append([]aichat.UIMessage(nil), stored.Messages...) - nextMessages = append(nextMessages, aichat.UIMessage{ - ID: "message-next", Role: "user", - Parts: []aichat.UIPart{{Type: "text", Text: "What happened?"}}, - }) - nextResponse := httptest.NewRecorder() - nextService.Handler().ServeHTTP(nextResponse, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ - ID: "chat-next", ThreadID: thread.ID, Model: "anthropic/claude-opus-5", - Messages: nextMessages, - })) - Expect(nextResponse.Code).To(Equal(http.StatusOK), nextResponse.Body.String()) - Expect(nextProvider.specs).To(HaveLen(1)) - Expect(nextProvider.specs[0].ToolApproval).To(BeNil()) - }) - - It("rejects unresolved tool parts outside an approval resume", func() { - provider := &fakeStreamingProvider{} - service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}}) - response := httptest.NewRecorder() - service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ - Model: "anthropic/claude-opus-5", - Messages: []aichat.UIMessage{ - {Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Continue."}}}, - {Role: "assistant", Parts: []aichat.UIPart{{ - Type: "dynamic-tool", ToolName: "example_tool", ToolCallID: "call-pending", - State: "approval-requested", Input: json.RawMessage(`{"id":"record-1"}`), - }}}, - }, - })) - - Expect(response.Code).To(Equal(http.StatusBadRequest)) - Expect(response.Body.String()).To(ContainSubstring(`tool call "call-pending" is not terminal`)) - Expect(provider.specs).To(BeEmpty()) - }) -}) diff --git a/pkg/aichat/database_threads.go b/pkg/aichat/database_threads.go new file mode 100644 index 00000000..8ca872e4 --- /dev/null +++ b/pkg/aichat/database_threads.go @@ -0,0 +1,427 @@ +package aichat + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" +) + +type DatabaseThreadStore struct { + db *database.DB +} + +func NewDatabaseThreadStore(db *database.DB) (*DatabaseThreadStore, error) { + if db == nil || db.Gorm() == nil { + return nil, fmt.Errorf("captain chat session store requires a database") + } + return &DatabaseThreadStore{db: db}, nil +} + +func (s *DatabaseThreadStore) Create(ctx context.Context, title string) (*Thread, error) { + metadata := map[string]any{"aichat": true} + // A caller who names a thread up front owns that name, so later automatic + // naming leaves it alone. + if strings.TrimSpace(title) != "" { + metadata[database.SessionTitleSourceKey] = string(database.SessionTitleUser) + } + record, err := s.db.CreateOrGetSession(ctx, database.CreateSessionInput{ + ID: uuid.New(), Source: "aichat", Provider: "captain", HostID: "local", + Title: strings.TrimSpace(title), Metadata: metadata, + }) + if err != nil { + return nil, err + } + return s.Get(ctx, record.ID.String()) +} + +func (s *DatabaseThreadStore) List(ctx context.Context) ([]*Thread, error) { + overviews, err := s.db.ListSessionOverviews(ctx, database.SessionOverviewFilter{Source: "aichat", RootsOnly: true}) + if err != nil { + return nil, err + } + threads := make([]*Thread, len(overviews)) + for i := range overviews { + aggregate, err := s.getSession(ctx, overviews[i]) + if err != nil { + return nil, err + } + threads[i] = threadFromSession(aggregate, overviews[i]) + } + return threads, nil +} + +func (s *DatabaseThreadStore) Get(ctx context.Context, id string) (*Thread, error) { + overview, err := s.getOverview(ctx, id) + if err != nil { + return nil, err + } + if overview.Source != "aichat" { + return nil, fmt.Errorf("captain chat session %s has source %q", overview.ID, overview.Source) + } + aggregate, err := s.getSession(ctx, *overview) + if err != nil { + return nil, err + } + return threadFromSession(aggregate, *overview), nil +} + +func (s *DatabaseThreadStore) GetSession(ctx context.Context, id string) (*session.Session, error) { + overview, err := s.getOverview(ctx, id) + if err != nil { + return nil, err + } + return s.getSession(ctx, *overview) +} + +func (s *DatabaseThreadStore) getOverview(ctx context.Context, id string) (*database.SessionOverview, error) { + parsed, err := uuid.Parse(strings.TrimSpace(id)) + if err != nil { + return nil, fmt.Errorf("captain chat session ID %q is not a UUID: %w", id, err) + } + overview, err := s.db.GetSessionOverviewByIdentity(ctx, parsed.String()) + if err != nil { + return nil, err + } + if overview.ID != parsed { + return nil, fmt.Errorf("captain session %s resolved to %s", parsed, overview.ID) + } + return overview, nil +} + +func (s *DatabaseThreadStore) getSession(ctx context.Context, overview database.SessionOverview) (*session.Session, error) { + messages, err := s.db.ListTranscriptMessages(ctx, database.TranscriptPage{SessionID: overview.ID}) + if err != nil { + return nil, err + } + turns, err := s.db.ListThreadTurns(ctx, overview.ID) + if err != nil { + return nil, err + } + requests, err := s.db.ListTurnRequests(ctx, database.TurnRequestFilter{SessionID: overview.ID}) + if err != nil { + return nil, err + } + costs, err := s.db.ListThreadCosts(ctx, overview.ID) + if err != nil { + return nil, err + } + agents, err := s.db.ListThreadAgents(ctx, overview.ID) + if err != nil { + return nil, err + } + aggregate := sessionFromOverview(overview) + aggregate.Root, aggregate.Agents = projectSessionAgents(agents) + if err := ApplyOverviewProjection(ctx, s.db, overview, aggregate); err != nil { + return nil, err + } + // The overview's own usage/cost is root-scoped; a thread's subagents spend + // against the same conversation, so roll the thread-wide figures on top. + applyThreadCosts(aggregate, costs) + aggregate.Messages, err = projectSessionMessages(messages) + if err != nil { + return nil, err + } + aggregate.Turns = projectSessionTurns(turns, aggregate.Messages) + aggregate.Requests, err = projectSessionRequests(requests) + if err != nil { + return nil, err + } + applyRequestState(aggregate) + return aggregate, nil +} + +func (s *DatabaseThreadStore) AppendMessage(ctx context.Context, id string, message UIMessage) error { + return s.putMessage(ctx, id, message, false) +} + +func (s *DatabaseThreadStore) ReplaceLastMessage(ctx context.Context, id string, message UIMessage) error { + thread, err := s.Get(ctx, id) + if err != nil { + return err + } + if err := validateLastMessageReplacement(thread.Messages, message); err != nil { + return err + } + return s.putMessage(ctx, id, message, true) +} + +func (s *DatabaseThreadStore) putMessage(ctx context.Context, id string, message UIMessage, replace bool) error { + sessionID, err := uuid.Parse(id) + if err != nil { + return fmt.Errorf("captain chat session ID %q is not a UUID: %w", id, err) + } + turnID, err := uuid.Parse(message.TurnID) + if err != nil { + return fmt.Errorf("captain chat message %q turn ID %q is not a UUID: %w", message.ID, message.TurnID, err) + } + parts, err := json.Marshal(message.Parts) + if err != nil { + return fmt.Errorf("encode Captain chat message %q: %w", message.ID, err) + } + return s.db.PutChatMessage(ctx, database.PutChatMessageInput{ + SessionID: sessionID, TurnID: turnID, ProviderMessageID: message.ID, + Role: message.Role, Parts: parts, Replace: replace, + }) +} + +func (s *DatabaseThreadStore) Delete(ctx context.Context, id string) error { + parsed, err := uuid.Parse(id) + if err != nil { + return fmt.Errorf("captain chat session ID %q is not a UUID: %w", id, err) + } + return s.db.DeleteChatSession(ctx, parsed) +} + +func (s *DatabaseThreadStore) SetProviderSession(ctx context.Context, id, providerSessionID string) error { + parsed, err := uuid.Parse(id) + if err != nil { + return fmt.Errorf("captain chat session ID %q is not a UUID: %w", id, err) + } + record, err := s.db.GetSession(ctx, parsed) + if err != nil { + return err + } + providerSessionID = strings.TrimSpace(providerSessionID) + if providerSessionID == "" { + return fmt.Errorf("provider session ID cannot be empty") + } + if record.ProviderSessionID != "" { + if record.ProviderSessionID == providerSessionID { + return nil + } + return fmt.Errorf("provider session ID is already bound to %q", record.ProviderSessionID) + } + _, err = s.db.UpdateSessionState(ctx, database.UpdateSessionStateInput{ + ID: parsed, ExpectedVersion: record.StateVersion, ProviderSessionID: &providerSessionID, + }) + return err +} + +func (s *DatabaseThreadStore) SetTitle(ctx context.Context, id string, update TitleUpdate) error { + parsed, err := uuid.Parse(id) + if err != nil { + return fmt.Errorf("captain chat session ID %q is not a UUID: %w", id, err) + } + title, err := normalizeTitle(update) + if err != nil { + return err + } + _, err = s.db.SetSessionTitle(ctx, database.SetSessionTitleInput{ + ID: parsed, Title: title, Source: database.SessionTitleSource(update.Source), + }) + return err +} + +func (s *DatabaseThreadStore) AddUsage(ctx context.Context, id string, _ TurnUsage) (*Thread, error) { + return s.Get(ctx, id) +} + +func sessionFromOverview(overview database.SessionOverview) *session.Session { + detail := &session.Session{ + ID: overview.ID.String(), ProviderSessionID: stringPointer(overview.ProviderSessionID), Revision: overview.StateVersion, + LifecycleStatus: overview.LifecycleStatus, ActivityState: overview.ActivityState, + HealthState: overview.HealthState, StateReason: stringPointer(overview.StateReason), + Source: overview.Source, Project: stringPointer(overview.Project), CWD: stringPointer(overview.CWD), + Slug: stringPointer(overview.Slug), Title: stringPointer(overview.Title), InitialPrompt: stringPointer(overview.InitialPrompt), + Version: stringPointer(overview.CLIVersion), Provider: overview.Provider, Backend: stringPointer(overview.Backend), + Model: stringPointer(overview.Model), ReasoningEffort: stringPointer(overview.Effort), + ExecutionMode: api.RuntimeMode(overview.ExecutionMode), + HistoryFile: stringPointer(overview.HistoryFile), StartedAt: overview.StartedAt, EndedAt: overview.EndedAt, + Usage: api.Usage{ + InputTokens: int(overview.InputTokens), OutputTokens: int(overview.OutputTokens), + ReasoningTokens: int(overview.ReasoningTokens), CacheReadTokens: int(overview.CacheReadTokens), + CacheWriteTokens: int(overview.CacheWriteTokens), + }, + // overview.CostUSD is the resolved total, which already falls back to list + // price per call — assigning it to ProviderCostUSD would make every + // reconstruction claim to be a billed figure. Carry both readings and let + // Cost.Total() resolve them the same way the view does. + Cost: api.Cost{ + Model: stringPointer(overview.Model), InputTokens: int(overview.InputTokens), OutputTokens: int(overview.OutputTokens), + ReasoningTokens: int(overview.ReasoningTokens), CacheReadTokens: int(overview.CacheReadTokens), + CacheWriteTokens: int(overview.CacheWriteTokens), TotalTokens: int(overview.TotalTokens), + InputCost: overview.InputCost, OutputCost: overview.OutputCost, ReasoningCost: overview.ReasoningCost, + CacheReadCost: overview.CacheReadCost, CacheWriteCost: overview.CacheWriteCost, + ProviderCostUSD: overview.ProviderCostUSD, + }, + } + return detail +} + +func projectSessionMessages(rows []database.TranscriptMessage) ([]session.Message, error) { + messages := make([]session.Message, len(rows)) + for i := range rows { + if err := json.Unmarshal(rows[i].Parts, &messages[i].Parts); err != nil { + return nil, fmt.Errorf("decode Captain message %s parts: %w", rows[i].ID, err) + } + messages[i].ID = stringPointer(rows[i].ProviderMessageID) + if messages[i].ID == "" { + messages[i].ID = rows[i].ID.String() + } + messages[i].Role = rows[i].Role + if rows[i].TurnID != nil { + messages[i].TurnID = rows[i].TurnID.String() + } + } + return messages, nil +} + +func projectSessionTurns(rows []database.SessionTurn, messages []session.Message) []session.Turn { + messageIDs := make(map[string][]string) + for _, message := range messages { + messageIDs[message.TurnID] = append(messageIDs[message.TurnID], message.ID) + } + turns := make([]session.Turn, len(rows)) + for i := range rows { + turns[i] = session.Turn{ + ID: rows[i].ID.String(), Status: rows[i].Status, Index: rows[i].TurnIndex, + StartedAt: rows[i].StartedAt, EndedAt: rows[i].EndedAt, + StopReason: stringPointer(rows[i].StopReason), Model: stringPointer(rows[i].Model), + Backend: stringPointer(rows[i].Backend), ReasoningEffort: stringPointer(rows[i].Effort), + MessageIDs: messageIDs[rows[i].ID.String()], + Usage: api.Usage{ + InputTokens: int(rows[i].InputTokens), OutputTokens: int(rows[i].OutputTokens), + ReasoningTokens: int(rows[i].ReasoningTokens), CacheReadTokens: int(rows[i].CacheReadTokens), + CacheWriteTokens: int(rows[i].CacheWriteTokens), + }, + Cost: api.Cost{ + Model: stringPointer(rows[i].Model), TotalTokens: int(rows[i].TotalTokens), + InputCost: rows[i].InputCost, OutputCost: rows[i].OutputCost, ReasoningCost: rows[i].ReasoningCost, + CacheReadCost: rows[i].CacheReadCost, CacheWriteCost: rows[i].CacheWriteCost, + ProviderCostUSD: rows[i].ProviderCostUSD, + }, + } + } + return turns +} + +func projectSessionRequests(rows []database.TurnRequest) ([]session.Request, error) { + requests := make([]session.Request, len(rows)) + for i := range rows { + input, err := json.Marshal(rows[i].Request["input"]) + if err != nil { + return nil, fmt.Errorf("encode Captain request %s input: %w", rows[i].ID, err) + } + updatedInput, err := json.Marshal(rows[i].Response["updatedInput"]) + if err != nil { + return nil, fmt.Errorf("encode Captain request %s updated input: %w", rows[i].ID, err) + } + requests[i] = session.Request{ + ID: rows[i].ID.String(), ToolCallID: rows[i].ToolCallID, Kind: rows[i].Kind, State: string(rows[i].State), + Tool: fmt.Sprint(rows[i].Request["tool"]), Input: input, RequestedBy: rows[i].RequestedBy, + ResolvedBy: rows[i].ResolvedBy, Reason: rows[i].Reason, Version: rows[i].Version, + ExpiresAt: rows[i].ExpiresAt, CreatedAt: rows[i].CreatedAt, ResolvedAt: rows[i].ResolvedAt, + } + if rows[i].Response != nil && rows[i].Response["updatedInput"] != nil { + requests[i].UpdatedInput = updatedInput + } + if rows[i].TurnID != nil { + requests[i].TurnID = rows[i].TurnID.String() + } + if rows[i].PromptRunID != nil { + requests[i].PromptRunID = rows[i].PromptRunID.String() + } + if rows[i].ModelCallID != nil { + requests[i].ModelCallID = rows[i].ModelCallID.String() + } + } + return requests, nil +} + +func applyRequestState(aggregate *session.Session) { + byID := make(map[string]session.Request, len(aggregate.Requests)) + for _, request := range aggregate.Requests { + byID[request.ID] = request + switch request.State { + case string(database.TurnRequestStateApproved): + aggregate.Approvals.Approved++ + case string(database.TurnRequestStateDenied): + aggregate.Approvals.Denied++ + aggregate.Approvals.Denials = append(aggregate.Approvals.Denials, session.Denial{ + ToolUseID: request.ToolCallID, Tool: request.Tool, Reason: request.Reason, + }) + } + } + for i := range aggregate.Messages { + for j := range aggregate.Messages[i].Parts { + part := &aggregate.Messages[i].Parts[j] + if part.Approval == nil { + continue + } + request, ok := byID[part.Approval.ID] + if !ok { + continue + } + switch request.State { + case string(database.TurnRequestStatePending): + part.State = session.ToolStateApprovalRequested + case string(database.TurnRequestStateApproved): + approved := true + if part.State == session.ToolStateApprovalRequested || part.State == session.ToolStateApprovalResponded { + part.State = session.ToolStateApprovalResponded + } + part.Approval.Approved = &approved + part.Approval.Reason = request.Reason + case string(database.TurnRequestStateDenied): + approved := false + part.State = session.ToolStateOutputDenied + part.Approval.Approved = &approved + part.Approval.Reason = request.Reason + case string(database.TurnRequestStateCancelled): + approved := false + part.State = session.ToolStateOutputDenied + part.Approval.Approved = &approved + part.Approval.Reason = request.Reason + } + } + } +} + +func threadFromSession(aggregate *session.Session, overview database.SessionOverview) *Thread { + messages := make([]UIMessage, len(aggregate.Messages)) + for i := range aggregate.Messages { + parts := make([]UIPart, len(aggregate.Messages[i].Parts)) + for j := range aggregate.Messages[i].Parts { + part := aggregate.Messages[i].Parts[j] + parts[j] = UIPart{ + Type: part.Type, Text: part.Text, MediaType: part.MediaType, URL: part.URL, Filename: part.Filename, + AttachmentID: part.AttachmentID, ToolName: part.ToolName, ToolCallID: part.ToolCallID, + State: part.State, Input: part.Input, Output: part.Output, ErrorText: part.ErrorText, Data: part.Data, + } + if part.Approval != nil { + parts[j].Approval = &Approval{ID: part.Approval.ID, Approved: part.Approval.Approved, Reason: part.Approval.Reason} + } + } + messages[i] = UIMessage{ + ID: aggregate.Messages[i].ID, Role: aggregate.Messages[i].Role, + Parts: parts, TurnID: aggregate.Messages[i].TurnID, + } + } + return &Thread{ + ID: aggregate.ID, Title: aggregate.Title, CreatedAt: overview.CreatedAt, UpdatedAt: overview.UpdatedAt, + Messages: messages, TotalInputTokens: aggregate.Usage.InputTokens, TotalOutputTokens: aggregate.Usage.OutputTokens, + TotalReasoningTokens: aggregate.Usage.ReasoningTokens, TotalCacheReadTokens: aggregate.Usage.CacheReadTokens, + TotalCacheWriteTokens: aggregate.Usage.CacheWriteTokens, TotalCostUSD: aggregate.Cost.Total(), + LastContextTokens: intPointer(overview.ContextTokens), ProviderSessionID: aggregate.ProviderSessionID, + } +} + +func stringPointer(value *string) string { + if value == nil { + return "" + } + return *value +} + +func intPointer(value *int64) int { + if value == nil { + return 0 + } + return int(*value) +} diff --git a/pkg/aichat/database_threads_integration_test.go b/pkg/aichat/database_threads_integration_test.go new file mode 100644 index 00000000..e2ca822d --- /dev/null +++ b/pkg/aichat/database_threads_integration_test.go @@ -0,0 +1,288 @@ +package aichat_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "time" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/flanksource/commons-db/dbtest" + "github.com/google/uuid" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Database chat sessions", func() { + It("keeps provider session identity immutable", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_provider_identity"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + store, err := aichat.NewDatabaseThreadStore(db) + Expect(err).NotTo(HaveOccurred()) + thread, err := store.Create(ctx, "Provider identity") + Expect(err).NotTo(HaveOccurred()) + + Expect(store.SetProviderSession(ctx, thread.ID, "provider-session-1")).To(Succeed()) + bound, err := store.GetSession(ctx, thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(store.SetProviderSession(ctx, thread.ID, "provider-session-1")).To(Succeed()) + replayed, err := store.GetSession(ctx, thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(replayed.Revision).To(Equal(bound.Revision)) + Expect(store.SetProviderSession(ctx, thread.ID, "provider-session-2")).To(MatchError(ContainSubstring( + `provider session ID is already bound to "provider-session-1"`, + ))) + }) + + It("projects messages, turns, usage, and durable approvals from one Captain session", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_session_projection"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + store, err := aichat.NewDatabaseThreadStore(db) + Expect(err).NotTo(HaveOccurred()) + thread, err := store.Create(ctx, "Accounts") + Expect(err).NotTo(HaveOccurred()) + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + execution, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: thread.ID, RequestID: "user-message-1", Title: thread.Title, + Spec: api.Spec{Model: api.Model{Name: "gemini", Backend: api.BackendGemini}.Capabilities()}, + }) + Expect(err).NotTo(HaveOccurred()) + + user := aichat.UIMessage{ + ID: "user-message-1", TurnID: execution.TurnID(), Role: "user", + Parts: []aichat.UIPart{{Type: "text", Text: "Edit the account"}}, + } + Expect(store.AppendMessage(ctx, thread.ID, user)).To(Succeed()) + permission, err := execution.Observe(ctx, api.Event{ + Kind: api.EventPermission, ToolCallID: "call-account-1", Tool: "accounts_edit", + Input: map[string]any{"id": "acc-1"}, + }) + Expect(err).NotTo(HaveOccurred()) + assistant := aichat.UIMessage{ + ID: execution.TurnID() + "-assistant", TurnID: execution.TurnID(), Role: "assistant", + Parts: []aichat.UIPart{{ + Type: "dynamic-tool", ToolName: "accounts_edit", ToolCallID: "call-account-1", + State: "approval-requested", Input: json.RawMessage(`{"id":"acc-1"}`), + Approval: &aichat.Approval{ID: permission.ApprovalID}, + }}, + } + Expect(store.AppendMessage(ctx, thread.ID, assistant)).To(Succeed()) + + aggregate, err := store.GetSession(ctx, thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(aggregate.ID).To(Equal(thread.ID)) + Expect(aggregate.Revision).To(BeNumerically(">", 0)) + Expect(aggregate.Messages).To(HaveLen(2)) + Expect(aggregate.Messages[0].TurnID).To(Equal(execution.TurnID())) + Expect(aggregate.Turns).To(HaveLen(1)) + Expect(aggregate.Requests).To(HaveLen(1)) + Expect(aggregate.Requests[0].ID).To(Equal(permission.ApprovalID)) + Expect(aggregate.Requests[0].TurnID).To(Equal(execution.TurnID())) + Expect(aggregate.Requests[0].PromptRunID).To(Equal(execution.PromptRunID())) + Expect(aggregate.Requests[0].ToolCallID).To(Equal("call-account-1")) + Expect(aggregate.Requests[0].Kind).To(Equal("tool_approval")) + Expect(aggregate.Requests[0].State).To(Equal("pending")) + Expect(aggregate.Requests[0].Tool).To(Equal("accounts_edit")) + Expect(aggregate.Requests[0].Input).To(MatchJSON(`{"id":"acc-1"}`)) + Expect(aggregate.Requests[0].RequestedBy).To(Equal("provider")) + Expect(aggregate.Messages[1].Parts[0].Approval.ID).To(Equal(permission.ApprovalID)) + + continuation, err := authority.ResolveToolApproval(ctx, aichat.ToolApprovalResolution{ + ThreadID: thread.ID, ApprovalID: permission.ApprovalID, Approved: false, Reason: "not now", + }) + Expect(err).NotTo(HaveOccurred()) + Expect(continuation).To(BeNil()) + resolved, err := store.GetSession(ctx, thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(resolved.Messages[1].Parts[0].State).To(Equal(session.ToolStateOutputDenied)) + Expect(*resolved.Messages[1].Parts[0].Approval.Approved).To(BeFalse()) + Expect(resolved.Messages[1].Parts[0].Approval.Reason).To(Equal("not now")) + Expect(resolved.Requests[0].ResolvedAt).NotTo(BeNil()) + Expect(*resolved.Requests[0].ResolvedAt).To(BeTemporally("~", time.Now(), time.Minute)) + }) + + It("resumes the provider after the final durable approval without another chat request", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_server_approval_resume"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + store, err := aichat.NewDatabaseThreadStore(db) + Expect(err).NotTo(HaveOccurred()) + thread, err := store.Create(ctx, "Accounts") + Expect(err).NotTo(HaveOccurred()) + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + + provider := &fakeStreamingProvider{backend: api.BackendGemini} + provider.execute = func(_ context.Context, spec api.Spec) (<-chan api.Event, error) { + var events []api.Event + if spec.ToolApproval == nil { + events = []api.Event{ + {Kind: api.EventToolUse, ToolCallID: "call-account-1", Tool: "accounts_edit", Input: map[string]any{"id": "acc-1"}}, + {Kind: api.EventPermission, ToolCallID: "call-account-1", Tool: "accounts_edit"}, + {Kind: api.EventResult, Success: true, ToolApproval: &api.ToolApprovalState{ + Messages: []api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Edit the account"}}}, + {Role: api.RoleAssistant, Parts: []api.Part{{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: "call-account-1", Name: "accounts_edit", Input: json.RawMessage(`{"id":"acc-1"}`), + }}}}, + }, + Calls: []api.ToolApprovalCall{{Request: api.ToolApprovalRequest{ + ToolCallID: "call-account-1", Tool: "accounts_edit", Input: json.RawMessage(`{"id":"acc-1"}`), + }}}, + ProviderCheckpoint: &api.ProviderCheckpoint{ + Codec: "test-provider", Version: 1, Payload: []byte("private provider state"), + }, + }}, + } + } else { + Expect(spec.ToolApproval.Decisions).To(HaveLen(1)) + Expect(spec.ToolApproval.State.ProviderCheckpoint.Payload).To(Equal([]byte("private provider state"))) + events = []api.Event{ + {Kind: api.EventToolUse, ToolCallID: "call-account-1", Tool: "accounts_edit", Input: map[string]any{"id": "acc-1"}}, + {Kind: api.EventToolResult, ToolCallID: "call-account-1", Tool: "accounts_edit", Success: true, Text: `{"updated":true}`}, + {Kind: api.EventText, Text: "Updated."}, + {Kind: api.EventResult, Success: true}, + } + } + stream := make(chan api.Event, len(events)) + for _, event := range events { + stream <- event + } + close(stream) + return stream, nil + } + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), Authority: authority, + Tools: aichat.StaticToolProvider([]api.ToolDefinition{{ + Name: "accounts_edit", DefaultPermission: api.ToolModeAsk, + Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}), + }) + + initial := httptest.NewRecorder() + service.Handler().ServeHTTP(initial, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "gemini", Backend: api.BackendGemini}, + Messages: []aichat.UIMessage{{ + ID: "user-message-1", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Edit the account"}}, + }}, + })) + Expect(initial.Code).To(Equal(http.StatusOK), initial.Body.String()) + suspended, err := store.GetSession(ctx, thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(suspended.Requests).To(HaveLen(1)) + Expect(suspended.Requests[0].State).To(Equal("pending")) + + approval := httptest.NewRecorder() + service.Handler().ServeHTTP(approval, requestJSON( + http.MethodPost, + "/api/chat/sessions/"+thread.ID+"/approvals/"+suspended.Requests[0].ID, + map[string]any{"approved": true}, + )) + Expect(approval.Code).To(Equal(http.StatusOK), approval.Body.String()) + var aggregate session.Session + Expect(json.Unmarshal(approval.Body.Bytes(), &aggregate)).To(Succeed()) + Expect(provider.specs).To(HaveLen(2)) + Expect(aggregate.Messages).To(HaveLen(2)) + Expect(aggregate.Messages[1].Parts[0].State).To(Equal(session.ToolStateOutputAvailable)) + Expect(aggregate.Messages[1].Parts[0].Output).To(MatchJSON(`{"updated":true}`)) + Expect(aggregate.Messages[1].Parts).NotTo(ContainElement(HaveField("Type", "data-tool-approval"))) + Expect(aggregate.Requests[0].State).To(Equal("approved")) + Expect(aggregate.Turns).To(HaveLen(1)) + Expect(aggregate.Turns[0].StopReason).To(Equal("stop")) + }) + + It("terminalizes a tool and preserves an approval persistence failure", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_approval_failure"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + Expect(db.Gorm().WithContext(ctx).Exec(` + ALTER TABLE captain_turn_requests + DROP CONSTRAINT captain_turn_requests_tool_approval_identity + `).Error).To(Succeed()) + Expect(db.Gorm().WithContext(ctx).Exec(` + ALTER TABLE captain_turn_requests + ADD CONSTRAINT captain_turn_requests_tool_approval_identity + CHECK (kind <> 'tool_approval' OR credential_id IS NOT NULL) + `).Error).To(Succeed()) + + store, err := aichat.NewDatabaseThreadStore(db) + Expect(err).NotTo(HaveOccurred()) + thread, err := store.Create(ctx, "Approval failure") + Expect(err).NotTo(HaveOccurred()) + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + provider := &fakeStreamingProvider{backend: api.BackendGemini, events: []api.Event{ + {Kind: api.EventToolUse, ToolCallID: "call-account-1", Tool: "accounts_edit", Input: map[string]any{"id": "acc-1"}}, + {Kind: api.EventPermission, ToolCallID: "call-account-1", Tool: "accounts_edit"}, + }} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), Authority: authority, + Tools: aichat.StaticToolProvider([]api.ToolDefinition{{ + Name: "accounts_edit", DefaultPermission: api.ToolModeAsk, + Handler: func(context.Context, map[string]any) (any, error) { + Fail("failed approval must not execute its tool") + return nil, nil + }, + }}), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "gemini", Backend: api.BackendGemini}, + Messages: []aichat.UIMessage{{ + ID: "user-approval-failure", Role: "user", + Parts: []aichat.UIPart{{Type: "text", Text: "Edit the account"}}, + }}, + })) + + Expect(response.Code).To(Equal(http.StatusOK), response.Body.String()) + parts := decodedDataLines(response.Body.String()) + Expect(partTypes(parts)).To(Equal([]string{ + "start", "start-step", "tool-input-available", "tool-output-error", + "error", "finish-step", "finish", + })) + Expect(parts[3]["errorText"]).To(ContainSubstring("captain_turn_requests_tool_approval_identity"), "wire tool error") + Expect(parts[4]["errorText"]).To(ContainSubstring("captain_turn_requests_tool_approval_identity"), "wire stream error") + + aggregate, err := store.GetSession(ctx, thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(aggregate.LifecycleStatus).To(Equal(string(database.SessionLifecycleFailed))) + Expect(aggregate.Messages).To(HaveLen(2)) + Expect(aggregate.Messages[1].Parts[0].State).To(Equal(session.ToolStateOutputError)) + Expect(aggregate.Messages[1].Parts[0].ErrorText).To(ContainSubstring("captain_turn_requests_tool_approval_identity"), "persisted tool error") + Expect(aggregate.Requests).To(BeEmpty()) + Expect(aggregate.Turns).To(HaveLen(1)) + var turnError string + Expect(db.Gorm().WithContext(ctx).Table("captain_turns"). + Select("error").Where("id = ?", aggregate.Turns[0].ID). + Row().Scan(&turnError)).To(Succeed()) + Expect(turnError).To(ContainSubstring("captain_turn_requests_tool_approval_identity"), "turn error") + + sessionID := uuid.MustParse(thread.ID) + runs, err := db.ListPromptRuns(ctx, database.PromptRunFilter{SessionID: &sessionID}) + Expect(err).NotTo(HaveOccurred()) + Expect(runs).To(HaveLen(1)) + Expect(runs[0].State).To(Equal(database.PromptRunStateFailed)) + Expect(runs[0].Error).To(ContainSubstring("captain_turn_requests_tool_approval_identity"), "prompt run error") + var modelCallError string + Expect(db.Gorm().WithContext(ctx).Table("captain_model_calls"). + Select("error").Where("prompt_run_id = ?", runs[0].ID). + Scan(&modelCallError).Error).To(Succeed()) + Expect(modelCallError).To(ContainSubstring("captain_turn_requests_tool_approval_identity"), "model call error") + }) +}) diff --git a/pkg/aichat/events.go b/pkg/aichat/events.go index 588b43f1..d13aac51 100644 --- a/pkg/aichat/events.go +++ b/pkg/aichat/events.go @@ -17,11 +17,13 @@ type toolState struct { type eventStream struct { writer *SSEWriter + messageID string blockType string blockID string nextBlock int tools map[string]toolState metadata *MessageMetadata + costs *TurnCosts sessionID string model string terminal bool @@ -31,6 +33,10 @@ type eventStream struct { // EventStreamOptions carries request state needed to finish the resumed UI message. type EventStreamOptions struct { ToolApproval *api.ToolApprovalResume + MessageID string + // Costs, when set, supplies the finish part's priced breakdown and the + // thread's cumulative cost. Nil for streams with no thread to accrue against. + Costs *TurnCosts } // WriteEventStream translates a Captain event channel into one complete UI Message Stream. @@ -42,7 +48,10 @@ func WriteEventStream(writer *SSEWriter, events <-chan api.Event, options EventS return fmt.Errorf("validate tool approval stream: %w", err) } } - stream := &eventStream{writer: writer, tools: map[string]toolState{}} + stream := &eventStream{ + writer: writer, messageID: options.MessageID, + tools: map[string]toolState{}, costs: options.Costs, + } if err := stream.start(); err != nil { return err } @@ -95,7 +104,7 @@ func (s *eventStream) approvalResults(resume *api.ToolApprovalResume) error { } func (s *eventStream) start() error { - if err := s.writer.WritePart(Part{Type: "start"}); err != nil { + if err := s.writer.WritePart(Part{Type: "start", MessageID: s.messageID}); err != nil { return err } return s.writer.WritePart(Part{Type: "start-step"}) @@ -128,11 +137,33 @@ func (s *eventStream) event(event api.Event) error { return s.result(event) case api.EventError: return s.providerError(event) + case api.EventInterrupted: + return s.interrupted(event) default: return fmt.Errorf("unsupported Captain event kind %q", event.Kind) } } +func (s *eventStream) interrupted(event api.Event) error { + if err := s.closeBlock(); err != nil { + return err + } + if err := s.terminalizeTools(event.Reason); err != nil { + return err + } + if err := s.writer.WritePart(Part{ + Type: "data-result", Data: map[string]any{"success": false, "interrupted": true}, + }); err != nil { + return err + } + success := false + s.metadata = &MessageMetadata{ + ProviderSessionID: s.sessionID, Model: s.model, Success: &success, Interrupted: true, + } + s.terminal = true + return nil +} + func (s *eventStream) delta(kind, delta string) error { if delta == "" { return nil @@ -189,10 +220,13 @@ func (s *eventStream) permission(event api.Event) error { if state.approvalRequested { return fmt.Errorf("duplicate permission for tool call %q", event.ToolCallID) } + if event.ApprovalID == "" { + return fmt.Errorf("permission for tool call %q has no durable approval ID", 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, + Type: "tool-approval-request", ApprovalID: event.ApprovalID, ToolCallID: event.ToolCallID, }) } @@ -242,7 +276,9 @@ func (s *eventStream) result(event api.Event) error { 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 { + if err := s.writer.WritePart(Part{ + Type: "data-result", Data: map[string]any{"success": event.Success, "waitingApproval": true}, + }); err != nil { return err } } else { @@ -265,12 +301,24 @@ func (s *eventStream) result(event api.Event) error { } if event.Usage != nil { s.metadata.Usage = usageMetadata(*event.Usage) - s.metadata.ContextTokens = event.Usage.InputTokens + s.metadata.ContextTokens = contextTokens(*event.Usage) + } + if s.costs != nil { + s.metadata.CostBreakdown = s.costs.Breakdown + s.metadata.ThreadCostUSD = s.costs.ThreadCostUSD } s.terminal = true return nil } +// contextTokens is the prompt's occupancy of the context window. The usage +// buckets are disjoint (pkg/api/cost.go), so the cached prefix counts too: +// agent backends report near-zero InputTokens against a six-figure cache read, +// and reporting input alone renders that as a context of single digits. +func contextTokens(usage api.Usage) int { + return usage.InputTokens + usage.CacheReadTokens + usage.CacheWriteTokens +} + func (s *eventStream) validateApprovalCorrelation(approval *api.ToolApprovalState) error { seen := make(map[string]bool) pending := approval.Pending() @@ -335,11 +383,33 @@ func (s *eventStream) providerError(event api.Event) error { if event.Error == "" { return fmt.Errorf("captain error event has no error message") } - s.tools = map[string]toolState{} + if err := s.terminalizeTools(event.Error); err != nil { + return err + } s.terminal = true return s.writer.WritePart(Part{Type: "error", ErrorText: event.Error}) } +func (s *eventStream) terminalizeTools(message string) error { + if message == "" { + message = "tool execution did not complete" + } + ids := make([]string, 0, len(s.tools)) + for id := range s.tools { + ids = append(ids, id) + } + sort.Strings(ids) + for _, id := range ids { + if err := s.writer.WritePart(Part{ + Type: "tool-output-error", ToolCallID: id, ErrorText: message, + }); err != nil { + return err + } + delete(s.tools, id) + } + return nil +} + func (s *eventStream) closeBlock() error { if s.blockType == "" { return nil @@ -397,11 +467,13 @@ func (s *eventStream) fail(cause error) error { if err := s.closeBlock(); err != nil { return err } + if err := s.terminalizeTools(cause.Error()); 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 } diff --git a/pkg/aichat/execution.go b/pkg/aichat/execution.go new file mode 100644 index 00000000..9a5a0304 --- /dev/null +++ b/pkg/aichat/execution.go @@ -0,0 +1,184 @@ +package aichat + +import ( + "context" + "errors" + "fmt" + + "github.com/flanksource/captain/pkg/api" +) + +// ExecutionRequest is the authoritative identity and resolved policy for one +// chat provider turn. +type ExecutionRequest struct { + ThreadID string + RequestID string + Title string + Spec api.Spec + Definitions []api.ToolDefinition +} + +// ToolApprovalResolution is the authenticated user's answer to one live +// caller-tool approval. +type ToolApprovalResolution struct { + ThreadID string + ApprovalID string + Approved bool + UpdatedInput map[string]any + Reason string +} + +type ApprovalContinuation struct { + Execution Execution + Spec api.Spec +} + +// ExecutionAuthority admits turns before provider launch and resolves approval +// decisions against the same durable identity. +type ExecutionAuthority interface { + Begin(context.Context, ExecutionRequest) (Execution, error) + ResolveToolApproval(context.Context, ToolApprovalResolution) (*ApprovalContinuation, error) +} + +// Execution is one admitted provider turn. Its caller-tool endpoint is already +// bound to the Captain session and prompt run. +type Execution interface { + CaptainSessionID() string + TurnID() string + PromptRunID() string + CallerTools() *api.CallerToolEndpoint + Events() <-chan api.Event + Observe(context.Context, api.Event) (api.Event, error) + Interrupt(context.Context, string) error + Close(context.Context) error +} + +func mergeExecutionEvents( + ctx context.Context, + provider <-chan api.Event, + approvals <-chan api.Event, + definitions []api.ToolDefinition, +) <-chan api.Event { + if approvals == nil { + return provider + } + askTools := make(map[string]bool) + for _, definition := range definitions { + askTools[definition.Name] = definition.NeedsApproval() + } + out := make(chan api.Event) + go func() { + defer close(out) + awaiting := make(map[string]bool) + pendingApprovals := make(map[string]api.Event) + deferred := make([]api.Event, 0) + send := func(event api.Event) bool { + select { + case out <- event: + return true + case <-ctx.Done(): + return false + } + } + flush := func() bool { + for _, event := range deferred { + if !send(event) { + return false + } + } + deferred = deferred[:0] + return true + } + for provider != nil || (len(awaiting) > 0 && approvals != nil) { + select { + case <-ctx.Done(): + return + case approval, ok := <-approvals: + if !ok { + approvals = nil + continue + } + if !awaiting[approval.ToolCallID] { + pendingApprovals[approval.ToolCallID] = approval + continue + } + if !send(approval) { + return + } + delete(awaiting, approval.ToolCallID) + if len(awaiting) == 0 && !flush() { + return + } + case event, ok := <-provider: + if !ok { + provider = nil + continue + } + if event.Kind == api.EventToolUse { + if !send(event) { + return + } + if askTools[event.Tool] { + awaiting[event.ToolCallID] = true + if approval, ok := pendingApprovals[event.ToolCallID]; ok { + if !send(approval) { + return + } + delete(pendingApprovals, event.ToolCallID) + delete(awaiting, event.ToolCallID) + } + } + continue + } + if len(awaiting) > 0 { + deferred = append(deferred, event) + continue + } + if !send(event) { + return + } + } + } + if len(awaiting) == 0 { + _ = flush() + } + }() + return out +} + +func observeExecutionEvents( + ctx context.Context, + execution Execution, + source <-chan api.Event, +) <-chan api.Event { + if execution == nil { + return source + } + out := make(chan api.Event) + go func() { + defer close(out) + for event := range source { + observed, err := execution.Observe(ctx, event) + if err != nil { + failure := api.Event{ + Kind: api.EventError, Error: err.Error(), Model: event.Model, SessionID: event.SessionID, + } + if event.Kind == api.EventError { + if event.Error != "" { + failure.Error = errors.Join(errors.New(event.Error), err).Error() + } + } else if terminal, terminalErr := execution.Observe(ctx, failure); terminalErr != nil { + failure.Error = errors.Join(err, fmt.Errorf("finish authoritative execution: %w", terminalErr)).Error() + } else { + failure = terminal + } + sendEvent(ctx, out, failure) + return + } + if !sendEvent(ctx, out, observed) { + return + } + } + }() + return out +} diff --git a/pkg/aichat/execution_authority_ginkgo_test.go b/pkg/aichat/execution_authority_ginkgo_test.go new file mode 100644 index 00000000..b7ea15c9 --- /dev/null +++ b/pkg/aichat/execution_authority_ginkgo_test.go @@ -0,0 +1,405 @@ +package aichat_test + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + . "github.com/onsi/gomega/gstruct" +) + +type fakeExecutionAuthority struct { + execution *fakeExecution + beginErr error + begins []aichat.ExecutionRequest + resolutions []aichat.ToolApprovalResolution +} + +func (f *fakeExecutionAuthority) Begin(_ context.Context, request aichat.ExecutionRequest) (aichat.Execution, error) { + f.begins = append(f.begins, request) + if f.beginErr != nil { + return nil, f.beginErr + } + f.execution.turnID = "turn-" + request.RequestID + return f.execution, nil +} + +func (f *fakeExecutionAuthority) ResolveToolApproval( + _ context.Context, + resolution aichat.ToolApprovalResolution, +) (*aichat.ApprovalContinuation, error) { + f.resolutions = append(f.resolutions, resolution) + return nil, nil +} + +type fakeExecution struct { + events chan api.Event + endpoint *api.CallerToolEndpoint + observed []api.Event + closed bool + turnID string + interrupts []string +} + +func (f *fakeExecution) Interrupt(_ context.Context, reason string) error { + f.interrupts = append(f.interrupts, reason) + return nil +} + +func (f *fakeExecution) CaptainSessionID() string { return "captain-session-1" } +func (f *fakeExecution) TurnID() string { + if f.turnID == "" { + return "turn-1" + } + return f.turnID +} +func (f *fakeExecution) PromptRunID() string { return "prompt-run-1" } +func (f *fakeExecution) CallerTools() *api.CallerToolEndpoint { + if f.endpoint == nil { + return nil + } + endpoint := *f.endpoint + return &endpoint +} +func (f *fakeExecution) Events() <-chan api.Event { return f.events } +func (f *fakeExecution) Observe(_ context.Context, event api.Event) (api.Event, error) { + if event.Kind == api.EventPermission && event.ApprovalID == "" { + event.ApprovalID = "0e5dc2fe-8b77-44e9-a3de-6a00298c8bde" + } + f.observed = append(f.observed, event) + return event, nil +} +func (f *fakeExecution) Close(context.Context) error { + f.closed = true + return nil +} + +var _ = Describe("Authoritative aichat execution", func() { + It("interrupts an active turn and finishes its stream without an error", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Interrupt") + Expect(err).NotTo(HaveOccurred()) + started := make(chan struct{}) + providerInterrupted := make(chan struct{}, 1) + provider := &fakeStreamingProvider{ + execute: func(ctx context.Context, _ api.Spec) (<-chan api.Event, error) { + events := make(chan api.Event, 2) + events <- api.Event{Kind: api.EventSystem, SessionID: "provider-session-1"} + events <- api.Event{Kind: api.EventText, Text: "partial"} + close(started) + go func() { + <-ctx.Done() + close(events) + }() + return events, nil + }, + interrupt: func(context.Context) error { + providerInterrupted <- struct{}{} + return nil + }, + } + execution := &fakeExecution{} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), + Authority: &fakeExecutionAuthority{execution: execution}, + }) + + chatResponse := httptest.NewRecorder() + chatDone := make(chan struct{}) + go func() { + defer close(chatDone) + service.Handler().ServeHTTP(chatResponse, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "sonnet", Backend: api.BackendClaudeAgent}, + Messages: []aichat.UIMessage{{ + ID: "user-message-interrupt", Role: "user", + Parts: []aichat.UIPart{{Type: "text", Text: "Start a long response"}}, + }}, + })) + }() + Eventually(started).Should(BeClosed()) + + interruptResponse := httptest.NewRecorder() + service.Handler().ServeHTTP(interruptResponse, httptest.NewRequest( + http.MethodPost, "/api/chat/sessions/"+thread.ID+"/interrupt", nil, + )) + Expect(interruptResponse.Code).To(Equal(http.StatusOK), interruptResponse.Body.String()) + Eventually(providerInterrupted).Should(Receive()) + Eventually(chatDone).Should(BeClosed()) + Expect(execution.interrupts).To(Equal([]string{"user"})) + Expect(chatResponse.Body.String()).To(ContainSubstring(`"interrupted":true`)) + Expect(chatResponse.Body.String()).NotTo(ContainSubstring(`"type":"error"`)) + + second := httptest.NewRecorder() + service.Handler().ServeHTTP(second, httptest.NewRequest( + http.MethodPost, "/api/chat/sessions/"+thread.ID+"/interrupt", nil, + )) + Expect(second.Code).To(Equal(http.StatusConflict)) + }) + + It("injects the run-bound endpoint and merges its live approval event", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Accounts") + Expect(err).NotTo(HaveOccurred()) + execution := &fakeExecution{ + events: make(chan api.Event, 1), + endpoint: &api.CallerToolEndpoint{ + Name: "captain", URL: "http://127.0.0.1:43210/mcp", + Headers: map[string]string{"Authorization": "Bearer secret"}, + }, + } + execution.events <- api.Event{ + Kind: api.EventPermission, Tool: "account_edit", + ToolCallID: "call-account-1", Input: map[string]any{"id": "acc-1"}, + } + close(execution.events) + authority := &fakeExecutionAuthority{execution: execution} + provider := &fakeStreamingProvider{ + backend: api.BackendClaudeAgent, + events: []api.Event{ + {Kind: api.EventToolUse, Tool: "account_edit", ToolCallID: "call-account-1", Input: map[string]any{"id": "acc-1"}}, + {Kind: api.EventToolResult, Tool: "account_edit", ToolCallID: "call-account-1", Text: `{"updated":true}`, Success: true}, + {Kind: api.EventResult, Success: true, SessionID: "provider-session-1"}, + }, + } + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), Authority: authority, + Tools: aichat.StaticToolProvider([]api.ToolDefinition{{ + Name: "account_edit", DefaultPermission: api.ToolModeAsk, + Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "sonnet", Backend: api.BackendClaudeAgent}, + Messages: []aichat.UIMessage{{ + ID: "user-message-1", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Edit the account"}}, + }}, + })) + + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(response.Body.String()).To(ContainSubstring(`"type":"tool-approval-request"`)) + Expect(authority.begins).To(HaveLen(1)) + Expect(authority.begins[0].ThreadID).To(Equal(thread.ID)) + Expect(authority.begins[0].RequestID).To(Equal("user-message-1")) + Expect(provider.specs).To(HaveLen(1)) + Expect(execution.observed).To(ContainElement( + MatchFields(IgnoreExtras, Fields{"Kind": Equal(api.EventResult)}), + )) + Expect(execution.closed).To(BeTrue()) + }) + + It("streams API-provider approvals without waiting for agent caller-tool events", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Accounts") + Expect(err).NotTo(HaveOccurred()) + calls := []api.ToolApprovalRequest{ + {ToolCallID: "call-account-1", Tool: "account_edit", Input: json.RawMessage(`{"id":"acc-1"}`)}, + {ToolCallID: "call-account-2", Tool: "account_edit", Input: json.RawMessage(`{"id":"acc-2"}`)}, + } + execution := &fakeExecution{events: make(chan api.Event)} + provider := &fakeStreamingProvider{backend: api.BackendGemini, events: []api.Event{ + {Kind: api.EventToolUse, Tool: "account_edit", ToolCallID: calls[0].ToolCallID, Input: map[string]any{"id": "acc-1"}}, + {Kind: api.EventToolUse, Tool: "account_edit", ToolCallID: calls[1].ToolCallID, Input: map[string]any{"id": "acc-2"}}, + {Kind: api.EventPermission, Tool: "account_edit", ToolCallID: calls[0].ToolCallID}, + {Kind: api.EventPermission, Tool: "account_edit", ToolCallID: calls[1].ToolCallID}, + {Kind: api.EventResult, Success: true, ToolApproval: pendingApprovalState(calls...)}, + }} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), + Authority: &fakeExecutionAuthority{execution: execution}, + Tools: aichat.StaticToolProvider([]api.ToolDefinition{{ + Name: "account_edit", DefaultPermission: api.ToolModeAsk, + Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "gemini", Backend: api.BackendGemini}, + Messages: []aichat.UIMessage{{ + ID: "user-message-approval", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Edit the accounts"}}, + }}, + })) + + parts := decodedDataLines(response.Body.String()) + Expect(response.Code).To(Equal(http.StatusOK), response.Body.String()) + Expect(partTypes(parts)).To(Equal([]string{ + "start", "start-step", "tool-input-available", "tool-input-available", + "tool-approval-request", "tool-approval-request", "data-result", "finish-step", "finish", + })) + Expect(parts[0]).To(HaveKeyWithValue("messageId", "turn-user-message-approval-assistant")) + stored, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(stored.Messages).To(HaveLen(2)) + Expect(stored.Messages[1].ID).To(Equal(parts[0]["messageId"])) + Expect(stored.Messages[1].Parts[0].State).To(Equal("approval-requested")) + Expect(stored.Messages[1].Parts[1].State).To(Equal("approval-requested")) + Expect(execution.observed).To(ContainElement(MatchFields(IgnoreExtras, Fields{"Kind": Equal(api.EventResult)}))) + }) + + It("admits sequential user messages in one Captain thread as distinct turns", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Accounts") + Expect(err).NotTo(HaveOccurred()) + execution := &fakeExecution{} + authority := &fakeExecutionAuthority{execution: execution} + provider := &fakeStreamingProvider{events: []api.Event{ + {Kind: api.EventText, Text: "Done."}, + {Kind: api.EventResult, Success: true}, + }} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), Authority: authority, + }) + + messages := []aichat.UIMessage{{ + ID: "user-message-1", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "List accounts"}}, + }} + first := httptest.NewRecorder() + service.Handler().ServeHTTP(first, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "gemini", Backend: api.BackendGemini}, Messages: messages, + })) + Expect(first.Code).To(Equal(http.StatusOK), first.Body.String()) + + persisted, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + messages = append(persisted.Messages, aichat.UIMessage{ + ID: "user-message-2", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "List contacts"}}, + }) + second := httptest.NewRecorder() + service.Handler().ServeHTTP(second, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "gemini", Backend: api.BackendGemini}, Messages: messages, + })) + Expect(second.Code).To(Equal(http.StatusOK), second.Body.String()) + + Expect(authority.begins).To(HaveLen(2)) + Expect(authority.begins[0].RequestID).To(Equal("user-message-1")) + Expect(authority.begins[1].RequestID).To(Equal("user-message-2")) + persisted, err = store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(persisted.Messages).To(HaveLen(4)) + Expect(persisted.Messages[1].ID).To(Equal("turn-user-message-1-assistant")) + Expect(persisted.Messages[3].ID).To(Equal("turn-user-message-2-assistant")) + }) + + It("regenerates the named assistant message without duplicating persisted history", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Accounts") + Expect(err).NotTo(HaveOccurred()) + user := aichat.UIMessage{ + ID: "user-message-1", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "List accounts"}}, + } + assistant := aichat.UIMessage{ + ID: "user-message-1-assistant", Role: "assistant", Parts: []aichat.UIPart{{Type: "text", Text: "Old answer"}}, + } + Expect(store.AppendMessage(context.Background(), thread.ID, user)).To(Succeed()) + Expect(store.AppendMessage(context.Background(), thread.ID, assistant)).To(Succeed()) + authority := &fakeExecutionAuthority{execution: &fakeExecution{}} + provider := &fakeStreamingProvider{events: []api.Event{ + {Kind: api.EventText, Text: "New answer"}, + {Kind: api.EventResult, Success: true}, + }} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), Authority: authority, + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "regenerate-message", MessageID: assistant.ID, + Runtime: &api.Model{Name: "gemini", Backend: api.BackendGemini}, Messages: []aichat.UIMessage{user}, + })) + + Expect(response.Code).To(Equal(http.StatusOK), response.Body.String()) + Expect(authority.begins).To(HaveLen(1)) + Expect(authority.begins[0].RequestID).To(Equal(assistant.ID)) + persisted, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(persisted.Messages).To(HaveLen(2)) + Expect(persisted.Messages[1].ID).To(Equal(assistant.ID)) + Expect(persisted.Messages[1].Parts).To(ContainElement(HaveField("Text", "New answer"))) + }) + + It("rejects a persisted chat id that differs from its Captain thread", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Accounts") + Expect(err).NotTo(HaveOccurred()) + service := aichat.NewService(aichat.ServiceOptions{Threads: aichat.FixedThreadStore(store)}) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: "different-chat", ThreadID: thread.ID, Trigger: "submit-message", + Messages: []aichat.UIMessage{{ID: "user-message-1", Role: "user"}}, + })) + + Expect(response.Code).To(Equal(http.StatusBadRequest)) + Expect(response.Body.String()).To(ContainSubstring("must match threadId")) + }) + + It("does not persist the user message when execution admission fails", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Accounts") + Expect(err).NotTo(HaveOccurred()) + authority := &fakeExecutionAuthority{beginErr: errors.New("duplicate prompt run")} + service := aichat.NewService(aichat.ServiceOptions{ + Threads: aichat.FixedThreadStore(store), Authority: authority, + Resolver: &fakeResolver{provider: &fakeStreamingProvider{backend: api.BackendGemini}}, + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", + Runtime: &api.Model{Name: "gemini", Backend: api.BackendGemini}, + Messages: []aichat.UIMessage{{ + ID: "user-message-1", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "List accounts"}}, + }}, + })) + + Expect(response.Code).To(Equal(http.StatusInternalServerError)) + persisted, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(persisted.Messages).To(BeEmpty()) + }) + + It("resolves an approval only after authorizing its thread", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Contacts") + Expect(err).NotTo(HaveOccurred()) + authority := &fakeExecutionAuthority{} + service := aichat.NewService(aichat.ServiceOptions{Threads: aichat.FixedThreadStore(store), Authority: authority}) + response := httptest.NewRecorder() + + approvalID := "0e5dc2fe-8b77-44e9-a3de-6a00298c8bde" + service.Handler().ServeHTTP(response, requestJSON( + http.MethodPost, + "/api/chat/sessions/"+thread.ID+"/approvals/"+approvalID, + map[string]any{"approved": true, "updatedInput": map[string]any{"name": "Acme"}}, + )) + + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(authority.resolutions).To(HaveLen(1)) + Expect(authority.resolutions[0].ThreadID).To(Equal(thread.ID)) + Expect(authority.resolutions[0].ApprovalID).To(Equal(approvalID)) + Expect(authority.resolutions[0].UpdatedInput).To(Equal(map[string]any{"name": "Acme"})) + + missing := httptest.NewRecorder() + service.Handler().ServeHTTP(missing, requestJSON( + http.MethodPost, + "/api/chat/sessions/missing/approvals/"+approvalID, + json.RawMessage(`{"approved":false}`), + )) + Expect(missing.Code).To(Equal(http.StatusNotFound)) + Expect(authority.resolutions).To(HaveLen(1)) + }) +}) diff --git a/pkg/aichat/execution_database.go b/pkg/aichat/execution_database.go new file mode 100644 index 00000000..014ddf34 --- /dev/null +++ b/pkg/aichat/execution_database.go @@ -0,0 +1,511 @@ +package aichat + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "time" + + "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/ai/callertools" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/google/uuid" +) + +const ( + callerToolApprovalTimeout = 5 * time.Minute + providerApprovalTimeout = 24 * time.Hour + approvalPollInterval = 100 * time.Millisecond +) + +type databaseExecution struct { + db *database.DB + ctx context.Context + session *database.Session + turn *database.ChatTurn + run *database.PromptRun + modelCallID uuid.UUID + model string + backend api.Backend + definitions []api.ToolDefinition + events chan api.Event + + mu sync.Mutex + finishMu sync.Mutex + credential *database.CallerToolCredential + runtime *callertools.Runtime + endpoint *api.CallerToolEndpoint + terminal bool + suspended bool + closed bool + providerID string + approvalIDs map[string]uuid.UUID + providerToolUses []api.Event + providerToolUseReady chan struct{} +} + +// finishModelCall persists a terminal model call with its priced cost breakdown. +// Pricing happens here, against the same model identity the call was created +// with, so the five per-bucket cost columns and the provider-reported total are +// both stored rather than the whole figure collapsing into output_cost. +func (e *databaseExecution) finishModelCall( + ctx context.Context, + status database.ModelCallStatus, + stopReason string, + event api.Event, +) error { + input := database.FinishChatModelCallInput{ + ID: e.modelCallID, Status: status, StopReason: stopReason, Event: event, + ContextWindowTokens: ai.ContextWindowFor(e.backend, e.model), + } + if event.Usage != nil { + cost := ai.PriceUsage(e.backend, e.model, *event.Usage, event.CostUSD) + input.Cost = &cost + } + return e.db.FinishChatModelCall(ctx, input) +} + +func (e *databaseExecution) CaptainSessionID() string { return e.session.ID.String() } +func (e *databaseExecution) TurnID() string { return e.turn.ID.String() } +func (e *databaseExecution) PromptRunID() string { return e.run.ID.String() } +func (e *databaseExecution) Events() <-chan api.Event { return e.events } + +func (e *databaseExecution) CallerTools() *api.CallerToolEndpoint { + e.mu.Lock() + defer e.mu.Unlock() + if e.endpoint == nil { + return nil + } + endpoint := *e.endpoint + endpoint.Headers = cloneStringValues(e.endpoint.Headers) + return &endpoint +} + +func (e *databaseExecution) startCallerTools(ctx context.Context, backend api.Backend) error { + var credentialID uuid.UUID + runtime, err := callertools.New(callertools.Options{ + Definitions: e.definitions, SessionID: e.session.ID.String(), + ApprovalTimeout: callerToolApprovalTimeout, + ValidateCredential: func(ctx context.Context) error { + if credentialID == uuid.Nil { + return fmt.Errorf("caller-tool credential has not been issued") + } + return e.db.ValidateCallerToolCredential(ctx, credentialID) + }, + CanUseTool: func(ctx context.Context, request api.PermissionRequest) (api.PermissionDecision, error) { + return e.requestApproval(ctx, credentialID, request) + }, + }) + if err != nil { + return err + } + policy := make(map[string]api.ToolMode, len(e.definitions)) + for _, definition := range e.definitions { + policy[definition.Name] = definition.DefaultPermission + } + credential, err := e.db.CreateCallerToolCredential(ctx, database.CreateCallerToolCredentialInput{ + SessionID: e.session.ID, PromptRunID: e.run.ID, Backend: backend, + SecretHash: runtime.CredentialHash(), Policy: policy, + }) + if err != nil { + _ = runtime.Close() + return err + } + credentialID = credential.ID + endpoint := runtime.Endpoint() + e.mu.Lock() + e.runtime = runtime + e.credential = credential + e.endpoint = &endpoint + e.mu.Unlock() + return nil +} + +func (e *databaseExecution) requestApproval( + ctx context.Context, + credentialID uuid.UUID, + request api.PermissionRequest, +) (api.PermissionDecision, error) { + if request.ToolUseIDGenerated { + toolUseID, err := e.claimProviderToolUse(ctx, request) + if err != nil { + return api.PermissionDecision{}, err + } + request.ToolUseID = toolUseID + } + expiresAt := time.Now().Add(callerToolApprovalTimeout) + pending, err := e.db.CreateToolApprovalRequest(ctx, database.CreateToolApprovalRequestInput{ + CredentialID: credentialID, SessionID: e.session.ID, TurnID: e.turn.ID, PromptRunID: e.run.ID, + ModelCallID: e.modelCallID, RequestedBy: "caller_tool", + ToolCallID: request.ToolUseID, Tool: request.Tool, Input: request.Input, + ExpiresAt: expiresAt, + }) + if err != nil { + return api.PermissionDecision{}, err + } + if err := e.markWaiting(ctx); err != nil { + return api.PermissionDecision{}, err + } + if err := e.emitApproval(ctx, pending.ID, request); err != nil { + return api.PermissionDecision{}, err + } + decision, err := e.waitForApproval(ctx, pending.ID) + restoreErr := e.markRunning(ctx) + return decision, errors.Join(err, restoreErr) +} + +func (e *databaseExecution) emitApproval(ctx context.Context, approvalID uuid.UUID, request api.PermissionRequest) error { + event := api.Event{ + Kind: api.EventPermission, Tool: request.Tool, + ToolCallID: request.ToolUseID, ApprovalID: approvalID.String(), Input: request.Input, + } + select { + case e.events <- event: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (e *databaseExecution) waitForApproval( + ctx context.Context, + requestID uuid.UUID, +) (api.PermissionDecision, error) { + ticker := time.NewTicker(approvalPollInterval) + defer ticker.Stop() + for { + request, err := e.db.GetTurnRequest(ctx, requestID) + if err != nil { + return api.PermissionDecision{}, err + } + switch request.State { + case database.TurnRequestStateApproved: + decision := api.PermissionDecision{Allow: true} + if updated, ok := request.Response["updatedInput"].(map[string]any); ok { + decision.UpdatedInput = updated + } + return decision, nil + case database.TurnRequestStateDenied: + message := request.Reason + if message == "" { + message = "tool call denied" + } + return api.PermissionDecision{Message: message}, nil + case database.TurnRequestStateExpired, database.TurnRequestStateCancelled: + return api.PermissionDecision{}, fmt.Errorf("tool approval %s", request.State) + } + if request.ExpiresAt != nil && !time.Now().Before(*request.ExpiresAt) { + if err := e.db.ExpireToolApprovalRequest(ctx, request.ID, database.TurnRequestStateExpired, "approval timed out"); err != nil { + return api.PermissionDecision{}, err + } + continue + } + if err := e.db.ValidateCallerToolCredential(ctx, *request.CredentialID); err != nil { + _ = e.db.ExpireToolApprovalRequest(ctx, request.ID, database.TurnRequestStateCancelled, err.Error()) + return api.PermissionDecision{}, err + } + select { + case <-ctx.Done(): + _ = e.db.ExpireToolApprovalRequest(context.Background(), request.ID, database.TurnRequestStateCancelled, ctx.Err().Error()) + return api.PermissionDecision{}, ctx.Err() + case <-ticker.C: + } + } +} + +func (e *databaseExecution) Observe(ctx context.Context, event api.Event) (api.Event, error) { + if event.SessionID != "" { + if err := e.bindProviderSession(ctx, event.SessionID); err != nil { + return event, err + } + } + switch event.Kind { + case api.EventToolUse: + e.rememberProviderToolUse(event) + return event, nil + case api.EventPermission: + if event.ApprovalID != "" { + return event, nil + } + approval, err := e.createProviderApproval(ctx, event) + if err != nil { + return event, err + } + event.ApprovalID = approval.ID.String() + return event, nil + case api.EventResult: + if event.ToolApproval != nil { + return event, e.suspend(ctx, *event.ToolApproval, event) + } + return event, e.finish(ctx, true, "", event) + case api.EventError: + return event, e.finish(ctx, false, event.Error, event) + default: + return event, nil + } +} + +func (e *databaseExecution) suspend(ctx context.Context, state api.ToolApprovalState, event api.Event) error { + if state.ProviderCheckpoint == nil { + return fmt.Errorf("provider tool approval ended without a private checkpoint") + } + for _, pending := range state.Pending() { + e.mu.Lock() + _, ok := e.approvalIDs[pending.ToolCallID] + e.mu.Unlock() + if !ok { + return fmt.Errorf("provider tool approval %q has no durable turn request", pending.ToolCallID) + } + } + if err := e.finishModelCall(ctx, database.ModelCallStatusSucceeded, "tool_approval", event); err != nil { + return err + } + checkpoint := database.PromptRunCheckpoint{ + Codec: state.ProviderCheckpoint.Codec, Version: state.ProviderCheckpoint.Version, + Payload: state.ProviderCheckpoint.Payload, + } + waiting := database.PromptRunStateWaiting + state.ProviderCheckpoint = nil + if err := e.updateRun(ctx, runUpdate{ + State: &waiting, ApprovalState: &state, ProviderCheckpoint: &checkpoint, + }); err != nil { + return err + } + if err := e.updateSessionActivity(ctx, database.SessionActivityApproval); err != nil { + return err + } + e.mu.Lock() + e.suspended = true + e.mu.Unlock() + return nil +} + +func (e *databaseExecution) Close(ctx context.Context) error { + e.mu.Lock() + if e.closed { + e.mu.Unlock() + return nil + } + e.closed = true + terminal := e.terminal + suspended := e.suspended + runtime := e.runtime + credential := e.credential + e.mu.Unlock() + + var errs []error + if !terminal && !suspended { + errs = append(errs, e.finish(ctx, false, "provider stream ended without a terminal event", api.Event{Kind: api.EventError, Error: "provider stream ended without a terminal event"})) + } + if credential != nil { + errs = append(errs, e.db.RevokeCallerToolCredential(ctx, credential.ID, "execution closed")) + } + if runtime != nil { + errs = append(errs, runtime.Close()) + } + return errors.Join(errs...) +} + +func (e *databaseExecution) Interrupt(ctx context.Context, reason string) error { + e.finishMu.Lock() + defer e.finishMu.Unlock() + e.mu.Lock() + if e.terminal { + e.mu.Unlock() + return nil + } + runtime := e.runtime + credential := e.credential + e.mu.Unlock() + + if runtime != nil { + runtime.Revoke() + } + var errs []error + errs = append(errs, e.finishModelCall(ctx, database.ModelCallStatusCancelled, "interrupt", + api.Event{Kind: api.EventInterrupted, Reason: reason})) + errs = append(errs, e.db.CancelPendingTurnRequests(ctx, e.session.ID, e.run.ID, "execution interrupted")) + if credential != nil { + errs = append(errs, e.db.RevokeCallerToolCredential(ctx, credential.ID, "execution interrupted")) + } + phase := database.PromptRunPhaseFinished + state := database.PromptRunStateCancelled + errs = append(errs, e.updateRun(ctx, runUpdate{ + Phase: &phase, State: &state, ClearApprovalState: true, ClearProviderCheckpoint: true, + })) + errs = append(errs, e.db.FinishChatTurn(ctx, e.turn.ID, database.TurnStatusInterrupted, "interrupt")) + errs = append(errs, e.updateSessionState( + ctx, database.SessionLifecycleInterrupted, database.SessionActivityIdle, reason, + )) + if err := errors.Join(errs...); err != nil { + return err + } + e.mu.Lock() + e.terminal = true + e.mu.Unlock() + return nil +} + +func (e *databaseExecution) bindProviderSession(ctx context.Context, providerID string) error { + e.mu.Lock() + defer e.mu.Unlock() + providerID = strings.TrimSpace(providerID) + if providerID == "" || providerID == e.providerID { + return nil + } + if e.providerID != "" { + return fmt.Errorf("provider session is already bound to %q", e.providerID) + } + session, err := e.db.GetSession(ctx, e.session.ID) + if err != nil { + return err + } + updated, err := e.db.UpdateSessionState(ctx, database.UpdateSessionStateInput{ + ID: session.ID, ExpectedVersion: session.StateVersion, ProviderSessionID: &providerID, + }) + if err != nil { + return err + } + e.session = updated + e.providerID = providerID + return nil +} + +func (e *databaseExecution) markRunning(ctx context.Context) error { + phase := database.PromptRunPhaseGenerate + state := database.PromptRunStateRunning + activity := database.SessionActivityWorking + if err := e.updateRun(ctx, runUpdate{Phase: &phase, State: &state}); err != nil { + return err + } + return e.updateSessionActivity(ctx, activity) +} + +func (e *databaseExecution) markWaiting(ctx context.Context) error { + state := database.PromptRunStateWaiting + activity := database.SessionActivityApproval + if err := e.updateRun(ctx, runUpdate{State: &state}); err != nil { + return err + } + return e.updateSessionActivity(ctx, activity) +} + +func (e *databaseExecution) finish(ctx context.Context, success bool, message string, event api.Event) error { + e.finishMu.Lock() + defer e.finishMu.Unlock() + e.mu.Lock() + if e.terminal { + e.mu.Unlock() + return nil + } + runtime := e.runtime + credential := e.credential + e.mu.Unlock() + + if runtime != nil { + runtime.Revoke() + } + var errs []error + callStatus := database.ModelCallStatusFailed + stopReason := "error" + if success { + callStatus = database.ModelCallStatusSucceeded + stopReason = "stop" + } + errs = append(errs, e.finishModelCall(ctx, callStatus, stopReason, event)) + if credential != nil { + errs = append(errs, e.db.RevokeCallerToolCredential(ctx, credential.ID, "prompt run terminal")) + } + phase := database.PromptRunPhaseFinished + state := database.PromptRunStateFailed + if success { + state = database.PromptRunStateSucceeded + } + errs = append(errs, e.updateRun(ctx, runUpdate{ + Phase: &phase, State: &state, Message: &message, + ClearApprovalState: true, ClearProviderCheckpoint: true, + })) + turnState := database.TurnStatusError + turnStopReason := message + if success { + turnState = database.TurnStatusEnded + turnStopReason = "stop" + } + errs = append(errs, e.db.FinishChatTurn(ctx, e.turn.ID, turnState, turnStopReason)) + lifecycle := database.SessionLifecycleFailed + if success { + lifecycle = database.SessionLifecycleSucceeded + } + errs = append(errs, e.updateSessionState(ctx, lifecycle, database.SessionActivityIdle, message)) + if err := errors.Join(errs...); err != nil { + return err + } + e.mu.Lock() + e.terminal = true + e.mu.Unlock() + return nil +} + +type runUpdate struct { + Phase *database.PromptRunPhase + State *database.PromptRunState + Message *string + ApprovalState *api.ToolApprovalState + ProviderCheckpoint *database.PromptRunCheckpoint + ClearApprovalState bool + ClearProviderCheckpoint bool +} + +func (e *databaseExecution) updateRun(ctx context.Context, update runUpdate) error { + e.mu.Lock() + defer e.mu.Unlock() + input := database.UpdatePromptRunInput{ + ID: e.run.ID, ExpectedVersion: e.run.Version, Phase: update.Phase, State: update.State, + } + if update.Message != nil && *update.Message != "" { + input.Error = update.Message + } + input.ApprovalState = update.ApprovalState + input.ProviderCheckpoint = update.ProviderCheckpoint + input.ClearApprovalState = update.ClearApprovalState + input.ClearProviderCheckpoint = update.ClearProviderCheckpoint + run, err := e.db.UpdatePromptRun(ctx, input) + if err != nil { + return err + } + e.run = run + return nil +} + +func (e *databaseExecution) updateSessionActivity( + ctx context.Context, + activity database.SessionActivityState, +) error { + return e.updateSessionState(ctx, database.SessionLifecycleRunning, activity, "") +} + +func (e *databaseExecution) updateSessionState( + ctx context.Context, + lifecycle database.SessionLifecycleStatus, + activity database.SessionActivityState, + reason string, +) error { + e.mu.Lock() + defer e.mu.Unlock() + session, err := e.db.GetSession(ctx, e.session.ID) + if err != nil { + return err + } + updated, err := e.db.UpdateSessionState(ctx, database.UpdateSessionStateInput{ + ID: session.ID, ExpectedVersion: session.StateVersion, + LifecycleStatus: &lifecycle, ActivityState: &activity, StateReason: &reason, + }) + if err != nil { + return err + } + e.session = updated + return nil +} diff --git a/pkg/aichat/execution_database_authority.go b/pkg/aichat/execution_database_authority.go new file mode 100644 index 00000000..1a38f9e7 --- /dev/null +++ b/pkg/aichat/execution_database_authority.go @@ -0,0 +1,323 @@ +package aichat + +import ( + "context" + "encoding/json" + "errors" + "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 DatabaseExecutionAuthority struct { + db *database.DB +} + +func NewDatabaseExecutionAuthority(db *database.DB) (*DatabaseExecutionAuthority, error) { + if db == nil || db.Gorm() == nil { + return nil, fmt.Errorf("captain execution authority requires a database") + } + return &DatabaseExecutionAuthority{db: db}, nil +} + +func (a *DatabaseExecutionAuthority) Begin( + ctx context.Context, + request ExecutionRequest, +) (Execution, error) { + sessionID, err := uuid.Parse(request.ThreadID) + if err != nil { + return nil, fmt.Errorf("chat thread ID %q is not a UUID: %w", request.ThreadID, err) + } + if request.Spec.Backend == "" { + return nil, fmt.Errorf("authoritative chat execution requires a resolved backend") + } + session, err := a.db.CreateOrGetSession(ctx, database.CreateSessionInput{ + ID: sessionID, Source: "aichat", Provider: ai.BackendToProvider(request.Spec.Backend), + HostID: "local", Title: request.Title, InitialPrompt: initialUserPrompt(request.Spec), + Metadata: map[string]any{"aichat": true}, + }) + if err != nil { + return nil, err + } + if session.Source != "aichat" { + return nil, fmt.Errorf("chat thread %s has incompatible source %q", request.ThreadID, session.Source) + } + renderedSpec, err := renderedSpecMap(request.Spec) + if err != nil { + return nil, err + } + var execution *databaseExecution + var recovered *database.ChatTurn + resumed := false + err = a.db.Transaction(ctx, func(tx *database.DB) error { + var createErr error + recovered, createErr = tx.RecoverIncompleteChatAdmission(ctx, database.RecoverIncompleteChatAdmissionInput{ + SessionID: session.ID, ProviderTurnID: request.RequestID, + }) + if createErr != nil { + return createErr + } + turn, created, createErr := tx.CreateChatTurn(ctx, database.CreateChatTurnInput{ + SessionID: session.ID, ProviderTurnID: request.RequestID, + }) + if createErr != nil { + return createErr + } + if !created && turn.Status != database.TurnStatusOpen { + return fmt.Errorf("chat turn %q already exists in state %s", request.RequestID, turn.Status) + } + resumed = !created + run, createErr := tx.CreatePromptRun(ctx, database.CreatePromptRunInput{ + SessionID: session.ID, TurnID: &turn.ID, AdmissionKey: executionAdmissionKey(request), + Origin: "aichat", RenderedSpec: renderedSpec, + Runtime: database.PromptRunRuntime{ + Mode: string(request.Spec.Mode), Driver: string(request.Spec.Backend), + Requested: runtimeSelection(request.Spec.Model), + Resolved: runtimeSelection(request.Spec.Model), + }, + PromptMarkdown: initialUserPrompt(request.Spec), + }) + if createErr != nil { + return createErr + } + if run.State != database.PromptRunStatePending { + return fmt.Errorf("chat request %q already has prompt run %s in state %s", request.RequestID, run.ID, run.State) + } + modelCallID, createErr := tx.CreateChatModelCall(ctx, database.CreateChatModelCallInput{ + TurnID: turn.ID, PromptRunID: run.ID, Model: request.Spec.Name, + Backend: string(request.Spec.Backend), Effort: string(request.Spec.Effort), + }) + if createErr != nil { + return createErr + } + execution = &databaseExecution{ + db: tx, ctx: ctx, session: session, turn: turn, run: run, modelCallID: modelCallID, + model: request.Spec.Name, backend: request.Spec.Backend, + events: make(chan api.Event, 16), definitions: append([]api.ToolDefinition(nil), request.Definitions...), + approvalIDs: map[string]uuid.UUID{}, providerToolUseReady: make(chan struct{}, 1), + } + return execution.markRunning(ctx) + }) + if err != nil { + return nil, err + } + execution.db = a.db + if recovered != nil { + serviceLog.Warnf("recovered incomplete chat admission turn %s for session %s", recovered.ID, session.ID) + } + if resumed { + serviceLog.Warnf("resumed incomplete chat admission turn %s for session %s", execution.turn.ID, session.ID) + } + if len(request.Definitions) > 0 && isAgentBackend(request.Spec.Backend) { + if err := execution.startCallerTools(ctx, request.Spec.Backend); err != nil { + _ = execution.Close(context.Background()) + return nil, err + } + } + return execution, nil +} + +func (a *DatabaseExecutionAuthority) ResolveToolApproval( + ctx context.Context, + resolution ToolApprovalResolution, +) (*ApprovalContinuation, error) { + sessionID, err := uuid.Parse(resolution.ThreadID) + if err != nil { + return nil, fmt.Errorf("chat thread ID %q is not a UUID: %w", resolution.ThreadID, err) + } + approvalID, err := uuid.Parse(resolution.ApprovalID) + if err != nil { + return nil, fmt.Errorf("tool approval ID %q is not a UUID: %w", resolution.ApprovalID, err) + } + request, err := a.db.ResolveToolApprovalRequest(ctx, database.ResolveToolApprovalRequestInput{ + SessionID: sessionID, RequestID: approvalID, + Approved: resolution.Approved, UpdatedInput: resolution.UpdatedInput, + ResolvedBy: "chat", Reason: resolution.Reason, + }) + if err != nil { + return nil, err + } + if request.CredentialID != nil { + return nil, nil + } + if request.PromptRunID == nil || request.TurnID == nil { + return nil, fmt.Errorf("provider approval %s has no prompt run or turn", request.ID) + } + requests, err := a.db.ListTurnRequests(ctx, database.TurnRequestFilter{ + SessionID: sessionID, PromptRunID: request.PromptRunID, + }) + if err != nil { + return nil, err + } + for _, item := range requests { + if item.State == database.TurnRequestStatePending { + return nil, nil + } + } + run, err := a.db.GetPromptRun(ctx, *request.PromptRunID) + if err != nil { + return nil, err + } + if run.State != database.PromptRunStateWaiting { + return nil, nil + } + if run.ApprovalState == nil || run.ProviderCheckpoint == nil { + return nil, fmt.Errorf("waiting prompt run %s has no durable approval state and provider checkpoint", run.ID) + } + state := *run.ApprovalState + state.ProviderCheckpoint = &api.ProviderCheckpoint{ + Codec: run.ProviderCheckpoint.Codec, Version: run.ProviderCheckpoint.Version, + Payload: append([]byte(nil), run.ProviderCheckpoint.Payload...), + } + decisions, err := approvalDecisions(state, requests) + if err != nil { + return nil, err + } + rendered, err := json.Marshal(run.RenderedSpec) + if err != nil { + return nil, fmt.Errorf("encode prompt run %s rendered spec: %w", run.ID, err) + } + var spec api.Spec + if err := json.Unmarshal(rendered, &spec); err != nil { + return nil, fmt.Errorf("decode prompt run %s rendered spec: %w", run.ID, err) + } + spec.Messages = nil + spec.Prompt.User = "" + spec.Prompt.System = "" + spec.Prompt.AppendSystem = "" + spec.Prompt.Attachments = nil + spec.ToolApproval = &api.ToolApprovalResume{State: state, Decisions: decisions} + turn, err := a.db.GetChatTurn(ctx, *request.TurnID) + if err != nil { + return nil, err + } + running := database.PromptRunStateRunning + phase := database.PromptRunPhaseGenerate + var resumed *database.PromptRun + var modelCallID uuid.UUID + err = a.db.Transaction(ctx, func(tx *database.DB) error { + var updateErr error + resumed, updateErr = tx.UpdatePromptRun(ctx, database.UpdatePromptRunInput{ + ID: run.ID, ExpectedVersion: run.Version, State: &running, Phase: &phase, + ClearApprovalState: true, ClearProviderCheckpoint: true, + }) + if updateErr != nil { + return updateErr + } + modelCallID, updateErr = tx.CreateChatModelCall(ctx, database.CreateChatModelCallInput{ + TurnID: turn.ID, PromptRunID: run.ID, Model: spec.Name, + Backend: string(spec.Backend), Effort: string(spec.Effort), + }) + return updateErr + }) + if err != nil { + if errors.Is(err, database.ErrPromptRunConflict) { + return nil, nil + } + return nil, err + } + sessionRecord, err := a.db.GetSession(ctx, sessionID) + if err != nil { + return nil, err + } + execution := &databaseExecution{ + db: a.db, ctx: ctx, session: sessionRecord, turn: turn, run: resumed, modelCallID: modelCallID, + model: spec.Name, backend: spec.Backend, + events: make(chan api.Event, 16), approvalIDs: map[string]uuid.UUID{}, + providerToolUseReady: make(chan struct{}, 1), + } + if err := execution.updateSessionActivity(ctx, database.SessionActivityWorking); err != nil { + return nil, err + } + return &ApprovalContinuation{Execution: execution, Spec: spec}, nil +} + +func approvalDecisions(state api.ToolApprovalState, requests []database.TurnRequest) ([]api.ToolApprovalDecision, error) { + byCall := make(map[string]database.TurnRequest, len(requests)) + for _, request := range requests { + byCall[request.ToolCallID] = request + } + decisions := make([]api.ToolApprovalDecision, 0, len(state.Pending())) + for _, pending := range state.Pending() { + request, ok := byCall[pending.ToolCallID] + if !ok { + return nil, fmt.Errorf("approval state tool call %q has no durable turn request", pending.ToolCallID) + } + decision := api.ToolApprovalDecision{ + ApprovalID: request.ID.String(), ToolCallID: pending.ToolCallID, Tool: pending.Tool, + } + switch request.State { + case database.TurnRequestStateApproved: + decision.Action = api.ToolApprovalApprove + if updated := request.Response["updatedInput"]; updated != nil { + encoded, err := json.Marshal(updated) + if err != nil { + return nil, fmt.Errorf("encode approval %s updated input: %w", request.ID, err) + } + decision.Input = encoded + } + case database.TurnRequestStateDenied: + decision.Action = api.ToolApprovalDeny + decision.Message = request.Reason + default: + return nil, fmt.Errorf("approval %s is in non-resumable state %s", request.ID, request.State) + } + decisions = append(decisions, decision) + } + return decisions, nil +} + +func runtimeSelection(model api.Model) database.PromptRunRuntimeSelection { + return database.PromptRunRuntimeSelection{ + Provider: ai.BackendToProvider(model.Backend), Backend: string(model.Backend), + Model: model.Name, Effort: string(model.Effort), + } +} + +func renderedSpecMap(spec api.Spec) (map[string]any, error) { + raw, err := json.Marshal(spec) + if err != nil { + return nil, fmt.Errorf("encode authoritative chat spec: %w", err) + } + var rendered map[string]any + if err := json.Unmarshal(raw, &rendered); err != nil { + return nil, fmt.Errorf("decode authoritative chat spec: %w", err) + } + return rendered, nil +} + +func executionAdmissionKey(request ExecutionRequest) string { + if strings.TrimSpace(request.RequestID) == "" { + return "" + } + return "aichat:" + request.ThreadID + ":" + request.RequestID +} + +func initialUserPrompt(spec api.Spec) string { + if spec.Prompt.User != "" { + return spec.Prompt.User + } + for i := len(spec.Messages) - 1; i >= 0; i-- { + if spec.Messages[i].Role != api.RoleUser { + continue + } + for _, part := range spec.Messages[i].Parts { + if part.Type == api.PartText && strings.TrimSpace(part.Text) != "" { + return part.Text + } + } + } + return "" +} + +func cloneStringValues(values map[string]string) map[string]string { + cloned := make(map[string]string, len(values)) + for key, value := range values { + cloned[key] = value + } + return cloned +} diff --git a/pkg/aichat/execution_database_correlation.go b/pkg/aichat/execution_database_correlation.go new file mode 100644 index 00000000..814384f8 --- /dev/null +++ b/pkg/aichat/execution_database_correlation.go @@ -0,0 +1,65 @@ +package aichat + +import ( + "context" + "fmt" + "reflect" + "time" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" +) + +const providerToolCorrelationTTL = 5 * time.Second + +func (e *databaseExecution) rememberProviderToolUse(event api.Event) { + e.mu.Lock() + e.providerToolUses = append(e.providerToolUses, event) + e.mu.Unlock() + select { + case e.providerToolUseReady <- struct{}{}: + default: + } +} + +func (e *databaseExecution) claimProviderToolUse(ctx context.Context, request api.PermissionRequest) (string, error) { + timer := time.NewTimer(providerToolCorrelationTTL) + defer timer.Stop() + for { + e.mu.Lock() + for i, event := range e.providerToolUses { + if event.Tool != request.Tool || !reflect.DeepEqual(event.Input, request.Input) { + continue + } + e.providerToolUses = append(e.providerToolUses[:i], e.providerToolUses[i+1:]...) + e.mu.Unlock() + return event.ToolCallID, nil + } + e.mu.Unlock() + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-timer.C: + return "", fmt.Errorf("caller tool %q did not match a provider tool use within %s", request.Tool, providerToolCorrelationTTL) + case <-e.providerToolUseReady: + } + } +} + +func (e *databaseExecution) createProviderApproval(ctx context.Context, event api.Event) (*database.TurnRequest, error) { + if event.ToolCallID == "" || event.Tool == "" { + return nil, fmt.Errorf("provider approval requires a tool call ID and tool name") + } + request, err := e.db.CreateToolApprovalRequest(ctx, database.CreateToolApprovalRequestInput{ + SessionID: e.session.ID, TurnID: e.turn.ID, PromptRunID: e.run.ID, ModelCallID: e.modelCallID, + ToolCallID: event.ToolCallID, Tool: event.Tool, Input: event.Input, + RequestedBy: "provider", ExpiresAt: time.Now().Add(providerApprovalTimeout), + }) + if err != nil { + return nil, err + } + e.mu.Lock() + e.approvalIDs[event.ToolCallID] = request.ID + e.mu.Unlock() + return request, nil +} diff --git a/pkg/aichat/execution_database_integration_test.go b/pkg/aichat/execution_database_integration_test.go new file mode 100644 index 00000000..17bbc377 --- /dev/null +++ b/pkg/aichat/execution_database_integration_test.go @@ -0,0 +1,411 @@ +package aichat_test + +import ( + "context" + "errors" + "sync/atomic" + "time" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/commons-db/dbtest" + "github.com/google/uuid" + mcpclient "github.com/mark3labs/mcp-go/client" + "github.com/mark3labs/mcp-go/client/transport" + "github.com/mark3labs/mcp-go/mcp" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Database execution authority", func() { + It("blocks an ask tool on its durable approval and revokes the credential at completion", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_execution"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + var calls atomic.Int32 + threadID := uuid.NewString() + execution, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "request-account-1", Title: "Accounts", + Spec: api.Spec{Model: api.Model{ + Name: "sonnet", Backend: api.BackendClaudeAgent, + }.Capabilities()}, + Definitions: []api.ToolDefinition{{ + Name: "account_edit", DefaultPermission: api.ToolModeAsk, + Handler: func(_ context.Context, input map[string]any) (any, error) { + calls.Add(1) + return input, nil + }, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(execution.Close) + + client := executionMCPClient(ctx, *execution.CallerTools()) + DeferCleanup(client.Close) + type callOutcome struct { + result *mcp.CallToolResult + err error + } + outcomes := make(chan callOutcome, 1) + _, err = execution.Observe(ctx, api.Event{ + Kind: api.EventToolUse, Tool: "account_edit", ToolCallID: "call-account-1", + Input: map[string]any{"name": "Draft"}, + }) + Expect(err).NotTo(HaveOccurred()) + go func() { + request := mcp.CallToolRequest{} + request.Params.Name = "account_edit" + request.Params.Arguments = map[string]any{"name": "Draft"} + result, callErr := client.CallTool(ctx, request) + outcomes <- callOutcome{result: result, err: callErr} + }() + + var approval api.Event + Eventually(execution.Events()).Should(Receive(&approval)) + Expect(approval.Kind).To(Equal(api.EventPermission)) + Expect(approval.ToolCallID).To(Equal("call-account-1")) + Expect(approval.ApprovalID).To(MatchRegexp(`^[0-9a-f-]{36}$`)) + Expect(calls.Load()).To(BeZero()) + + continuation, err := authority.ResolveToolApproval(ctx, aichat.ToolApprovalResolution{ + ThreadID: threadID, ApprovalID: approval.ApprovalID, Approved: true, + UpdatedInput: map[string]any{"name": "Approved"}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(continuation).To(BeNil()) + var outcome callOutcome + Eventually(outcomes).Should(Receive(&outcome)) + Expect(outcome.err).NotTo(HaveOccurred()) + Expect(outcome.result.IsError).To(BeFalse()) + Expect(outcome.result.StructuredContent).To(Equal(map[string]any{"name": "Approved"})) + Expect(calls.Load()).To(Equal(int32(1))) + + _, err = execution.Observe(ctx, api.Event{ + Kind: api.EventResult, Success: true, SessionID: "provider-session-1", + }) + Expect(err).NotTo(HaveOccurred()) + Expect(execution.Close(ctx)).To(Succeed()) + + runID := uuid.MustParse(execution.PromptRunID()) + run, err := db.GetPromptRun(ctx, runID) + Expect(err).NotTo(HaveOccurred()) + Expect(run.State).To(Equal(database.PromptRunStateSucceeded)) + var credential struct { + RevokedAt *time.Time + } + Expect(db.Gorm().WithContext(ctx). + Table("captain_session_mcp_credentials"). + Select("revoked_at"). + Where("prompt_run_id = ?", runID). + Scan(&credential).Error).To(Succeed()) + Expect(credential.RevokedAt).NotTo(BeNil()) + }) + + It("creates distinct prompt runs for sequential turn identities and rejects a replay", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_sequential_turns"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + threadID := uuid.NewString() + spec := api.Spec{Model: api.Model{Name: "gemini", Backend: api.BackendGemini}.Capabilities()} + for _, turnID := range []string{"user-message-1", "user-message-2"} { + execution, beginErr := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: turnID, Title: "Accounts", Spec: spec, + }) + Expect(beginErr).NotTo(HaveOccurred()) + _, err = execution.Observe(ctx, api.Event{Kind: api.EventResult, Success: true}) + Expect(err).NotTo(HaveOccurred()) + Expect(execution.Close(ctx)).To(Succeed()) + } + + sessionID := uuid.MustParse(threadID) + runs, err := db.ListPromptRuns(ctx, database.PromptRunFilter{SessionID: &sessionID}) + Expect(err).NotTo(HaveOccurred()) + Expect(runs).To(HaveLen(2)) + Expect(runs).To(ConsistOf( + HaveField("AdmissionKey", "aichat:"+threadID+":user-message-1"), + HaveField("AdmissionKey", "aichat:"+threadID+":user-message-2"), + )) + + _, err = authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "user-message-1", Title: "Accounts", Spec: spec, + }) + Expect(err).To(MatchError(ContainSubstring("already exists in state ended"))) + }) + + It("rolls back an admission when its model call cannot be created", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_atomic_admission"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + threadID := uuid.NewString() + _, err = authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "invalid-model-call", Title: "Atomic admission", + Spec: api.Spec{Model: api.Model{Backend: api.BackendOpenAI}.Capabilities()}, + }) + Expect(err).To(MatchError(ContainSubstring("model"))) + + sessionID := uuid.MustParse(threadID) + turns, err := db.ListThreadTurns(ctx, sessionID) + Expect(err).NotTo(HaveOccurred()) + Expect(turns).To(BeEmpty()) + runs, err := db.ListPromptRuns(ctx, database.PromptRunFilter{SessionID: &sessionID}) + Expect(err).NotTo(HaveOccurred()) + Expect(runs).To(BeEmpty()) + + execution, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "valid-model-call", Title: "Atomic admission", + Spec: api.Spec{Model: api.Model{Name: "gpt", Backend: api.BackendOpenAI}.Capabilities()}, + }) + Expect(err).NotTo(HaveOccurred()) + _, err = execution.Observe(ctx, api.Event{Kind: api.EventResult, Success: true}) + Expect(err).NotTo(HaveOccurred()) + Expect(execution.Close(ctx)).To(Succeed()) + }) + + It("terminalizes an incomplete admission before opening the next turn", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_incomplete_admission"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + session, err := db.CreateOrGetSession(ctx, database.CreateSessionInput{ + ID: uuid.New(), Source: "aichat", Provider: "openai", HostID: "local", + }) + Expect(err).NotTo(HaveOccurred()) + incompleteTurn, created, err := db.CreateChatTurn(ctx, database.CreateChatTurnInput{ + SessionID: session.ID, ProviderTurnID: "incomplete-model-call", + }) + Expect(err).NotTo(HaveOccurred()) + Expect(created).To(BeTrue()) + incompleteRun, err := db.CreatePromptRun(ctx, database.CreatePromptRunInput{ + SessionID: session.ID, TurnID: &incompleteTurn.ID, + AdmissionKey: "aichat:" + session.ID.String() + ":incomplete-model-call", + }) + Expect(err).NotTo(HaveOccurred()) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + execution, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: session.ID.String(), RequestID: "valid-model-call", Title: "Recovered admission", + Spec: api.Spec{Model: api.Model{Name: "gpt", Backend: api.BackendOpenAI}.Capabilities()}, + }) + Expect(err).NotTo(HaveOccurred()) + + incompleteTurn, err = db.GetChatTurn(ctx, incompleteTurn.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(incompleteTurn.Status).To(Equal(database.TurnStatusError)) + incompleteRun, err = db.GetPromptRun(ctx, incompleteRun.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(incompleteRun.State).To(Equal(database.PromptRunStateFailed)) + Expect(incompleteRun.Error).To(Equal("chat execution admission did not complete")) + + _, err = execution.Observe(ctx, api.Event{Kind: api.EventResult, Success: true}) + Expect(err).NotTo(HaveOccurred()) + Expect(execution.Close(ctx)).To(Succeed()) + }) + + It("resumes an incomplete admission when the same request is retried", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_incomplete_retry"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + requestID := "retried-model-call" + session, err := db.CreateOrGetSession(ctx, database.CreateSessionInput{ + ID: uuid.New(), Source: "aichat", Provider: "openai", HostID: "local", + }) + Expect(err).NotTo(HaveOccurred()) + incompleteTurn, created, err := db.CreateChatTurn(ctx, database.CreateChatTurnInput{ + SessionID: session.ID, ProviderTurnID: requestID, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(created).To(BeTrue()) + incompleteRun, err := db.CreatePromptRun(ctx, database.CreatePromptRunInput{ + SessionID: session.ID, TurnID: &incompleteTurn.ID, + AdmissionKey: "aichat:" + session.ID.String() + ":" + requestID, + }) + Expect(err).NotTo(HaveOccurred()) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + execution, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: session.ID.String(), RequestID: requestID, Title: "Retried admission", + Spec: api.Spec{Model: api.Model{Name: "gpt", Backend: api.BackendOpenAI}.Capabilities()}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(execution.TurnID()).To(Equal(incompleteTurn.ID.String())) + Expect(execution.PromptRunID()).To(Equal(incompleteRun.ID.String())) + + _, err = execution.Observe(ctx, api.Event{Kind: api.EventResult, Success: true}) + Expect(err).NotTo(HaveOccurred()) + Expect(execution.Close(ctx)).To(Succeed()) + turns, err := db.ListThreadTurns(ctx, session.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(turns).To(HaveLen(1)) + Expect(turns[0].Status).To(Equal(string(database.TurnStatusEnded))) + }) + + It("rejects a second admission while the first turn is running", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_active_admission"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + threadID := uuid.NewString() + spec := api.Spec{Model: api.Model{Name: "gpt", Backend: api.BackendOpenAI}.Capabilities()} + execution, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "running-model-call", Title: "Active admission", Spec: spec, + }) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(execution.Close) + + _, err = authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "second-model-call", Title: "Active admission", Spec: spec, + }) + Expect(errors.Is(err, database.ErrOpenChatTurn)).To(BeTrue()) + Expect(err.Error()).NotTo(ContainSubstring("duplicate key")) + }) + + It("records an interruption and admits a later turn on the same Captain session", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_interrupt_resume"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + threadID := uuid.NewString() + spec := api.Spec{Model: api.Model{Name: "gpt", Backend: api.BackendOpenAI}.Capabilities()} + interrupted, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "user-message-interrupted", Title: "Interrupt", Spec: spec, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(interrupted.Interrupt(ctx, "user")).To(Succeed()) + Expect(interrupted.Close(ctx)).To(Succeed()) + + run, err := db.GetPromptRun(ctx, uuid.MustParse(interrupted.PromptRunID())) + Expect(err).NotTo(HaveOccurred()) + Expect(run.State).To(Equal(database.PromptRunStateCancelled)) + sessionRecord, err := db.GetSession(ctx, uuid.MustParse(threadID)) + Expect(err).NotTo(HaveOccurred()) + Expect(sessionRecord.LifecycleStatus).To(Equal(database.SessionLifecycleInterrupted)) + Expect(sessionRecord.ActivityState).To(Equal(database.SessionActivityIdle)) + turns, err := db.ListThreadTurns(ctx, uuid.MustParse(threadID)) + Expect(err).NotTo(HaveOccurred()) + Expect(turns).To(HaveLen(1)) + Expect(turns[0].Status).To(Equal(string(database.TurnStatusInterrupted))) + Expect(turns[0].StopReason).NotTo(BeNil()) + Expect(*turns[0].StopReason).To(Equal("interrupt")) + var modelCall struct{ Status, StopReason string } + Expect(db.Gorm().WithContext(ctx).Table("captain_model_calls"). + Select("status, stop_reason").Where("prompt_run_id = ?", run.ID). + Scan(&modelCall).Error).To(Succeed()) + Expect(modelCall.Status).To(Equal(string(database.ModelCallStatusCancelled))) + Expect(modelCall.StopReason).To(Equal("interrupt")) + + resumed, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: threadID, RequestID: "user-message-resumed", Title: "Interrupt", Spec: spec, + }) + Expect(err).NotTo(HaveOccurred()) + _, err = resumed.Observe(ctx, api.Event{Kind: api.EventResult, Success: true}) + Expect(err).NotTo(HaveOccurred()) + Expect(resumed.Close(ctx)).To(Succeed()) + sessionRecord, err = db.GetSession(ctx, uuid.MustParse(threadID)) + Expect(err).NotTo(HaveOccurred()) + Expect(sessionRecord.LifecycleStatus).To(Equal(database.SessionLifecycleSucceeded)) + Expect(sessionRecord.ActivityState).To(Equal(database.SessionActivityIdle)) + }) + + It("keeps an interrupted API run waiting for its durable approvals", func(ctx SpecContext) { + testDB := dbtest.ForGinkgo(dbtest.Options{Name: "captain_aichat_waiting_approval"}) + db, err := database.Open(ctx, database.WithDSN(testDB.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(db.Close) + + authority, err := aichat.NewDatabaseExecutionAuthority(db) + Expect(err).NotTo(HaveOccurred()) + execution, err := authority.Begin(ctx, aichat.ExecutionRequest{ + ThreadID: uuid.NewString(), RequestID: "user-message-approval", Title: "Accounts", + Spec: api.Spec{ + Model: api.Model{Name: "gemini", Backend: api.BackendGemini}.Capabilities(), + Messages: []api.Message{{Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Edit the account"}}}}, + }, + }) + Expect(err).NotTo(HaveOccurred()) + + permission, err := execution.Observe(ctx, api.Event{ + Kind: api.EventPermission, ToolCallID: "call-account-approval", Tool: "account_edit", + Input: map[string]any{"id": "acc-1"}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(permission.ApprovalID).To(MatchRegexp(`^[0-9a-f-]{36}$`)) + + _, err = execution.Observe(ctx, api.Event{ + Kind: api.EventResult, Success: true, + ToolApproval: &api.ToolApprovalState{ + ProviderCheckpoint: &api.ProviderCheckpoint{ + Codec: "test-checkpoint", Version: 1, Payload: []byte("private provider state"), + }, + Messages: []api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Edit the account"}}}, + {Role: api.RoleAssistant, Parts: []api.Part{{Type: api.PartToolRequest, ToolRequest: &api.ToolRequest{ + ToolCallID: "call-account-approval", Name: "account_edit", Input: []byte(`{"id":"acc-1"}`), + }}}}, + }, + Calls: []api.ToolApprovalCall{{Request: api.ToolApprovalRequest{ + ToolCallID: "call-account-approval", Tool: "account_edit", Input: []byte(`{"id":"acc-1"}`), + }}}, + }, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(execution.Close(ctx)).To(Succeed()) + + run, err := db.GetPromptRun(ctx, uuid.MustParse(execution.PromptRunID())) + Expect(err).NotTo(HaveOccurred()) + Expect(run.State).To(Equal(database.PromptRunStateWaiting)) + requests, err := db.ListTurnRequests(ctx, database.TurnRequestFilter{SessionID: run.SessionID, PromptRunID: &run.ID}) + Expect(err).NotTo(HaveOccurred()) + Expect(requests).To(HaveLen(1)) + continuation, err := authority.ResolveToolApproval(ctx, aichat.ToolApprovalResolution{ + ThreadID: run.SessionID.String(), ApprovalID: requests[0].ID.String(), Approved: true, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(continuation).NotTo(BeNil()) + Expect(continuation.Execution.TurnID()).To(Equal(execution.TurnID())) + Expect(continuation.Spec.ToolApproval).NotTo(BeNil()) + Expect(continuation.Spec.ToolApproval.State.ProviderCheckpoint).NotTo(BeNil()) + Expect(continuation.Spec.ToolApproval.Decisions).To(HaveLen(1)) + Expect(continuation.Spec.ToolApproval.Decisions[0].ApprovalID).To(Equal(requests[0].ID.String())) + Expect(continuation.Spec.Messages).To(BeEmpty()) + Expect(continuation.Spec.Prompt.User).To(BeEmpty()) + Expect(continuation.Execution.Close(ctx)).To(Succeed()) + }) +}) + +func executionMCPClient(ctx context.Context, endpoint api.CallerToolEndpoint) *mcpclient.Client { + channel, err := transport.NewStreamableHTTP(endpoint.URL, transport.WithHTTPHeaders(endpoint.Headers)) + Expect(err).NotTo(HaveOccurred()) + client := mcpclient.NewClient(channel) + Expect(client.Start(ctx)).To(Succeed()) + request := mcp.InitializeRequest{} + request.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION + request.Params.ClientInfo = mcp.Implementation{Name: "captain-authority-test", Version: "1.0.0"} + _, err = client.Initialize(ctx, request) + Expect(err).NotTo(HaveOccurred()) + return client +} diff --git a/pkg/aichat/interrupt.go b/pkg/aichat/interrupt.go new file mode 100644 index 00000000..84b6a4bd --- /dev/null +++ b/pkg/aichat/interrupt.go @@ -0,0 +1,196 @@ +package aichat + +import ( + "context" + "errors" + "fmt" + "net/http" + "sync" + + "github.com/flanksource/captain/pkg/api" +) + +var errNoActiveTurn = errors.New("chat session has no active turn") + +type activeTurn struct { + provider api.StreamingProvider + execution Execution + ctx context.Context + cancel context.CancelFunc + + mu sync.Mutex + interrupting bool + interrupted bool + done bool + signal chan struct{} + aborted chan struct{} + emitted chan struct{} +} + +func newActiveTurn(ctx context.Context, provider api.StreamingProvider, execution Execution, cancel context.CancelFunc) *activeTurn { + return &activeTurn{ + provider: provider, execution: execution, ctx: ctx, cancel: cancel, + signal: make(chan struct{}), aborted: make(chan struct{}), emitted: make(chan struct{}), + } +} + +func (t *activeTurn) stream(source <-chan api.Event) <-chan api.Event { + out := make(chan api.Event) + go func() { + defer close(out) + emitInterrupted := func() { + select { + case out <- api.Event{Kind: api.EventInterrupted, Reason: "user"}: + case <-t.ctx.Done(): + } + close(t.emitted) + } + for { + select { + case <-t.signal: + emitInterrupted() + return + case event, ok := <-source: + if !ok { + t.mu.Lock() + interrupting := t.interrupting + if !interrupting { + t.done = true + } + t.mu.Unlock() + if interrupting { + select { + case <-t.signal: + emitInterrupted() + case <-t.aborted: + case <-t.ctx.Done(): + } + } + return + } + t.mu.Lock() + interrupted := t.interrupted + t.mu.Unlock() + if !interrupted { + select { + case out <- event: + case <-t.ctx.Done(): + return + } + } + case <-t.ctx.Done(): + return + } + } + }() + return out +} + +func (t *activeTurn) interrupt(ctx context.Context) error { + t.mu.Lock() + if t.done || t.interrupting || t.interrupted { + t.mu.Unlock() + return errNoActiveTurn + } + t.interrupting = true + t.mu.Unlock() + + if provider, ok := api.ProviderAs[api.InterruptibleProvider](t.provider); ok { + if err := provider.Interrupt(ctx); err != nil { + t.abortInterrupt() + return fmt.Errorf("interrupt provider turn: %w", err) + } + } + if t.execution != nil { + if err := t.execution.Interrupt(ctx, "user"); err != nil { + t.abortInterrupt() + return fmt.Errorf("interrupt authoritative execution: %w", err) + } + } + + t.mu.Lock() + t.interrupted = true + close(t.signal) + t.mu.Unlock() + select { + case <-t.emitted: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (t *activeTurn) abortInterrupt() { + t.mu.Lock() + t.interrupting = false + close(t.aborted) + t.mu.Unlock() +} + +func (t *activeTurn) finish() { + t.cancel() + t.mu.Lock() + t.done = true + t.mu.Unlock() +} + +func (s *Service) registerActiveTurn(threadID string, turn *activeTurn) error { + s.activeMu.Lock() + defer s.activeMu.Unlock() + if _, exists := s.active[threadID]; exists { + return fmt.Errorf("chat session %s already has an active turn", threadID) + } + s.active[threadID] = turn + return nil +} + +func (s *Service) unregisterActiveTurn(threadID string, turn *activeTurn) { + turn.finish() + s.activeMu.Lock() + if s.active[threadID] == turn { + delete(s.active, threadID) + } + s.activeMu.Unlock() +} + +func (s *Service) handleInterrupt(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w, request) + if store == nil { + return + } + threadID := request.PathValue("id") + if _, err := store.Get(request.Context(), threadID); err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + s.activeMu.Lock() + turn := s.active[threadID] + s.activeMu.Unlock() + if turn == nil { + http.Error(w, errNoActiveTurn.Error(), http.StatusConflict) + return + } + if err := turn.interrupt(request.Context()); err != nil { + status := http.StatusBadGateway + if errors.Is(err, errNoActiveTurn) { + status = http.StatusConflict + } + http.Error(w, err.Error(), status) + return + } + if sessions, ok := store.(SessionReader); ok { + aggregate, err := sessions.GetSession(request.Context(), threadID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + _ = writeJSON(w, http.StatusOK, aggregate) + return + } + thread, err := store.Get(request.Context(), threadID) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + _ = writeJSON(w, http.StatusOK, thread) +} diff --git a/pkg/aichat/messages.go b/pkg/aichat/messages.go index d7ae1f08..0665549f 100644 --- a/pkg/aichat/messages.go +++ b/pkg/aichat/messages.go @@ -69,19 +69,9 @@ func (s *Service) resolveAttachments(ctx context.Context, messages []UIMessage) } 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() + override, err := chatModel(request, settings.Spec.Model) if err != nil { - return api.Spec{}, fmt.Errorf("invalid chat model %q: %w", model, err) + return api.Spec{}, err } spec := settings.Spec.Merge(api.Spec{ Model: override, @@ -104,10 +94,21 @@ func requestSpec(request ChatRequest, settings RuntimeSettings, attachments map[ 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...) + if isAgentBackend(spec.Backend) { + user, promptAttachments, err := agentPrompt(messages, request.ProviderSessionID != "") + if err != nil { + return api.Spec{}, err + } + spec.Messages = nil + spec.Prompt.System = system + spec.Prompt.User = user + spec.Prompt.Attachments = promptAttachments + } else { + if system != "" { + messages = append([]api.Message{{Role: api.RoleSystem, Parts: []api.Part{{Type: api.PartText, Text: system}}}}, messages...) + } + spec.Messages = messages } - spec.Messages = messages } else { spec.Messages = nil } @@ -117,6 +118,48 @@ func requestSpec(request ChatRequest, settings RuntimeSettings, attachments map[ return spec, nil } +func chatModel(request ChatRequest, fallback api.Model) (api.Model, error) { + selected := fallback + if request.Runtime != nil { + selected = *request.Runtime + } else if model := strings.TrimSpace(request.Model); model != "" { + selected = api.Model{Name: model} + } + if strings.TrimSpace(selected.Name) == "" { + return api.Model{}, fmt.Errorf("chat model is required") + } + if request.ReasoningEffort != "" { + if selected.Effort != "" && selected.Effort != request.ReasoningEffort { + return api.Model{}, fmt.Errorf("chat runtime effort %q conflicts with reasoning effort %q", selected.Effort, request.ReasoningEffort) + } + selected.Effort = request.ReasoningEffort + } + if request.Temperature != nil { + if selected.Temperature != nil && *selected.Temperature != *request.Temperature { + return api.Model{}, fmt.Errorf("chat runtime temperature conflicts with request temperature") + } + selected.Temperature = request.Temperature + } + expanded, err := selected.Expand() + if err != nil { + return api.Model{}, fmt.Errorf("invalid chat runtime: %w", err) + } + if request.Runtime != nil && strings.TrimSpace(request.Model) != "" { + legacy, err := api.Model{Name: strings.TrimSpace(request.Model)}.Expand() + if err != nil { + return api.Model{}, fmt.Errorf("invalid chat model %q: %w", request.Model, err) + } + if legacy.Name != expanded.Name || legacy.Backend != expanded.Backend { + return api.Model{}, fmt.Errorf("chat model %q conflicts with structured runtime %s/%s", request.Model, expanded.Backend, expanded.Name) + } + } + return expanded, nil +} + +func isAgentBackend(backend api.Backend) bool { + return backend == api.BackendClaudeAgent || backend == api.BackendCodexAgent +} + func canonicalMessages(messages []UIMessage, attachments map[partLocation]api.AttachmentRef) ([]api.Message, error) { out := make([]api.Message, 0, len(messages)) for messageIndex, message := range messages { @@ -136,6 +179,9 @@ func canonicalMessages(messages []UIMessage, attachments map[partLocation]api.At } } if len(parts) == 0 { + if role == api.RoleAssistant { + continue + } return nil, fmt.Errorf("message %d (%s) has no provider content", messageIndex+1, role) } out = append(out, api.Message{Role: role, Parts: parts}) @@ -175,7 +221,7 @@ func canonicalPart(role api.MessageRole, part UIPart, attachment api.AttachmentR result.ToolResult.Output = nil result.ToolResult.Error = "tool execution denied" if part.Approval != nil && part.Approval.Reason != "" { - result.ToolResult.Error = part.Approval.Reason + result.ToolResult.Error += ": " + part.Approval.Reason } } return request, result, nil diff --git a/pkg/aichat/messages_ginkgo_test.go b/pkg/aichat/messages_ginkgo_test.go new file mode 100644 index 00000000..9872aec2 --- /dev/null +++ b/pkg/aichat/messages_ginkgo_test.go @@ -0,0 +1,105 @@ +package aichat_test + +import ( + "encoding/json" + "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" +) + +var _ = Describe("chat message projection", func() { + DescribeTable("omits assistant messages without provider content", + func(parts []aichat.UIPart) { + 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", + Messages: []aichat.UIMessage{ + {Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Review the transaction."}}}, + {Role: "assistant", Parts: parts}, + {Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Try again."}}}, + }, + })) + + Expect(response.Code).To(Equal(http.StatusOK), response.Body.String()) + Expect(provider.specs).To(HaveLen(1)) + Expect(provider.specs[0].Messages).To(Equal([]api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Review the transaction."}}}, + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Try again."}}}, + })) + }, + Entry("empty assistant shell", []aichat.UIPart{}), + Entry("step boundary only", []aichat.UIPart{{Type: "step-start"}}), + Entry("persisted error data only", []aichat.UIPart{{ + Type: "data-error", Data: json.RawMessage(`{"error":"provider disconnected"}`), + }}), + ) + + It("retains assistant provider content alongside UI-only parts", func() { + 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", + Messages: []aichat.UIMessage{ + {Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Review the transaction."}}}, + {Role: "assistant", Parts: []aichat.UIPart{ + {Type: "step-start"}, + {Type: "text", Text: "The stream started."}, + {Type: "data-result", Data: json.RawMessage(`{"success":true}`)}, + }}, + {Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Continue."}}}, + }, + })) + + Expect(response.Code).To(Equal(http.StatusOK), response.Body.String()) + Expect(provider.specs).To(HaveLen(1)) + Expect(provider.specs[0].Messages).To(Equal([]api.Message{ + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Review the transaction."}}}, + {Role: api.RoleAssistant, Parts: []api.Part{{Type: api.PartText, Text: "The stream started."}}}, + {Role: api.RoleUser, Parts: []api.Part{{Type: api.PartText, Text: "Continue."}}}, + })) + }) + + It("rejects a user message without provider content", func() { + provider := &fakeStreamingProvider{} + 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", + Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "step-start"}}}}, + })) + + Expect(response.Code).To(Equal(http.StatusBadRequest)) + Expect(response.Body.String()).To(ContainSubstring("message 1 (user) has no provider content")) + Expect(provider.specs).To(BeEmpty()) + }) + + It("rejects unsupported assistant parts", func() { + provider := &fakeStreamingProvider{} + 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", + Messages: []aichat.UIMessage{ + {Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Review the transaction."}}}, + {Role: "assistant", Parts: []aichat.UIPart{{Type: "finish-step"}}}, + {Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "Continue."}}}, + }, + })) + + Expect(response.Code).To(Equal(http.StatusBadRequest)) + Expect(response.Body.String()).To(ContainSubstring(`message 2 part 1: unsupported AI SDK part type "finish-step"`)) + Expect(provider.specs).To(BeEmpty()) + }) +}) diff --git a/pkg/aichat/persistence.go b/pkg/aichat/persistence.go index 676741f9..8e1df882 100644 --- a/pkg/aichat/persistence.go +++ b/pkg/aichat/persistence.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "reflect" "strings" "github.com/flanksource/captain/pkg/api" @@ -18,20 +19,18 @@ type assistantMessageBuilder struct { } type assistantMessageBuilderOptions struct { - ChatID string - Seed *UIMessage - Resume *api.ToolApprovalResume + MessageID string + TurnID string + Replace bool + Seed *UIMessage + Resume *api.ToolApprovalResume } func newAssistantMessageBuilder(options assistantMessageBuilderOptions) (*assistantMessageBuilder, error) { - id := "" - if options.ChatID != "" { - id = options.ChatID + "-assistant" - } builder := &assistantMessageBuilder{ - message: UIMessage{ID: id, Role: string(api.RoleAssistant), Parts: []UIPart{}}, + message: UIMessage{ID: options.MessageID, TurnID: options.TurnID, Role: string(api.RoleAssistant), Parts: []UIPart{}}, toolParts: map[string]int{}, - replace: options.Resume != nil, + replace: options.Replace || options.Resume != nil, } if options.Resume == nil { return builder, nil @@ -76,7 +75,18 @@ func newAssistantMessageBuilder(options assistantMessageBuilderOptions) (*assist return builder, nil } -func (s *Service) persistedEvents(ctx context.Context, request ChatRequest, source <-chan api.Event) <-chan api.Event { +// persistedEventOptions carries the turn identity a persisted stream needs: +// which thread and turn it belongs to, which model produced it (for pricing), +// and where to record the resulting costs for the finish part. +type persistedEventOptions struct { + Request ChatRequest + TurnID string + Model api.Model + Costs *TurnCosts +} + +func (s *Service) persistedEvents(ctx context.Context, options persistedEventOptions, source <-chan api.Event) <-chan api.Event { + request, turnID := options.Request, options.TurnID if request.ThreadID == "" { return source } @@ -88,8 +98,10 @@ func (s *Service) persistedEvents(ctx context.Context, request ChatRequest, sour sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error()}) return } + messageID := assistantMessageID(request, turnID) + replace := request.Trigger == "regenerate-message" builder, err := newAssistantMessageBuilder(assistantMessageBuilderOptions{ - ChatID: request.ID, Seed: seed, Resume: request.ToolApproval, + MessageID: messageID, TurnID: turnID, Replace: replace, Seed: seed, Resume: request.ToolApproval, }) if err != nil { sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error()}) @@ -102,13 +114,13 @@ func (s *Service) persistedEvents(ctx context.Context, request ChatRequest, sour return } if !persisted && event.Kind != api.EventResult && event.Kind != api.EventError && event.SessionID != "" { - if err := s.persistEvent(ctx, request.ThreadID, event); err != nil { + if err := s.persistEvent(ctx, request.ThreadID, event, options.Model, options.Costs); 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 { + if !persisted && (event.Kind == api.EventResult || event.Kind == api.EventError || event.Kind == api.EventInterrupted) { + if err := s.persistCompletedTurn(ctx, request.ThreadID, builder, event, options); err != nil { sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error(), Model: event.Model}) return } @@ -137,7 +149,11 @@ func (s *Service) approvalPersistenceSeed(ctx context.Context, request ChatReque return &last, nil } } - thread, err := s.options.Threads.Get(ctx, request.ThreadID) + store, err := s.threads(ctx) + if err != nil { + return nil, err + } + thread, err := store.Get(ctx, request.ThreadID) if err != nil { return nil, fmt.Errorf("load suspended assistant message: %w", err) } @@ -157,8 +173,14 @@ func sendEvent(ctx context.Context, target chan<- api.Event, event api.Event) bo } } -func (s *Service) persistCompletedTurn(ctx context.Context, threadID string, builder *assistantMessageBuilder, event api.Event) error { - if err := s.persistEvent(ctx, threadID, event); err != nil { +func (s *Service) persistCompletedTurn( + ctx context.Context, + threadID string, + builder *assistantMessageBuilder, + event api.Event, + options persistedEventOptions, +) error { + if err := s.persistEvent(ctx, threadID, event, options.Model, options.Costs); err != nil { return err } if len(builder.message.Parts) == 0 { @@ -171,10 +193,22 @@ func (s *Service) persistCompletedTurn(ctx context.Context, threadID string, bui } func (s *Service) persistAssistantMessage(ctx context.Context, threadID string, builder *assistantMessageBuilder) error { + store, err := s.threads(ctx) + if err != nil { + return err + } if builder.replace { - return s.options.Threads.ReplaceLastMessage(ctx, threadID, builder.message) + err = store.ReplaceLastMessage(ctx, threadID, builder.message) + } else { + err = store.AppendMessage(ctx, threadID, builder.message) + } + if err != nil { + return err } - return s.options.Threads.AppendMessage(ctx, threadID, builder.message) + // Backends that name sessions themselves (Claude's SessionTitle) report it as + // a tool call rather than through Captain's own tool handler. + s.setThreadTitle(ctx, threadID, TitleUpdate{Title: agentTitle(builder.message), Source: TitleSourceAI}) + return nil } func (b *assistantMessageBuilder) apply(event api.Event) error { @@ -198,16 +232,53 @@ func (b *assistantMessageBuilder) apply(event api.Event) error { case api.EventResult: return b.result(event) case api.EventError: + b.terminalizeTools(event.Error) 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.EventInterrupted: + return b.interrupted(event.Reason) case api.EventSystem: } return nil } +func (b *assistantMessageBuilder) interrupted(reason string) error { + b.terminalizeTools(reason) + data, err := json.Marshal(map[string]bool{"success": false, "interrupted": true}) + if err != nil { + return err + } + b.message.Parts = append(b.message.Parts, UIPart{Type: "data-result", Data: data}) + success := false + b.message.Metadata = &MessageMetadata{ + ProviderSessionID: b.sessionID, Model: b.model, Success: &success, Interrupted: true, + } + return nil +} + +func (b *assistantMessageBuilder) terminalizeTools(message string) { + if message == "" { + message = "tool execution did not complete" + } + for i := range b.message.Parts { + part := &b.message.Parts[i] + if !part.IsTool() { + continue + } + switch part.State { + case "output-available", "output-error", "output-denied": + continue + } + part.State = "output-error" + part.Output = nil + part.ErrorText = message + part.Approval = nil + } +} + func (b *assistantMessageBuilder) appendText(partType, text string) { if text == "" { return @@ -256,12 +327,15 @@ func (b *assistantMessageBuilder) toolUse(event api.Event) error { } func (b *assistantMessageBuilder) permission(event api.Event) error { + if event.ApprovalID == "" { + return fmt.Errorf("persist tool approval %q has no durable approval ID", event.ToolCallID) + } part, err := b.toolPart(event.ToolCallID, event.Tool) if err != nil { return err } part.State = "approval-requested" - part.Approval = &Approval{ID: event.ToolCallID} + part.Approval = &Approval{ID: event.ApprovalID} return nil } @@ -302,11 +376,10 @@ 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) + data, err = json.Marshal(map[string]bool{"success": event.Success, "waitingApproval": true}) if err != nil { - return fmt.Errorf("marshal tool approval state: %w", err) + return fmt.Errorf("marshal tool approval result state: %w", err) } } else if len(data) == 0 { var err error @@ -322,7 +395,7 @@ func (b *assistantMessageBuilder) result(event api.Event) error { } if event.Usage != nil { b.message.Metadata.Usage = usageMetadata(*event.Usage) - b.message.Metadata.ContextTokens = event.Usage.InputTokens + b.message.Metadata.ContextTokens = contextTokens(*event.Usage) } return nil } @@ -351,12 +424,15 @@ func applyApprovalDecision(part *UIPart, decision api.ToolApprovalDecision) erro case api.ToolApprovalApprove: approved := true part.State = "approval-responded" - part.Approval = &Approval{ID: decision.ToolCallID, Approved: &approved} + part.Approval = &Approval{ID: decision.ApprovalID, Approved: &approved} + if len(decision.Input) > 0 { + part.Input = append(json.RawMessage(nil), decision.Input...) + } case api.ToolApprovalDeny: approved := false part.State = "output-denied" part.Approval = &Approval{ - ID: decision.ToolCallID, Approved: &approved, Reason: decision.Message, + ID: decision.ApprovalID, Approved: &approved, Reason: decision.Message, } case api.ToolApprovalRespond: if decision.Result == nil { @@ -381,3 +457,13 @@ func jsonValue(text string) (json.RawMessage, error) { payload, err := json.Marshal(text) return payload, err } + +func equalPartJSON(left, right json.RawMessage) bool { + if len(left) == 0 || len(right) == 0 { + return len(left) == len(right) + } + var leftValue, rightValue any + return json.Unmarshal(left, &leftValue) == nil && + json.Unmarshal(right, &rightValue) == nil && + reflect.DeepEqual(leftValue, rightValue) +} diff --git a/pkg/aichat/provider_config.go b/pkg/aichat/provider_config.go index 0094108a..2d37b888 100644 --- a/pkg/aichat/provider_config.go +++ b/pkg/aichat/provider_config.go @@ -39,28 +39,57 @@ func (s *Service) annotateConfiguredModels(ctx context.Context, models ModelCata configured[ai.BackendToProvider(backend)] = true } for i := range models { - models[i].Configured = models[i].Configured || configured[models[i].Provider] + if configured[models[i].Provider] && models[i].Availability.State == api.AvailabilityMissingCredential { + models[i].Configured = true + models[i].Availability = api.Available() + } + } + return nil +} + +func (s *Service) annotateConfiguredRuntimes(ctx context.Context, runtimes []api.RuntimeFamily) 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[api.Backend]bool, len(backends)) + for _, backend := range backends { + if backend == "" { + return fmt.Errorf("configured chat provider backend is required") + } + configured[backend] = true + } + for familyIndex := range runtimes { + for modeIndex := range runtimes[familyIndex].Modes { + mode := &runtimes[familyIndex].Modes[modeIndex] + if configured[api.Backend(mode.Backend)] && mode.Availability.State == api.AvailabilityMissingCredential { + mode.Availability = api.Available() + } + } } return nil } -func (s *Service) resolveProvider(ctx context.Context, config api.Config) (api.StreamingProvider, error) { +func (s *Service) prepareProviderConfig(ctx context.Context, config api.Config) (api.Config, error) { if s.options.ProviderConfig != nil { resolved, err := ai.ResolveModelSelectors(config.Model) if err != nil { - return nil, fmt.Errorf("resolve chat model: %w", err) + return api.Config{}, 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) + return api.Config{}, 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)", + return api.Config{}, 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) + return config, nil } diff --git a/pkg/aichat/resolver.go b/pkg/aichat/resolver.go index 9db27023..d7dd818e 100644 --- a/pkg/aichat/resolver.go +++ b/pkg/aichat/resolver.go @@ -14,6 +14,7 @@ import ( // canonical resolver; tests may replace it with a fake provider. type Resolver interface { Models(context.Context) (ModelCatalogResponse, error) + Runtimes(context.Context) ([]api.RuntimeFamily, error) Provider(context.Context, api.Config) (api.StreamingProvider, error) } @@ -35,6 +36,10 @@ func (captainResolver) Models(_ context.Context) (ModelCatalogResponse, error) { return ai.LiveCatalogInfo(configured) } +func (captainResolver) Runtimes(_ context.Context) ([]api.RuntimeFamily, error) { + return ai.LiveRuntimeCatalog() +} + func (captainResolver) Provider(_ context.Context, config api.Config) (api.StreamingProvider, error) { provider, err := ai.NewProvider(config) if err != nil { diff --git a/pkg/aichat/service.go b/pkg/aichat/service.go index fd1ecab0..d8ec5944 100644 --- a/pkg/aichat/service.go +++ b/pkg/aichat/service.go @@ -6,7 +6,10 @@ import ( "fmt" "net/http" "strings" + "sync" + "time" + "github.com/flanksource/captain/pkg/ai" aitools "github.com/flanksource/captain/pkg/ai/tools" "github.com/flanksource/captain/pkg/api" "github.com/flanksource/commons/logger" @@ -45,6 +48,24 @@ func (f RuntimeSettingsProviderFunc) RuntimeSettings(ctx context.Context) (Runti return f(ctx) } +// ThreadStoreProvider supplies the request-scoped thread store. Applications +// that can serve more than one database resolve it per request; a fixed store +// is expressed as a provider that ignores the context. +type ThreadStoreProvider interface { + ThreadStore(context.Context) (ThreadStore, error) +} + +type ThreadStoreProviderFunc func(context.Context) (ThreadStore, error) + +func (f ThreadStoreProviderFunc) ThreadStore(ctx context.Context) (ThreadStore, error) { + return f(ctx) +} + +// FixedThreadStore adapts a single store to the provider interface. +func FixedThreadStore(store ThreadStore) ThreadStoreProvider { + return ThreadStoreProviderFunc(func(context.Context) (ThreadStore, error) { return store, nil }) +} + // ServiceOptions injects every application-owned chat dependency. A nil // Resolver uses Captain's canonical model/provider resolver. type ServiceOptions struct { @@ -54,13 +75,16 @@ type ServiceOptions struct { Tools ToolProvider MCP ToolProvider Attachments AttachmentResolver - Threads ThreadStore + Threads ThreadStoreProvider + Authority ExecutionAuthority } // Service is Captain's AI SDK-compatible HTTP chat service. type Service struct { options ServiceOptions resolver Resolver + activeMu sync.Mutex + active map[string]*activeTurn } func NewService(options ServiceOptions) *Service { @@ -68,18 +92,34 @@ func NewService(options ServiceOptions) *Service { if resolver == nil { resolver = captainResolver{} } - return &Service{options: options, resolver: resolver} + return &Service{options: options, resolver: resolver, active: map[string]*activeTurn{}} } 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/runtimes", s.handleRuntimes) mux.HandleFunc("GET /api/chat/tools", s.handleTools) s.registerThreadRoutes(mux) return mux } +func (s *Service) handleRuntimes(w http.ResponseWriter, request *http.Request) { + runtimes, err := s.resolver.Runtimes(request.Context()) + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + if err := s.annotateConfiguredRuntimes(request.Context(), runtimes); err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + if err := writeJSON(w, http.StatusOK, runtimes); err != nil { + serviceLog.Errorf("write chat runtimes response: %v", err) + } +} + func (s *Service) handleModels(w http.ResponseWriter, request *http.Request) { models, err := s.resolver.Models(request.Context()) if err != nil { @@ -112,7 +152,8 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { http.Error(w, fmt.Sprintf("invalid chat request: %v", err), http.StatusBadRequest) return } - if err := resolveToolApproval(&chat); err != nil { + turnID, err := chatTurnID(chat) + if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } @@ -125,7 +166,12 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { http.Error(w, err.Error(), requestErrorStatus(err)) return } - if err := s.resolveThreadSession(request.Context(), &chat); err != nil { + thread, err := s.resolveThreadSession(request.Context(), &chat) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + if err := validateThreadTurn(chat, thread); err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } @@ -147,16 +193,62 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { http.Error(w, err.Error(), http.StatusInternalServerError) return } - if err := s.persistIncoming(request.Context(), chat); err != nil { + definitions, err := aitools.ResolveDefinitions(set.Definitions, spec.ToolPreferences) + if err != nil { http.Error(w, err.Error(), http.StatusBadRequest) return } + // Bound to this thread and appended after resolution: the conversation's own + // name is not a preference the user manages. Chats that carry no caller tools + // at all stay that way — they are named from their opening message instead. + if chat.ThreadID != "" && len(definitions) > 0 && s.options.Threads != nil { + definitions = append(definitions, s.sessionTitleTool(chat.ThreadID)) + appendSessionTitleInstruction(&spec) + } 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) + config.CaptainSessionID = chat.ThreadID + config.Tools = definitions + config, err = s.prepareProviderConfig(request.Context(), config) + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + spec.Model = config.Model + var execution Execution + var callerToolEvents <-chan api.Event + if s.options.Authority != nil && chat.ThreadID != "" { + title := "" + if thread != nil { + title = thread.Title + } + execution, err = s.options.Authority.Begin(request.Context(), ExecutionRequest{ + ThreadID: chat.ThreadID, RequestID: turnID, Title: title, + Spec: spec, Definitions: definitions, + }) + if err != nil { + http.Error(w, fmt.Sprintf("admit chat execution: %v", err), http.StatusInternalServerError) + return + } + defer closeExecution(execution) + turnID = execution.TurnID() + if chat.Trigger == "submit-message" && chat.MessageID == "" && len(chat.Messages) > 0 { + chat.Messages[len(chat.Messages)-1].TurnID = turnID + } + config.CaptainSessionID = execution.CaptainSessionID() + config.CallerTools = execution.CallerTools() + if config.CallerTools != nil { + callerToolEvents = execution.Events() + } + } + if len(definitions) > 0 && isAgentBackend(config.Model.Backend) && + (execution == nil || config.CallerTools == nil) { + http.Error(w, "agent caller tools require an authoritative Captain execution", http.StatusServiceUnavailable) + return + } + provider, err := s.resolver.Provider(request.Context(), config) if err != nil { http.Error(w, err.Error(), http.StatusServiceUnavailable) return @@ -166,13 +258,17 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { serviceLog.Errorf("close chat provider: %v", closeErr) } }() - if len(set.Definitions) > 0 { + if len(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 } } + if err := s.persistIncoming(request.Context(), chat); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } streamContext, cancel := context.WithCancel(request.Context()) defer cancel() events, err := provider.ExecuteStream(streamContext, spec) @@ -180,18 +276,95 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { http.Error(w, err.Error(), http.StatusBadGateway) return } + events = mergeExecutionEvents(streamContext, events, callerToolEvents, definitions) + if chat.ThreadID != "" { + active := newActiveTurn(streamContext, provider, execution, cancel) + if err := s.registerActiveTurn(chat.ThreadID, active); err != nil { + http.Error(w, err.Error(), http.StatusConflict) + return + } + defer s.unregisterActiveTurn(chat.ThreadID, active) + events = active.stream(events) + } + events = observeExecutionEvents(streamContext, execution, events) writer, err := NewSSEWriter(w) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } - if err := WriteEventStream(writer, s.persistedEvents(streamContext, chat, events), EventStreamOptions{ + costs := &TurnCosts{} + persisted := s.persistedEvents(streamContext, persistedEventOptions{ + Request: chat, TurnID: turnID, Model: config.Model, Costs: costs, + }, events) + if err := WriteEventStream(writer, persisted, EventStreamOptions{ ToolApproval: chat.ToolApproval, + MessageID: assistantMessageID(chat, turnID), + Costs: costs, }); err != nil { serviceLog.Errorf("stream chat response: %v", err) } } +func assistantMessageID(request ChatRequest, turnID string) string { + if request.MessageID != "" { + return request.MessageID + } + if turnID != "" { + return turnID + "-assistant" + } + return "" +} + +func chatTurnID(request ChatRequest) (string, error) { + if request.ThreadID == "" { + return "", nil + } + if request.ID != request.ThreadID { + return "", fmt.Errorf("chat id %q must match threadId %q", request.ID, request.ThreadID) + } + switch request.Trigger { + case "submit-message": + if request.MessageID != "" { + return "", fmt.Errorf( + "submit-message cannot include messageId %q; resolve approvals through /api/chat/sessions/{id}/approvals/{approvalID}", + request.MessageID, + ) + } + if len(request.Messages) == 0 { + return "", fmt.Errorf("submit-message requires a final user message") + } + last := request.Messages[len(request.Messages)-1] + if !strings.EqualFold(last.Role, string(api.RoleUser)) { + return "", fmt.Errorf("submit-message must end with a user message") + } + if last.ID == "" { + return "", fmt.Errorf("submit-message final user message requires an id") + } + return last.ID, nil + case "regenerate-message": + if request.MessageID == "" { + return "", fmt.Errorf("regenerate-message requires messageId") + } + return request.MessageID, nil + default: + return "", fmt.Errorf("unsupported chat trigger %q", request.Trigger) + } +} + +func validateThreadTurn(request ChatRequest, thread *Thread) error { + if request.Trigger != "regenerate-message" || thread == nil { + return nil + } + if len(thread.Messages) == 0 { + return fmt.Errorf("regenerate-message messageId %q has no persisted assistant message", request.MessageID) + } + last := thread.Messages[len(thread.Messages)-1] + if !strings.EqualFold(last.Role, string(api.RoleAssistant)) || last.ID != request.MessageID { + return fmt.Errorf("regenerate-message messageId %q must match the final persisted assistant message", request.MessageID) + } + return nil +} + func (s *Service) runtimeSettings(ctx context.Context) (RuntimeSettings, error) { if s.options.Settings == nil { return RuntimeSettings{}, nil @@ -199,44 +372,75 @@ func (s *Service) runtimeSettings(ctx context.Context) (RuntimeSettings, error) return s.options.Settings.RuntimeSettings(ctx) } -func (s *Service) resolveThreadSession(ctx context.Context, request *ChatRequest) error { +func (s *Service) resolveThreadSession(ctx context.Context, request *ChatRequest) (*Thread, error) { if request.ThreadID == "" { - return nil + return nil, nil } - if s.options.Threads == nil { - return fmt.Errorf("thread persistence is not configured") + store, err := s.threads(ctx) + if err != nil { + return nil, err } - thread, err := s.options.Threads.Get(ctx, request.ThreadID) + thread, err := store.Get(ctx, request.ThreadID) if err != nil { - return err + return nil, err } if request.ProviderSessionID == "" { request.ProviderSessionID = thread.ProviderSessionID } - return nil + return thread, nil +} + +func closeExecution(execution Execution) { + if execution == nil { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := execution.Close(ctx); err != nil { + serviceLog.Errorf("close authoritative chat execution: %v", err) + } } func (s *Service) persistIncoming(ctx context.Context, request ChatRequest) error { - if request.ThreadID == "" || len(request.Messages) == 0 { + if request.ThreadID == "" || request.Trigger != "submit-message" || request.MessageID != "" || 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) + if !strings.EqualFold(last.Role, string(api.RoleUser)) { + return nil + } + store, err := s.threads(ctx) + if err != nil { + return err } + if err := store.AppendMessage(ctx, request.ThreadID, last); err != nil { + return err + } + // Names an as-yet-unnamed thread after the message that opened it. The store + // keeps this from displacing a title the agent or the user already chose. + s.setThreadTitle(ctx, request.ThreadID, TitleUpdate{ + Title: derivedTitle(request.Messages), Source: TitleSourceDerived, + }) return nil } -func (s *Service) persistEvent(ctx context.Context, threadID string, event api.Event) error { +// persistEvent accrues a completed turn against its thread. The thread returned +// by AddUsage carries the conversation's running total, which is recorded on +// costs so the finish part can report cumulative rather than per-turn spend. +func (s *Service) persistEvent(ctx context.Context, threadID string, event api.Event, model api.Model, costs *TurnCosts) error { + store, err := s.threads(ctx) + if err != nil { + return err + } if event.SessionID != "" { - if err := s.options.Threads.SetProviderSession(ctx, threadID, event.SessionID); err != nil { + if err := store.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{ + thread, err := store.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, @@ -244,5 +448,26 @@ func (s *Service) persistEvent(ctx context.Context, threadID string, event api.E if err != nil { return fmt.Errorf("persist thread usage: %w", err) } + if costs != nil { + costs.Breakdown = costBreakdownMetadata(model, *event.Usage, event.CostUSD) + if thread != nil { + costs.ThreadCostUSD = thread.TotalCostUSD + } + } return nil } + +func costBreakdownMetadata(model api.Model, usage api.Usage, providerCostUSD float64) *CostBreakdownMetadata { + cost := ai.PriceUsage(model.Backend, model.Name, usage, providerCostUSD) + return &CostBreakdownMetadata{ + Model: cost.Model, + InputUSD: cost.InputCost, + OutputUSD: cost.OutputCost, + ReasoningUSD: cost.ReasoningCost, + CacheReadUSD: cost.CacheReadCost, + // genkit reports no cache-write tokens on the API backends, so this + // stays zero there rather than being silently omitted. + CacheWriteUSD: cost.CacheWriteCost, + TotalUSD: cost.Total(), + } +} diff --git a/pkg/aichat/service_ginkgo_test.go b/pkg/aichat/service_ginkgo_test.go index 299ad6f5..2af07692 100644 --- a/pkg/aichat/service_ginkgo_test.go +++ b/pkg/aichat/service_ginkgo_test.go @@ -1,7 +1,6 @@ package aichat_test import ( - "bytes" "context" "encoding/json" "fmt" @@ -17,6 +16,7 @@ import ( type fakeResolver struct { models aichat.ModelCatalogResponse + runtimes []api.RuntimeFamily provider *fakeStreamingProvider configs []api.Config } @@ -47,15 +47,22 @@ func (f *fakeResolver) Models(context.Context) (aichat.ModelCatalogResponse, err return f.models, nil } +func (f *fakeResolver) Runtimes(context.Context) ([]api.RuntimeFamily, error) { + return f.runtimes, 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) + events []api.Event + specs []api.Spec + execute func(context.Context, api.Spec) (<-chan api.Event, error) + backend api.Backend + supportsCallerTools *bool + interrupt func(context.Context) error } func (f *fakeStreamingProvider) Execute(context.Context, api.Spec) (*api.Response, error) { @@ -75,34 +82,29 @@ func (f *fakeStreamingProvider) ExecuteStream(ctx context.Context, spec api.Spec 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")}) +func (f *fakeStreamingProvider) GetModel() string { return "test-model" } +func (f *fakeStreamingProvider) GetBackend() api.Backend { + if f.backend != "" { + return f.backend } - return refs, nil + return api.BackendOpenAI +} +func (f *fakeStreamingProvider) SupportsCallerTools() bool { + return f.supportsCallerTools == nil || *f.supportsCallerTools } -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) +func (f *fakeStreamingProvider) Interrupt(ctx context.Context) error { + if f.interrupt == nil { + return nil + } + return f.interrupt(ctx) } 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"}, + {ID: "anthropic/claude-sonnet", Provider: "anthropic", Label: "Claude", Availability: api.Availability{State: api.AvailabilityMissingCredential, Reason: "No Claude API credentials.", Remediation: "Configure credentials."}}, + {ID: "openai/gpt", Provider: "openai", Label: "GPT", Availability: api.Availability{State: api.AvailabilityMissingCredential, Reason: "No OpenAI API credentials.", Remediation: "Configure credentials."}}, }} source := &fakeProviderConfigSource{backends: []api.Backend{api.BackendOpenAI}} service := aichat.NewService(aichat.ServiceOptions{Resolver: resolver, ProviderConfig: source}) @@ -115,6 +117,28 @@ var _ = Describe("Captain aichat service", func() { Expect(models).To(HaveLen(2)) Expect(models[0].Configured).To(BeFalse()) Expect(models[1].Configured).To(BeTrue()) + Expect(models[0].Availability.State).To(Equal(api.AvailabilityMissingCredential)) + Expect(models[1].Availability).To(Equal(api.Available())) + }) + + It("annotates runtime modes with request-scoped configured providers", func() { + resolver := &fakeResolver{runtimes: []api.RuntimeFamily{{ + Family: "codex", Provider: "openai", CatalogPrefix: "openai", + Modes: []api.RuntimeModeEntry{{ + Mode: "api", Backend: string(api.BackendOpenAI), Kind: "api", + Availability: api.Availability{State: api.AvailabilityMissingCredential, Reason: "No OpenAI API credentials.", Remediation: "Configure credentials."}, + }}, + }}} + 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/runtimes", nil)) + Expect(response.Code).To(Equal(http.StatusOK)) + var runtimes []api.RuntimeFamily + Expect(json.Unmarshal(response.Body.Bytes(), &runtimes)).To(Succeed()) + Expect(runtimes).To(HaveLen(1)) + Expect(runtimes[0].Modes[0].Availability).To(Equal(api.Available())) }) It("applies request-scoped credentials after canonical model selection", func() { @@ -219,6 +243,8 @@ var _ = Describe("Captain aichat service", func() { 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, + Availability: api.Available(), + Runtime: api.Model{Name: "test-model", Backend: api.BackendOpenAI}, }}} service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, @@ -240,7 +266,7 @@ var _ = Describe("Captain aichat service", func() { 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}]`)) + Expect(models.Body.String()).To(MatchJSON(`[{"id":"openai/test-model","provider":"openai","label":"Test","runtime":{"model":"test-model","backend":"openai"},"reasoning":false,"temperature":false,"configured":true,"availability":{"state":"available"},"contextWindow":0,"inputMediaTypes":null}]`)) tools := httptest.NewRecorder() service.Handler().ServeHTTP(tools, httptest.NewRequest(http.MethodGet, "/api/chat/tools", nil)) @@ -309,48 +335,106 @@ var _ = Describe("Captain aichat service", func() { 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"}`), - }}}, + It("adapts canonical chat messages into an agent prompt", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "Agent chat") + Expect(err).NotTo(HaveOccurred()) + provider := &fakeStreamingProvider{ + backend: api.BackendClaudeAgent, + events: []api.Event{{Kind: api.EventResult, Success: true}}, } - 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}}) + resolver := &fakeResolver{provider: provider} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, Threads: aichat.FixedThreadStore(store), + Settings: aichat.RuntimeSettingsProviderFunc(func(context.Context) (aichat.RuntimeSettings, error) { + return aichat.RuntimeSettings{System: "Use accounting tools."}, nil + }), + }) + response := httptest.NewRecorder() service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ - Model: "openai/test-model", ToolApproval: resume, + ID: thread.ID, + Trigger: "submit-message", + Runtime: &api.Model{Name: "sonnet", Backend: api.BackendClaudeAgent}, + ThreadID: thread.ID, + ProviderSessionID: "provider-session-1", + Messages: []aichat.UIMessage{{ + ID: "message-agent-user", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "inspect the invoice"}}, + }}, })) 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()) + Expect(provider.specs[0].Prompt.System).To(Equal("Use accounting tools.")) + Expect(provider.specs[0].Prompt.User).To(Equal("inspect the invoice")) + Expect(resolver.configs[0].Model.Backend).To(Equal(api.BackendClaudeAgent)) + Expect(resolver.configs[0].CaptainSessionID).To(Equal(thread.ID)) + Expect(resolver.configs[0].SessionID).To(Equal("provider-session-1")) + }) + + It("treats an all-off resolved tool set as no caller tools", func() { + supported := false + provider := &fakeStreamingProvider{ + events: []api.Event{{Kind: api.EventResult, Success: true}}, + supportsCallerTools: &supported, + } + resolver := &fakeResolver{provider: provider} + service := aichat.NewService(aichat.ServiceOptions{ + Resolver: resolver, + Tools: aichat.StaticToolProvider([]api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}), + }) + + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + Model: "openai/test-model", + ToolPreferences: api.ToolPreferences{"invoice_get": api.ToolModeOff}, + Messages: []aichat.UIMessage{{ + Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "hello"}}, + }}, + })) + + Expect(response.Code).To(Equal(http.StatusOK)) + Expect(resolver.configs).To(HaveLen(1)) + Expect(resolver.configs[0].Tools).To(BeEmpty()) + }) + + It("rejects AI SDK approval continuations outside the Captain session endpoint", func() { + provider := &fakeStreamingProvider{} + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}}) + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: "session-1", ThreadID: "session-1", Trigger: "submit-message", MessageID: "assistant-1", + Messages: []aichat.UIMessage{{ + ID: "assistant-1", Role: "assistant", Parts: []aichat.UIPart{{ + Type: "dynamic-tool", ToolName: "invoice_pay", ToolCallID: "call-1", + State: "approval-responded", Approval: &aichat.Approval{ID: "approval-1"}, + }}, + }}, + })) + + Expect(response.Code).To(Equal(http.StatusBadRequest)) + Expect(response.Body.String()).To(ContainSubstring("resolve approvals through /api/chat/sessions/{id}/approvals/{approvalID}")) + Expect(provider.specs).To(BeEmpty()) }) It("serves thread CRUD through the injected persistence store", func() { - service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{}, Threads: aichat.NewMemoryThreadStore()}) + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{}, Threads: aichat.FixedThreadStore(aichat.NewMemoryThreadStore())}) create := httptest.NewRecorder() - service.Handler().ServeHTTP(create, requestJSON(http.MethodPost, "/api/chat/threads", map[string]string{"title": "Review"})) + service.Handler().ServeHTTP(create, requestJSON(http.MethodPost, "/api/chat/sessions", 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)) + service.Handler().ServeHTTP(get, httptest.NewRequest(http.MethodGet, "/api/chat/sessions/"+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)) + service.Handler().ServeHTTP(remove, httptest.NewRequest(http.MethodDelete, "/api/chat/sessions/"+thread.ID, nil)) Expect(remove.Code).To(Equal(http.StatusNoContent)) }) @@ -365,11 +449,11 @@ var _ = Describe("Captain aichat service", func() { {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}) + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(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"}}}}, + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", Model: "openai/test-model", + Messages: []aichat.UIMessage{{ID: "message-review-user", Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "inspect"}}}}, })) Expect(response.Code).To(Equal(http.StatusOK)) @@ -377,6 +461,7 @@ var _ = Describe("Captain aichat service", func() { Expect(err).NotTo(HaveOccurred()) Expect(stored.Messages).To(HaveLen(2)) assistant := stored.Messages[1] + Expect(assistant.ID).To(Equal("message-review-user-assistant")) Expect(assistant.Role).To(Equal("assistant")) Expect(assistant.Parts).To(HaveLen(3)) Expect(assistant.Parts[0].Type).To(Equal("text")) @@ -418,11 +503,11 @@ var _ = Describe("Captain aichat service", func() { }() return events, nil } - service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}, Threads: store}) + service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(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"}}}}, + ID: thread.ID, ThreadID: thread.ID, Trigger: "submit-message", Model: "openai/test-model", + Messages: []aichat.UIMessage{{ID: "message-cancel-user", 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/service_helpers_ginkgo_test.go b/pkg/aichat/service_helpers_ginkgo_test.go new file mode 100644 index 00000000..18edeff6 --- /dev/null +++ b/pkg/aichat/service_helpers_ginkgo_test.go @@ -0,0 +1,33 @@ +package aichat_test + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + + . "github.com/onsi/gomega" + + "github.com/flanksource/captain/pkg/aichat" + "github.com/flanksource/captain/pkg/api" +) + +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) +} diff --git a/pkg/aichat/session_overview.go b/pkg/aichat/session_overview.go new file mode 100644 index 00000000..49a3aa5e --- /dev/null +++ b/pkg/aichat/session_overview.go @@ -0,0 +1,226 @@ +package aichat + +import ( + "context" + "sort" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" +) + +// OverviewProjectionStore is the read surface the cross-branch projection needs. +type OverviewProjectionStore interface { + ListThreadSessionOverviews(context.Context, uuid.UUID) ([]database.SessionOverview, error) + ListPlans(context.Context, database.PlanFilter) ([]database.Plan, error) +} + +// ApplyOverviewProjection fills a session aggregate with everything that is a +// property of the stored row rather than of the branch that produced it: the +// monitor-owned metadata projection, the git blob, the context window, the +// authoritative plan, and the thread file rollup. +// +// Three branches build the same aggregate — the database branch, a transcript +// re-parse, and a prompt run — and each used to carry its own subset, so a +// session's changed files, git branch and plan appeared or vanished depending on +// which one served the request. Fields already set win: a transcript-derived +// value is fresher than the stored copy. +// +// Approvals are deliberately not projected. applyRequestState derives those from +// captain_turn_requests, where the stored transcript copy counts every +// operational tool use as an approval and reads as "200 approved". +func ApplyOverviewProjection( + ctx context.Context, + db OverviewProjectionStore, + overview database.SessionOverview, + detail *session.Session, +) error { + applyOverviewMetadata(overview, detail) + + plans, err := db.ListPlans(ctx, database.PlanFilter{SourceSessionID: &overview.ID}) + if err != nil { + return err + } + // captain_plans holds the approved revision, which the transcript copy does + // not know about, so it outranks whatever the branch already found. + if plan := planFromNative(plans); plan != nil { + detail.Plan = plan + } + + rootID := overview.ID + if overview.RootSessionID != nil { + rootID = *overview.RootSessionID + } + thread, err := db.ListThreadSessionOverviews(ctx, rootID) + if err != nil { + return err + } + detail.Files = threadFiles(overview.ID, detail.Files, thread) + return nil +} + +// applyOverviewMetadata fills the fields the stored row carries directly, +// leaving anything the branch already resolved alone. +func applyOverviewMetadata(overview database.SessionOverview, detail *session.Session) { + metadata := session.DecodeMetadata(overview.Metadata) + if detail.Model == "" { + detail.Model = metadata.Model + } + if detail.Provider == "" { + detail.Provider = metadata.Provider + } + if len(detail.Files.Read) == 0 && len(detail.Files.Written) == 0 { + detail.Files = metadata.Files + } + if len(detail.Todos) == 0 { + detail.Todos = metadata.Todos + } + if detail.Plan == nil { + detail.Plan = metadata.Plan + } + if detail.Git == (session.GitState{}) { + detail.Git = session.DecodeGitState(overview.Git) + } + if detail.Context == nil { + detail.Context = overviewContext(overview) + } +} + +// overviewContext projects the context-window reading the overview already +// computes for the session list, so the detail view stops reporting none while +// the list row beside it reports a percentage. +func overviewContext(overview database.SessionOverview) *session.Context { + if overview.ContextTokens == nil && overview.ContextWindowTokens == nil && overview.ContextFreePercent == nil { + return nil + } + context := &session.Context{} + if overview.ContextTokens != nil { + context.UsedTokens = int(*overview.ContextTokens) + } + if overview.ContextWindowTokens != nil { + context.WindowTokens = int(*overview.ContextWindowTokens) + } + if overview.ContextFreePercent != nil { + context.FreePercent = *overview.ContextFreePercent + } + return context +} + +// projectSessionAgents rebuilds the sub-agent hierarchy from the thread's agent +// rows, returning the root node and the flat index (root first) the transcript +// parser produces. Rows arrive root-first and parent-before-child from +// ListThreadAgents' ordering, but parentage is resolved through the index so a +// child whose parent is missing from the slice still lands in the flat list +// rather than disappearing. +func projectSessionAgents(rows []database.SessionAgent) (*session.Agent, []*session.Agent) { + if len(rows) == 0 { + return nil, nil + } + byID := make(map[string]*session.Agent, len(rows)) + agents := make([]*session.Agent, 0, len(rows)) + for _, row := range rows { + agent := &session.Agent{ + ID: row.SessionID.String(), Type: stringPointer(row.AgentType), Desc: stringPointer(row.Description), + IsRoot: row.IsRoot, HistoryFile: stringPointer(row.HistoryFile), + Usage: api.Usage{ + InputTokens: int(row.InputTokens), OutputTokens: int(row.OutputTokens), + ReasoningTokens: int(row.ReasoningTokens), CacheReadTokens: int(row.CacheReadTokens), + CacheWriteTokens: int(row.CacheWriteTokens), + }, + Cost: api.Cost{ + InputTokens: int(row.InputTokens), OutputTokens: int(row.OutputTokens), + ReasoningTokens: int(row.ReasoningTokens), CacheReadTokens: int(row.CacheReadTokens), + CacheWriteTokens: int(row.CacheWriteTokens), TotalTokens: int(row.TotalTokens), + ProviderCostUSD: row.CostUSD, + }, + } + if row.ParentSessionID != nil { + agent.ParentID = row.ParentSessionID.String() + } + byID[agent.ID] = agent + agents = append(agents, agent) + } + + var root *session.Agent + for _, agent := range agents { + if agent.IsRoot && root == nil { + root = agent + } + if parent, ok := byID[agent.ParentID]; ok && agent.ParentID != agent.ID { + parent.Children = append(parent.Children, agent) + } + } + return root, agents +} + +// threadFiles rolls a parent's changed files up from the sub-agents that did the +// work. A session that spawns sub-agents edits nothing itself — the row exists +// to own the thread — so reporting only its own (empty) set hides the whole +// thread's output. A leaf keeps its own set, or the hierarchy would stop +// distinguishing which sub-agent touched what. +func threadFiles(sessionID uuid.UUID, own session.ChangedFiles, thread []database.SessionOverview) session.ChangedFiles { + children := map[uuid.UUID][]database.SessionOverview{} + for _, row := range thread { + if row.ParentSessionID != nil { + children[*row.ParentSessionID] = append(children[*row.ParentSessionID], row) + } + } + read, written := own.Read, own.Written + descendants := 0 + for queue := children[sessionID]; len(queue) > 0; { + row := queue[0] + queue = append(queue[1:], children[row.ID]...) + descendants++ + files := session.DecodeMetadata(row.Metadata).Files + read, written = append(read, files.Read...), append(written, files.Written...) + } + if descendants == 0 { + return own + } + return session.ChangedFiles{Read: sortedUnique(read), Written: sortedUnique(written)} +} + +func sortedUnique(values []string) []string { + if len(values) == 0 { + return nil + } + seen := make(map[string]struct{}, len(values)) + unique := make([]string, 0, len(values)) + for _, value := range values { + if _, ok := seen[value]; ok || value == "" { + continue + } + seen[value] = struct{}{} + unique = append(unique, value) + } + sort.Strings(unique) + return unique +} + +// planFromNative projects the authoritative plan for a session: the approved +// revision when one exists, otherwise the latest. Plans are ordered newest +// first, so the first approved row wins and an unapproved plan only supplies +// content when nothing was ever approved. +func planFromNative(plans []database.Plan) *session.Plan { + var selected *database.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 + } + revision := selected.ApprovedRevision + if revision == nil { + revision = selected.LatestRevision + } + return &session.Plan{ + Path: selected.Path, Slug: selected.Slug, Content: revision.PlanMarkdown, Explicit: true, + } +} diff --git a/pkg/aichat/session_overview_ginkgo_test.go b/pkg/aichat/session_overview_ginkgo_test.go new file mode 100644 index 00000000..832640fa --- /dev/null +++ b/pkg/aichat/session_overview_ginkgo_test.go @@ -0,0 +1,133 @@ +package aichat + +import ( + "github.com/flanksource/captain/pkg/claude/tools" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" + "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/segmentio/encoding/json" +) + +// Three branches build the same session aggregate — the database branch, a +// transcript re-parse, and a prompt run — so anything only one of them maps +// appears or vanishes depending on which served the request. An ingested session +// takes the database branch, which is why its changed files, git branch and +// context window used to disappear the moment it gained messages. +var _ = ginkgo.Describe("applyOverviewMetadata", func() { + contextTokens := int64(48_000) + windowTokens := int64(200_000) + freePercent := 76 + + overview := func() database.SessionOverview { + return database.SessionOverview{ + ID: uuid.MustParse("860305d7-e8cd-41b8-b5ee-332dc74d4a41"), + Git: json.RawMessage(`{"branch":"feat/todo-prompt-runs","commit":"0c6f277","worktree":"/tmp/wt"}`), + Metadata: json.RawMessage(`{ + "model": "gpt-5-codex", + "provider": "openai", + "files": {"read": ["todos/plans.go"], "written": ["todos/outcome.go", "todos/provider.go"]}, + "todos": [{"text": "ship it", "status": "pending"}], + "plan": {"path": "/plans/stored.md", "slug": "stored"}, + "tags": ["ignored-sibling-key"] + }`)} + } + + project := func(overview database.SessionOverview) *session.Session { + detail := &session.Session{} + applyOverviewMetadata(overview, detail) + return detail + } + + ginkgo.It("projects the changed files stored by the monitor", func() { + Expect(project(overview()).Files).To(Equal(session.ChangedFiles{ + Read: []string{"todos/plans.go"}, + Written: []string{"todos/outcome.go", "todos/provider.go"}, + })) + }) + + ginkgo.It("projects todos, which no reader previously declared", func() { + Expect(project(overview()).Todos).To(Equal( + []tools.TodoItem{{Text: "ship it", Status: "pending"}})) + }) + + ginkgo.It("projects the git state the row already carries", func() { + Expect(project(overview()).Git).To(Equal(session.GitState{ + Branch: "feat/todo-prompt-runs", Commit: "0c6f277", Worktree: "/tmp/wt", + })) + }) + + ginkgo.It("projects the context window from the overview's own columns", func() { + row := overview() + row.ContextTokens, row.ContextWindowTokens, row.ContextFreePercent = &contextTokens, &windowTokens, &freePercent + + Expect(project(row).Context).To(Equal(&session.Context{ + UsedTokens: 48_000, WindowTokens: 200_000, FreePercent: 76, + })) + }) + + ginkgo.It("projects the stored plan when the branch found none", func() { + Expect(project(overview()).Plan).To(Equal(&session.Plan{Path: "/plans/stored.md", Slug: "stored"})) + }) + + ginkgo.DescribeTable("leaves a value the branch already resolved alone", + func(seed *session.Session, assert func(*session.Session)) { + applyOverviewMetadata(overview(), seed) + assert(seed) + }, + ginkgo.Entry("transcript-derived files", + &session.Session{Files: session.ChangedFiles{Written: []string{"from/transcript.go"}}}, + func(s *session.Session) { + Expect(s.Files.Written).To(Equal([]string{"from/transcript.go"})) + Expect(s.Files.Read).To(BeEmpty()) + }), + ginkgo.Entry("transcript-derived git", + &session.Session{Git: session.GitState{Branch: "main"}}, + func(s *session.Session) { Expect(s.Git.Branch).To(Equal("main")) }), + ginkgo.Entry("model resolved from the overview column", + &session.Session{Model: "claude-opus-5"}, + func(s *session.Session) { Expect(s.Model).To(Equal("claude-opus-5")) }), + ginkgo.Entry("a plan the branch already recovered", + &session.Session{Plan: &session.Plan{Slug: "from-transcript"}}, + func(s *session.Session) { Expect(s.Plan.Slug).To(Equal("from-transcript")) }), + ) + + ginkgo.It("falls back to the stored model and provider when the columns are empty", func() { + detail := project(overview()) + + Expect(detail.Model).To(Equal("gpt-5-codex")) + Expect(detail.Provider).To(Equal("openai")) + }) + + ginkgo.It("leaves approvals to the turn-request rows, not the transcript count", func() { + // The stored copy counts every operational tool use as approved, which + // reads as "200 approved" on a 200-tool-call session. applyRequestState + // derives the real figure from captain_turn_requests. + row := overview() + row.Metadata = json.RawMessage(`{"approvals": {"approved": 200, "denied": 0}}`) + + Expect(project(row).Approvals).To(Equal(session.ApprovalStats{})) + }) + + ginkgo.DescribeTable("survives a row with nothing to project", + func(mutate func(*database.SessionOverview)) { + row := database.SessionOverview{ID: uuid.New()} + mutate(&row) + + detail := project(row) + + Expect(detail.Files).To(Equal(session.ChangedFiles{})) + Expect(detail.Git).To(Equal(session.GitState{})) + Expect(detail.Context).To(BeNil()) + Expect(detail.Plan).To(BeNil()) + }, + ginkgo.Entry("no metadata or git at all", func(*database.SessionOverview) {}), + ginkgo.Entry("empty json objects", func(r *database.SessionOverview) { + r.Metadata, r.Git = json.RawMessage(`{}`), json.RawMessage(`{}`) + }), + ginkgo.Entry("malformed json", func(r *database.SessionOverview) { + r.Metadata, r.Git = json.RawMessage(`{"files":`), json.RawMessage(`nope`) + }), + ) +}) diff --git a/pkg/aichat/session_projection_ginkgo_test.go b/pkg/aichat/session_projection_ginkgo_test.go new file mode 100644 index 00000000..2ae6b8c2 --- /dev/null +++ b/pkg/aichat/session_projection_ginkgo_test.go @@ -0,0 +1,102 @@ +package aichat + +import ( + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" + "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = ginkgo.Describe("projectSessionAgents", func() { + rootID := uuid.MustParse("11111111-1111-4111-8111-111111111111") + childID := uuid.MustParse("22222222-2222-4222-8222-222222222222") + grandchildID := uuid.MustParse("33333333-3333-4333-8333-333333333333") + + explorer := "Explore" + rows := []database.SessionAgent{ + {SessionID: rootID, IsRoot: true, InputTokens: 100, OutputTokens: 10, TotalTokens: 110, CostUSD: 0.5}, + {SessionID: childID, ParentSessionID: &rootID, AgentType: &explorer, InputTokens: 40}, + {SessionID: grandchildID, ParentSessionID: &childID}, + } + + ginkgo.It("returns nothing for a session with no agent rows", func() { + root, agents := projectSessionAgents(nil) + + Expect(root).To(BeNil()) + Expect(agents).To(BeEmpty()) + }) + + ginkgo.It("nests each agent under its parent", func() { + root, agents := projectSessionAgents(rows) + + Expect(agents).To(HaveLen(3)) + Expect(root.ID).To(Equal(rootID.String())) + Expect(root.Children).To(HaveLen(1)) + Expect(root.Children[0].ID).To(Equal(childID.String())) + Expect(root.Children[0].Type).To(Equal("Explore")) + Expect(root.Children[0].Children[0].ID).To(Equal(grandchildID.String())) + }) + + ginkgo.It("carries each agent's own usage and cost", func() { + root, _ := projectSessionAgents(rows) + + Expect(root.Usage.InputTokens).To(Equal(100)) + Expect(root.Cost.TotalTokens).To(Equal(110)) + Expect(root.Cost.ProviderCostUSD).To(Equal(0.5)) + }) + + ginkgo.It("keeps an orphan in the flat index rather than dropping it", func() { + missingParent := uuid.MustParse("44444444-4444-4444-8444-444444444444") + + _, agents := projectSessionAgents([]database.SessionAgent{ + {SessionID: childID, ParentSessionID: &missingParent}, + }) + + Expect(agents).To(HaveLen(1)) + Expect(agents[0].ParentID).To(Equal(missingParent.String())) + }) + + ginkgo.It("does not make a self-parented row its own child", func() { + root, _ := projectSessionAgents([]database.SessionAgent{ + {SessionID: rootID, ParentSessionID: &rootID, IsRoot: true}, + }) + + Expect(root.Children).To(BeEmpty()) + }) +}) + +var _ = ginkgo.Describe("planFromNative", func() { + approved := database.Plan{ + Path: "/plans/approved.md", Slug: "approved", + ApprovedRevision: &database.PlanRevision{PlanMarkdown: "# approved"}, + LatestRevision: &database.PlanRevision{PlanMarkdown: "# newer draft"}, + } + draft := database.Plan{ + Path: "/plans/draft.md", Slug: "draft", + LatestRevision: &database.PlanRevision{PlanMarkdown: "# draft"}, + } + + ginkgo.It("has no plan when the session never persisted one", func() { + Expect(planFromNative(nil)).To(BeNil()) + }) + + ginkgo.It("prefers the approved revision over the newer draft on the same plan", func() { + Expect(planFromNative([]database.Plan{approved})).To(Equal(&session.Plan{ + Path: "/plans/approved.md", Slug: "approved", Content: "# approved", Explicit: true, + })) + }) + + ginkgo.It("prefers an approved plan over a newer unapproved one", func() { + // ListPlans orders newest first, so the draft is the more recent row. + Expect(planFromNative([]database.Plan{draft, approved}).Slug).To(Equal("approved")) + }) + + ginkgo.It("falls back to the latest revision when nothing was approved", func() { + Expect(planFromNative([]database.Plan{draft}).Content).To(Equal("# draft")) + }) + + ginkgo.It("ignores a plan row that has no revision at all", func() { + Expect(planFromNative([]database.Plan{{Path: "/plans/empty.md"}})).To(BeNil()) + }) +}) diff --git a/pkg/aichat/session_thread_files_ginkgo_test.go b/pkg/aichat/session_thread_files_ginkgo_test.go new file mode 100644 index 00000000..c69979a4 --- /dev/null +++ b/pkg/aichat/session_thread_files_ginkgo_test.go @@ -0,0 +1,75 @@ +package aichat + +import ( + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" + "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/segmentio/encoding/json" +) + +var _ = ginkgo.Describe("threadFiles", func() { + parentID := uuid.MustParse("2d33df99-654b-5b25-9d3e-b4d3a31e7cb5") + childID := uuid.MustParse("860305d7-e8cd-41b8-b5ee-332dc74d4a41") + grandchildID := uuid.MustParse("55555555-5555-4555-8555-555555555555") + + filesRow := func(id uuid.UUID, parent *uuid.UUID, written ...string) database.SessionOverview { + row := database.SessionOverview{ID: id, ParentSessionID: parent} + blob, err := json.Marshal(map[string]any{ + "files": session.ChangedFiles{Written: written}, + }) + Expect(err).NotTo(HaveOccurred()) + row.Metadata = blob + return row + } + + thread := []database.SessionOverview{ + {ID: parentID}, + filesRow(childID, &parentID, "todos/outcome.go", "todos/plans.go"), + filesRow(grandchildID, &childID, "todos/provider.go", "todos/plans.go"), + } + + ginkgo.It("gives a parent that edited nothing its sub-agents' files", func() { + // The reported bug: the thread's parent row carries no files metadata, so + // selecting it in the hierarchy showed an empty Files tab. + Expect(threadFiles(parentID, session.ChangedFiles{}, thread).Written).To(Equal( + []string{"todos/outcome.go", "todos/plans.go", "todos/provider.go"})) + }) + + ginkgo.It("includes descendants more than one level down", func() { + Expect(threadFiles(parentID, session.ChangedFiles{}, thread).Written).To( + ContainElement("todos/provider.go")) + }) + + ginkgo.It("merges the parent's own edits with its descendants'", func() { + own := session.ChangedFiles{Written: []string{"cmd/gavel/todos_plan.go"}} + + Expect(threadFiles(parentID, own, thread).Written).To(Equal([]string{ + "cmd/gavel/todos_plan.go", "todos/outcome.go", "todos/plans.go", "todos/provider.go", + })) + }) + + ginkgo.It("leaves a leaf reporting only its own files", func() { + own := session.ChangedFiles{Written: []string{"todos/provider.go", "todos/plans.go"}} + + // Unsorted and unmerged: a leaf's set is returned untouched so the + // hierarchy still shows which sub-agent touched what. + Expect(threadFiles(grandchildID, own, thread)).To(Equal(own)) + }) + + ginkgo.It("does not roll a sibling's files onto a leaf", func() { + Expect(threadFiles(grandchildID, session.ChangedFiles{}, thread).Written).To(BeEmpty()) + }) + + ginkgo.It("returns a mid-tree node its own subtree, not the whole thread", func() { + Expect(threadFiles(childID, session.ChangedFiles{Written: []string{"todos/outcome.go"}}, thread).Written).To( + Equal([]string{"todos/outcome.go", "todos/plans.go", "todos/provider.go"})) + }) + + ginkgo.It("returns the session's own files when it is the only row", func() { + own := session.ChangedFiles{Read: []string{"a.go"}} + + Expect(threadFiles(parentID, own, []database.SessionOverview{{ID: parentID}})).To(Equal(own)) + }) +}) diff --git a/pkg/aichat/session_title.go b/pkg/aichat/session_title.go new file mode 100644 index 00000000..995b9608 --- /dev/null +++ b/pkg/aichat/session_title.go @@ -0,0 +1,167 @@ +package aichat + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/session" +) + +// sessionTitleInstruction rides the system prompt whenever the naming tool is +// exposed. Without it the model almost never volunteers a title, and the +// conversation falls back to its opening message. +const sessionTitleInstruction = "Call " + session.TitleToolName + + " once, as early as you can, with a short title (at most eight words) describing what this conversation is about." + +const sessionTitleInput = "aiTitle" + +// appendSessionTitleInstruction adds the naming instruction to whichever system +// prompt the request carries. The two request modes are mutually exclusive, so +// message-mode specs must stay message-mode: agent backends prompt through +// Prompt, everything else through a leading system message. +func appendSessionTitleInstruction(spec *api.Spec) { + if len(spec.Messages) == 0 { + spec.Prompt.AppendSystem = strings.TrimSpace(spec.Prompt.AppendSystem + "\n\n" + sessionTitleInstruction) + return + } + if spec.Messages[0].Role == api.RoleSystem { + parts := spec.Messages[0].Parts + if len(parts) > 0 && parts[len(parts)-1].Type == api.PartText { + parts[len(parts)-1].Text += "\n\n" + sessionTitleInstruction + return + } + } + spec.Messages = append([]api.Message{{ + Role: api.RoleSystem, Parts: []api.Part{{Type: api.PartText, Text: sessionTitleInstruction}}, + }}, spec.Messages...) +} + +func normalizeTitle(update TitleUpdate) (string, error) { + title := strings.Join(strings.Fields(update.Title), " ") + if title == "" { + return "", fmt.Errorf("chat thread title cannot be empty") + } + if titleRank(update.Source) == 0 { + return "", fmt.Errorf("unknown chat thread title source %q", update.Source) + } + return title, nil +} + +func titleRank(source TitleSource) int { + switch source { + case TitleSourceDerived: + return 1 + case TitleSourceAI: + return 2 + case TitleSourceUser: + return 3 + default: + return 0 + } +} + +// titleWins applies the naming precedence: a person's title is final, the +// agent's replaces one inferred from the opening message, and an inferred title +// only ever fills a blank. +func titleWins(current string, stored, incoming TitleSource) bool { + if strings.TrimSpace(current) == "" { + return true + } + return titleRank(incoming) > titleRank(stored) +} + +// derivedTitle names a conversation after the first thing the user asked. +func derivedTitle(messages []UIMessage) string { + for _, message := range messages { + if !strings.EqualFold(message.Role, string(api.RoleUser)) { + continue + } + for _, part := range message.Parts { + if part.Type != "text" { + continue + } + if title := session.DeriveTitle(part.Text); title != "" { + return title + } + } + } + return "" +} + +// agentTitle reads the title an agent gave itself, from either the tool Captain +// exposes or an equivalent call the backend emits on its own. +func agentTitle(message UIMessage) string { + title := "" + for _, part := range message.Parts { + if part.ToolName != session.TitleToolName || len(part.Input) == 0 { + continue + } + input := struct { + AITitle string `json:"aiTitle"` + }{} + if err := json.Unmarshal(part.Input, &input); err != nil { + serviceLog.Warnf("decode %s input: %v", session.TitleToolName, err) + continue + } + if strings.TrimSpace(input.AITitle) != "" { + title = input.AITitle + } + } + return title +} + +// setThreadTitle applies a title, treating a losing or malformed one as +// nothing to do: naming a conversation must never fail its turn. +func (s *Service) setThreadTitle(ctx context.Context, threadID string, update TitleUpdate) { + if threadID == "" || strings.TrimSpace(update.Title) == "" { + return + } + store, err := s.threads(ctx) + if err != nil { + serviceLog.Warnf("title chat thread %q: %v", threadID, err) + return + } + if err := store.SetTitle(ctx, threadID, update); err != nil { + serviceLog.Warnf("title chat thread %q: %v", threadID, err) + } +} + +// sessionTitleTool lets the model name the conversation it is in. It is bound +// to one thread and injected per request, so it never appears in the catalog +// the user manages tool preferences against. +func (s *Service) sessionTitleTool(threadID string) api.ToolDefinition { + readOnly := true + return api.ToolDefinition{ + Name: session.TitleToolName, + Description: "Name the current conversation. " + sessionTitleInstruction, + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + sessionTitleInput: map[string]any{ + "type": "string", + "description": "A short title for this conversation, at most eight words.", + }, + }, + "required": []any{sessionTitleInput}, + }, + ReadOnlyHint: &readOnly, + DefaultPermission: api.ToolModeOn, + Handler: func(ctx context.Context, input map[string]any) (any, error) { + title, _ := input[sessionTitleInput].(string) + if strings.TrimSpace(title) == "" { + return nil, fmt.Errorf("%s requires a non-empty %s", session.TitleToolName, sessionTitleInput) + } + store, err := s.threads(ctx) + if err != nil { + return nil, err + } + if err := store.SetTitle(ctx, threadID, TitleUpdate{Title: title, Source: TitleSourceAI}); err != nil { + return nil, err + } + return map[string]any{"title": title}, nil + }, + } +} diff --git a/pkg/aichat/session_title_ginkgo_test.go b/pkg/aichat/session_title_ginkgo_test.go new file mode 100644 index 00000000..7f87620f --- /dev/null +++ b/pkg/aichat/session_title_ginkgo_test.go @@ -0,0 +1,146 @@ +package aichat_test + +import ( + "context" + "encoding/json" + "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" + "github.com/flanksource/captain/pkg/session" +) + +var _ = Describe("chat session titles", func() { + const opener = "Update the category and tax_category dimensions on all accounts missing them" + + newService := func(store aichat.ThreadStore, events []api.Event, tools []api.ToolDefinition) (*aichat.Service, *fakeStreamingProvider) { + provider := &fakeStreamingProvider{events: append(events, api.Event{Kind: api.EventResult, Success: true})} + options := aichat.ServiceOptions{ + Resolver: &fakeResolver{provider: provider}, Threads: aichat.FixedThreadStore(store), + } + if len(tools) > 0 { + options.Tools = aichat.StaticToolProvider(tools) + } + return aichat.NewService(options), provider + } + + submit := func(service *aichat.Service, threadID, text string) *httptest.ResponseRecorder { + response := httptest.NewRecorder() + service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ + ID: threadID, ThreadID: threadID, Trigger: "submit-message", Model: "openai/test-model", + Messages: []aichat.UIMessage{{ + ID: "message-" + threadID, Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: text}}, + }}, + })) + return response + } + + It("names an unnamed thread after the message that opened it", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "") + Expect(err).NotTo(HaveOccurred()) + service, _ := newService(store, nil, nil) + + Expect(submit(service, thread.ID, opener).Code).To(Equal(http.StatusOK)) + + named, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(named.Title).To(Equal(opener)) + }) + + It("prefers the title the agent gives itself over the one inferred from the opener", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "") + Expect(err).NotTo(HaveOccurred()) + service, _ := newService(store, []api.Event{{ + Kind: api.EventToolUse, Tool: session.TitleToolName, ToolCallID: "call-title", + Input: map[string]any{"aiTitle": "Account dimension backfill"}, + }, { + Kind: api.EventToolResult, Tool: session.TitleToolName, ToolCallID: "call-title", + Text: `{"title":"Account dimension backfill"}`, Success: true, + }}, nil) + + Expect(submit(service, thread.ID, opener).Code).To(Equal(http.StatusOK)) + + named, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(named.Title).To(Equal("Account dimension backfill")) + }) + + It("offers the naming tool and its instruction to a chat that already carries tools", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "") + Expect(err).NotTo(HaveOccurred()) + service, provider := newService(store, nil, []api.ToolDefinition{{ + Name: "invoice_get", DefaultPermission: api.ToolModeOn, + Handler: func(context.Context, map[string]any) (any, error) { return nil, nil }, + }}) + + Expect(submit(service, thread.ID, opener).Code).To(Equal(http.StatusOK)) + + Expect(provider.specs).To(HaveLen(1)) + Expect(provider.specs[0].Messages[0].Role).To(Equal(api.RoleSystem)) + Expect(provider.specs[0].Messages[0].Parts[0].Text).To(ContainSubstring(session.TitleToolName)) + Expect(provider.specs[0].Prompt.AppendSystem).To(BeEmpty(), "message-mode requests must stay message-mode") + }) + + It("leaves a toolless chat without caller tools", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "") + Expect(err).NotTo(HaveOccurred()) + service, provider := newService(store, nil, nil) + + Expect(submit(service, thread.ID, opener).Code).To(Equal(http.StatusOK)) + + Expect(provider.specs).To(HaveLen(1)) + Expect(provider.specs[0].Prompt.AppendSystem).To(BeEmpty()) + }) + + It("renames a thread through the sessions endpoint and keeps that name", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "") + Expect(err).NotTo(HaveOccurred()) + service, _ := newService(store, []api.Event{{ + Kind: api.EventToolUse, Tool: session.TitleToolName, ToolCallID: "call-title", + Input: map[string]any{"aiTitle": "Account dimension backfill"}, + }, { + Kind: api.EventToolResult, Tool: session.TitleToolName, ToolCallID: "call-title", + Text: `{"title":"Account dimension backfill"}`, Success: true, + }}, nil) + + rename := httptest.NewRecorder() + service.Handler().ServeHTTP(rename, requestJSON(http.MethodPatch, "/api/chat/sessions/"+thread.ID, + map[string]string{"title": " FY25 dimension cleanup "})) + Expect(rename.Code).To(Equal(http.StatusOK)) + var renamed aichat.Thread + Expect(json.Unmarshal(rename.Body.Bytes(), &renamed)).To(Succeed()) + Expect(renamed.Title).To(Equal("FY25 dimension cleanup")) + + Expect(submit(service, thread.ID, opener).Code).To(Equal(http.StatusOK)) + + stored, err := store.Get(context.Background(), thread.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(stored.Title).To(Equal("FY25 dimension cleanup")) + }) + + It("rejects a blank rename and an unknown thread", func() { + store := aichat.NewMemoryThreadStore() + thread, err := store.Create(context.Background(), "") + Expect(err).NotTo(HaveOccurred()) + service, _ := newService(store, nil, nil) + + blank := httptest.NewRecorder() + service.Handler().ServeHTTP(blank, requestJSON(http.MethodPatch, "/api/chat/sessions/"+thread.ID, + map[string]string{"title": " "})) + Expect(blank.Code).To(Equal(http.StatusBadRequest)) + + missing := httptest.NewRecorder() + service.Handler().ServeHTTP(missing, requestJSON(http.MethodPatch, "/api/chat/sessions/does-not-exist", + map[string]string{"title": "anything"})) + Expect(missing.Code).To(Equal(http.StatusNotFound)) + }) +}) diff --git a/pkg/aichat/session_title_test.go b/pkg/aichat/session_title_test.go new file mode 100644 index 00000000..a718ed1d --- /dev/null +++ b/pkg/aichat/session_title_test.go @@ -0,0 +1,64 @@ +package aichat + +import ( + "encoding/json" + "testing" + + "github.com/flanksource/captain/pkg/session" +) + +func TestAgentTitleReadsTheLastNamingCall(t *testing.T) { + input := func(title string) json.RawMessage { + payload, err := json.Marshal(map[string]any{"aiTitle": title}) + if err != nil { + t.Fatalf("marshal tool input: %v", err) + } + return payload + } + message := UIMessage{Role: "assistant", Parts: []UIPart{ + {Type: "text", Text: "working"}, + {Type: "dynamic-tool", ToolName: "invoice_get", ToolCallID: "call-1", Input: input("not a title")}, + {Type: "dynamic-tool", ToolName: session.TitleToolName, ToolCallID: "call-2", Input: input("First guess")}, + {Type: "dynamic-tool", ToolName: session.TitleToolName, ToolCallID: "call-3", Input: input("Account dimension backfill")}, + }} + if got := agentTitle(message); got != "Account dimension backfill" { + t.Fatalf("agentTitle = %q, want the last %s call's title", got, session.TitleToolName) + } + if got := agentTitle(UIMessage{Role: "assistant", Parts: []UIPart{{Type: "text", Text: "no tools here"}}}); got != "" { + t.Fatalf("agentTitle = %q, want empty when the agent never named the thread", got) + } +} + +func TestDerivedTitleUsesTheFirstUserText(t *testing.T) { + messages := []UIMessage{ + {Role: "assistant", Parts: []UIPart{{Type: "text", Text: "How can I help?"}}}, + {Role: "user", Parts: []UIPart{ + {Type: "file", Filename: "trial-balance.csv"}, + {Type: "text", Text: " Reconcile\n the trial balance "}, + }}, + {Role: "user", Parts: []UIPart{{Type: "text", Text: "and then post it"}}}, + } + if got := derivedTitle(messages); got != "Reconcile the trial balance" { + t.Fatalf("derivedTitle = %q, want the first user message collapsed", got) + } +} + +func TestTitleWinsFollowsNamingPrecedence(t *testing.T) { + for _, tc := range []struct { + name string + current string + stored TitleSource + incoming TitleSource + want bool + }{ + {name: "anything names a blank thread", current: "", stored: "", incoming: TitleSourceDerived, want: true}, + {name: "derived does not overwrite", current: "Named", stored: TitleSourceDerived, incoming: TitleSourceDerived, want: false}, + {name: "agent replaces derived", current: "Named", stored: TitleSourceDerived, incoming: TitleSourceAI, want: true}, + {name: "agent does not replace a person", current: "Named", stored: TitleSourceUser, incoming: TitleSourceAI, want: false}, + {name: "a person always wins", current: "Named", stored: TitleSourceAI, incoming: TitleSourceUser, want: true}, + } { + if got := titleWins(tc.current, tc.stored, tc.incoming); got != tc.want { + t.Errorf("%s: titleWins = %v, want %v", tc.name, got, tc.want) + } + } +} diff --git a/pkg/aichat/sse.go b/pkg/aichat/sse.go index 5f9c29f8..7ac7faca 100644 --- a/pkg/aichat/sse.go +++ b/pkg/aichat/sse.go @@ -36,14 +36,43 @@ type UsageMetadata struct { TotalTokens int `json:"totalTokens"` } +// CostBreakdownMetadata is one turn's cost split across the disjoint token +// buckets, priced by ai.PriceUsage. Field names match the frontend's +// ChatCostBreakdown so the UI renders per-bucket rows instead of "-". +type CostBreakdownMetadata struct { + Model string `json:"model,omitempty"` + InputUSD float64 `json:"inputUsd"` + OutputUSD float64 `json:"outputUsd"` + ReasoningUSD float64 `json:"reasoningUsd"` + CacheReadUSD float64 `json:"cacheReadUsd"` + // CacheWriteUSD is structurally zero on the API backends: genkit's usage + // type carries no cache-write field (pkg/ai/provider/genkit/mapping.go). + CacheWriteUSD float64 `json:"cacheWriteUsd"` + TotalUSD float64 `json:"totalUsd"` +} + // 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"` + ProviderSessionID string `json:"providerSessionId,omitempty"` + Model string `json:"model,omitempty"` + Usage *UsageMetadata `json:"usage,omitempty"` + Cost float64 `json:"cost,omitempty"` + CostBreakdown *CostBreakdownMetadata `json:"costBreakdown,omitempty"` + // ThreadCostUSD is the conversation's cumulative spend. Cost above is this + // turn alone; a UI showing a running total must read this field, since the + // two differ by the number of turns taken. + ThreadCostUSD float64 `json:"threadCostUsd,omitempty"` + ContextTokens int `json:"contextTokens,omitempty"` + Success *bool `json:"success,omitempty"` + Interrupted bool `json:"interrupted,omitempty"` +} + +// TurnCosts is written by the persistence layer as a turn completes and read by +// the event stream when it writes the finish part. The event channel between +// them orders the write before the read. +type TurnCosts struct { + Breakdown *CostBreakdownMetadata + ThreadCostUSD float64 } // SSEWriter writes AI SDK v6 chunks using Server-Sent Events framing. diff --git a/pkg/aichat/stream_ginkgo_test.go b/pkg/aichat/stream_ginkgo_test.go index 7e81a597..3ae863db 100644 --- a/pkg/aichat/stream_ginkgo_test.go +++ b/pkg/aichat/stream_ginkgo_test.go @@ -29,6 +29,10 @@ func (w *nonFlushWriter) Write([]byte) (int, error) { return 0, nil } func (w *nonFlushWriter) WriteHeader(statusCode int) {} func recordEvents(events ...api.Event) (*flushRecorder, error) { + return recordEventsWithOptions(aichat.EventStreamOptions{}, events...) +} + +func recordEventsWithOptions(options aichat.EventStreamOptions, events ...api.Event) (*flushRecorder, error) { recorder := &flushRecorder{ResponseRecorder: httptest.NewRecorder()} writer, err := aichat.NewSSEWriter(recorder) if err != nil { @@ -39,7 +43,7 @@ func recordEvents(events ...api.Event) (*flushRecorder, error) { channel <- event } close(channel) - return recorder, aichat.WriteEventStream(writer, channel, aichat.EventStreamOptions{}) + return recorder, aichat.WriteEventStream(writer, channel, options) } func decodedDataLines(body string) []map[string]any { @@ -81,6 +85,12 @@ func pendingApprovalState(calls ...api.ToolApprovalRequest) *api.ToolApprovalSta } } +func approvalEvent(callID, tool string) api.Event { + return api.Event{ + Kind: api.EventPermission, ToolCallID: callID, Tool: tool, ApprovalID: "approval-" + callID, + } +} + var _ = Describe("AI SDK v6 event stream", func() { It("rejects a response writer that cannot stream", func() { _, err := aichat.NewSSEWriter(&nonFlushWriter{header: http.Header{}}) @@ -107,7 +117,7 @@ var _ = Describe("AI SDK v6 event stream", func() { 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"}, + approvalEvent("call-1", "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"}`)}, @@ -130,7 +140,7 @@ var _ = Describe("AI SDK v6 event stream", func() { HaveKeyWithValue("dynamic", true), )) Expect(parts[10]).To(SatisfyAll( - HaveKeyWithValue("approvalId", "call-1"), + HaveKeyWithValue("approvalId", "approval-call-1"), HaveKeyWithValue("toolCallId", "call-1"), )) Expect(parts[11]["output"]).To(Equal(map[string]any{"output": `{"status":"draft"}`})) @@ -142,30 +152,93 @@ var _ = Describe("AI SDK v6 event stream", func() { "inputTokens": 100.0, "outputTokens": 40.0, "reasoningTokens": 10.0, "cacheReadTokens": 5.0, "cacheWriteTokens": 0.0, "totalTokens": 155.0, }, - "cost": 0.0125, - "contextTokens": 100.0, + "cost": 0.0125, + // Context occupancy is the whole prompt: input plus the cached + // prefix, not input alone. + "contextTokens": 105.0, "success": true, })) }) + It("rides the priced breakdown and the thread's cumulative cost on the finish part", func() { + // Without these the UI has nothing to render: every per-bucket Cost row + // falls back to "-", and "Thread total" silently degrades to this one + // turn's cost, which understated a real 9-call thread by 10x. + costs := &aichat.TurnCosts{ + Breakdown: &aichat.CostBreakdownMetadata{ + Model: "claude-sonnet", InputUSD: 0.0003, OutputUSD: 0.006, + ReasoningUSD: 0.0015, CacheReadUSD: 0.0000015, TotalUSD: 0.0125, + }, + ThreadCostUSD: 7.358646, + } + recorder, err := recordEventsWithOptions( + aichat.EventStreamOptions{Costs: costs}, + api.Event{Kind: api.EventText, Text: "done"}, + api.Event{ + Kind: api.EventResult, SessionID: "session-1", Model: "claude-sonnet", + Usage: &api.Usage{InputTokens: 100, OutputTokens: 40}, CostUSD: 0.0125, Success: true, + }, + ) + Expect(err).NotTo(HaveOccurred()) + + parts := decodedDataLines(recorder.Body.String()) + metadata, ok := parts[len(parts)-1]["messageMetadata"].(map[string]any) + Expect(ok).To(BeTrue()) + Expect(metadata).To(HaveKeyWithValue("threadCostUsd", 7.358646)) + Expect(metadata).To(HaveKeyWithValue("cost", 0.0125), + "the per-turn cost stays alongside the cumulative total, not replaced by it") + // cacheWriteUsd is emitted as an explicit 0 rather than omitted: on the + // API backends genkit carries no cache-write tokens at all, and a + // missing key would render as "-" (unknown) instead of "$0". + Expect(metadata["costBreakdown"]).To(Equal(map[string]any{ + "model": "claude-sonnet", "inputUsd": 0.0003, "outputUsd": 0.006, + "reasoningUsd": 0.0015, "cacheReadUsd": 0.0000015, "cacheWriteUsd": 0.0, + "totalUsd": 0.0125, + })) + }) + 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.EventToolUse, ToolCallID: "call-1", Tool: "accounts_edit"}, 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", + "tool-input-available", "tool-output-error", "error", "finish-step", "finish", })) - Expect(decodedDataLines(recorder.Body.String())[5]).To(HaveKeyWithValue("errorText", "provider disconnected")) + parts := decodedDataLines(recorder.Body.String()) + Expect(parts[6]).To(SatisfyAll( + HaveKeyWithValue("toolCallId", "call-1"), + HaveKeyWithValue("errorText", "provider disconnected"), + )) + Expect(parts[7]).To(HaveKeyWithValue("errorText", "provider disconnected")) Expect(recorder.Body.String()).To(HaveSuffix("data: [DONE]\n\n")) }) + It("finishes an interrupted turn without rendering a provider error", func() { + recorder, err := recordEvents( + api.Event{Kind: api.EventText, Text: "partial"}, + api.Event{Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "accounts_edit"}, + api.Event{Kind: api.EventInterrupted, Reason: "user"}, + ) + Expect(err).NotTo(HaveOccurred()) + parts := decodedDataLines(recorder.Body.String()) + Expect(partTypes(parts)).To(Equal([]string{ + "start", "start-step", "text-start", "text-delta", "text-end", + "tool-input-available", "tool-output-error", "data-result", "finish-step", "finish", + })) + Expect(parts[6]).To(HaveKeyWithValue("errorText", "user")) + Expect(parts[7]["data"]).To(Equal(map[string]any{"success": false, "interrupted": true})) + Expect(parts[9]["messageMetadata"]).To(HaveKeyWithValue("interrupted", true)) + Expect(recorder.Body.String()).NotTo(ContainSubstring(`"type":"error"`)) + }) + 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"}, + approvalEvent("call-1", "invoice_update"), ) Expect(err).NotTo(HaveOccurred()) Expect(partTypes(decodedDataLines(recorder.Body.String()))).To(Equal([]string{ @@ -179,23 +252,23 @@ var _ = Describe("AI SDK v6 event stream", 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"}, + approvalEvent("call-1", "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", + "data-result", "finish-step", "finish", })) - Expect(parts[4]["data"]).To(HaveKeyWithValue("calls", HaveLen(1))) + Expect(parts[4]["data"]).To(Equal(map[string]any{"success": true, "waitingApproval": true})) }) 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"}, + approvalEvent("call-1", "invoice_update"), api.Event{Kind: api.EventResult, Success: true, ToolApproval: state}, ) Expect(err).To(MatchError(message)) @@ -222,9 +295,9 @@ var _ = Describe("AI SDK v6 event stream", func() { }) _, 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"}, + approvalEvent("call-1", "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"}, + approvalEvent("call-2", "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`)) @@ -244,12 +317,12 @@ var _ = Describe("AI SDK v6 event stream", func() { {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 permission", []api.Event{approvalEvent("call-1", "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"}, + approvalEvent("call-1", "invoice_update"), + approvalEvent("call-1", "invoice_update"), }, `duplicate permission for tool call "call-1"`), Entry("mismatched tool name", []api.Event{ {Kind: api.EventToolUse, ToolCallID: "call-1", Tool: "invoice_get"}, @@ -257,7 +330,7 @@ var _ = Describe("AI SDK v6 event stream", func() { }, `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"}, + approvalEvent("call-1", "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/thread_costs.go b/pkg/aichat/thread_costs.go new file mode 100644 index 00000000..7312f6a4 --- /dev/null +++ b/pkg/aichat/thread_costs.go @@ -0,0 +1,114 @@ +package aichat + +import ( + "context" + "fmt" + "sort" + "strings" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" +) + +// ThreadCosts is the cost breakdown for one chat thread, scoped to the root +// session plus every subagent session beneath it. +// +// ByModel is the meaningful axis today: a chat thread can switch backends +// mid-conversation (an API turn followed by an agent turn bills at very +// different rates), and only this split makes that visible. ByAgent carries the +// sub-session dimension, which collapses to a single root row until chat +// threads start recording subagents as child sessions. +type ThreadCosts struct { + ThreadID string `json:"threadId"` + TotalCostUSD float64 `json:"totalCostUsd"` + ByModel []database.SessionCost `json:"byModel"` + ByAgent []database.SessionAgent `json:"byAgent"` +} + +// ThreadCostReader is implemented by thread stores that can report a thread's +// cost breakdown. The in-memory store cannot, so the route reports that rather +// than serving zeros that look like a free conversation. +type ThreadCostReader interface { + GetThreadCosts(context.Context, string) (*ThreadCosts, error) +} + +func (s *DatabaseThreadStore) GetThreadCosts(ctx context.Context, id string) (*ThreadCosts, error) { + rootID, err := uuid.Parse(strings.TrimSpace(id)) + if err != nil { + return nil, fmt.Errorf("captain chat session ID %q is not a UUID: %w", id, err) + } + costs, err := s.db.ListThreadCosts(ctx, rootID) + if err != nil { + return nil, err + } + agents, err := s.db.ListThreadAgents(ctx, rootID) + if err != nil { + return nil, err + } + return &ThreadCosts{ + ThreadID: rootID.String(), + TotalCostUSD: threadTotalCostUSD(agents), + ByModel: costs, + ByAgent: agents, + }, nil +} + +// threadTotalCostUSD sums every session in the thread rather than reading the +// root's own total: captain_session_overview scopes cost to `t.session_id = +// s.id`, so a root row's figure excludes its subagents' spend. +func threadTotalCostUSD(agents []database.SessionAgent) float64 { + var total float64 + for i := range agents { + total += agents[i].CostUSD + } + return total +} + +// applyThreadCosts replaces a session aggregate's root-scoped usage and cost +// with thread-wide totals, and fills the per-model breakdown that the DB path +// otherwise leaves empty. +func applyThreadCosts(aggregate *session.Session, rows []database.SessionCost) { + if len(rows) == 0 { + return + } + costs := make(api.Costs, len(rows)) + for i := range rows { + costs[i] = api.Cost{ + Model: rows[i].Model, + InputTokens: int(rows[i].InputTokens), + OutputTokens: int(rows[i].OutputTokens), + ReasoningTokens: int(rows[i].ReasoningTokens), + CacheReadTokens: int(rows[i].CacheReadTokens), + CacheWriteTokens: int(rows[i].CacheWriteTokens), + TotalTokens: int(rows[i].TotalTokens), + InputCost: rows[i].InputCost, + OutputCost: rows[i].OutputCost, + ReasoningCost: rows[i].ReasoningCost, + CacheReadCost: rows[i].CacheReadCost, + CacheWriteCost: rows[i].CacheWriteCost, + // Only the providers' own reported share, not the view's TotalCost: + // TotalCost falls back to the list-priced buckets per call, so passing + // it here would make every reconstruction claim to be a billed figure. + // Cost.Total() resolves the two exactly as the view's CASE does. + ProviderCostUSD: rows[i].ProviderCostUSD, + } + } + total := costs.Sum() + aggregate.Cost = total + aggregate.Usage = api.Usage{ + InputTokens: total.InputTokens, OutputTokens: total.OutputTokens, + ReasoningTokens: total.ReasoningTokens, CacheReadTokens: total.CacheReadTokens, + CacheWriteTokens: total.CacheWriteTokens, + } + byModel := costs.ByModel() + aggregate.ToolCosts = make(api.Costs, 0, len(byModel)) + for model, cost := range byModel { + cost.Model = model + aggregate.ToolCosts = append(aggregate.ToolCosts, cost) + } + sort.Slice(aggregate.ToolCosts, func(i, j int) bool { + return aggregate.ToolCosts[i].Model < aggregate.ToolCosts[j].Model + }) +} diff --git a/pkg/aichat/thread_costs_ginkgo_test.go b/pkg/aichat/thread_costs_ginkgo_test.go new file mode 100644 index 00000000..7204ada1 --- /dev/null +++ b/pkg/aichat/thread_costs_ginkgo_test.go @@ -0,0 +1,61 @@ +package aichat + +import ( + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/session" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// captain_session_costs.total_cost resolves provider-reported against list-price +// per underlying call, so it is non-zero either way. Carrying it as the +// aggregate's ProviderCostUSD — which three separate projections did — makes a +// pure reconstruction indistinguishable from a billed figure, and every renderer +// that marks estimates then silently presents one as the other. +var _ = Describe("thread cost aggregation", func() { + const ( + inputCost = 0.75 + outputCost = 1.25 + cacheReadCost = 0.50 + bucketTotal = inputCost + outputCost + cacheReadCost + providerBilled = 3.10 + ) + + listPriced := database.SessionCost{ + Model: "claude-opus-5", InputTokens: 1000, OutputTokens: 200, CacheReadTokens: 5000, + InputCost: inputCost, OutputCost: outputCost, CacheReadCost: cacheReadCost, + TotalCost: bucketTotal, ProviderCostUSD: 0, + } + // A provider-reported call stores its buckets too; total_cost prefers the + // billed figure over their sum, and so must the aggregate. + providerReported := database.SessionCost{ + Model: "claude-opus-5", InputTokens: 1000, OutputTokens: 200, CacheReadTokens: 5000, + InputCost: inputCost, OutputCost: outputCost, CacheReadCost: cacheReadCost, + TotalCost: providerBilled, ProviderCostUSD: providerBilled, + } + + It("leaves a list-priced thread with no provider cost, so it renders as an estimate", func() { + aggregate := &session.Session{} + applyThreadCosts(aggregate, []database.SessionCost{listPriced}) + + Expect(aggregate.Cost.ProviderCostUSD).To(BeZero()) + Expect(aggregate.Cost.Total()).To(BeNumerically("~", bucketTotal, 1e-9)) + }) + + It("keeps a provider-reported thread's billed total in preference to its buckets", func() { + aggregate := &session.Session{} + applyThreadCosts(aggregate, []database.SessionCost{providerReported}) + + Expect(aggregate.Cost.ProviderCostUSD).To(BeNumerically("~", providerBilled, 1e-9)) + Expect(aggregate.Cost.Total()).To(BeNumerically("~", providerBilled, 1e-9)) + }) + + It("reports usage summed across the thread's models", func() { + aggregate := &session.Session{} + applyThreadCosts(aggregate, []database.SessionCost{listPriced, providerReported}) + + Expect(aggregate.Usage.InputTokens).To(Equal(2000)) + Expect(aggregate.Usage.OutputTokens).To(Equal(400)) + Expect(aggregate.Usage.CacheReadTokens).To(Equal(10000)) + }) +}) diff --git a/pkg/aichat/threads.go b/pkg/aichat/threads.go index 6830618b..f6c0873a 100644 --- a/pkg/aichat/threads.go +++ b/pkg/aichat/threads.go @@ -7,6 +7,9 @@ import ( "strings" "sync" "time" + + "github.com/flanksource/captain/pkg/session" + "github.com/google/uuid" ) type Thread struct { @@ -35,6 +38,22 @@ type TurnUsage struct { CostUSD float64 } +// TitleSource records who named a conversation. It decides whether a later +// writer may rename it: a person's title outranks the agent's, which outranks +// one inferred from the opening message. +type TitleSource string + +const ( + TitleSourceDerived TitleSource = "derived" + TitleSourceAI TitleSource = "ai" + TitleSourceUser TitleSource = "user" +) + +type TitleUpdate struct { + Title string + Source TitleSource +} + // ThreadStore is the persistence boundary for chat history, provider session // identity, and cumulative usage. Implementations must be concurrency-safe. type ThreadStore interface { @@ -45,26 +64,33 @@ type ThreadStore interface { ReplaceLastMessage(context.Context, string, UIMessage) error Delete(context.Context, string) error SetProviderSession(context.Context, string, string) error + SetTitle(context.Context, string, TitleUpdate) error AddUsage(context.Context, string, TurnUsage) (*Thread, error) } +type SessionReader interface { + GetSession(context.Context, string) (*session.Session, error) +} + type memoryThreadStore struct { - mu sync.Mutex - seq int - threads map[string]*Thread + mu sync.Mutex + threads map[string]*Thread + titleSources map[string]TitleSource } func NewMemoryThreadStore() ThreadStore { - return &memoryThreadStore{threads: map[string]*Thread{}} + return &memoryThreadStore{threads: map[string]*Thread{}, titleSources: map[string]TitleSource{}} } 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{}} + thread := &Thread{ID: uuid.NewString(), Title: title, CreatedAt: now, UpdatedAt: now, Messages: []UIMessage{}} s.threads[thread.ID] = thread + if strings.TrimSpace(title) != "" { + s.titleSources[thread.ID] = TitleSourceUser + } return cloneThread(thread), nil } @@ -133,11 +159,45 @@ func (s *memoryThreadStore) SetProviderSession(_ context.Context, id, sessionID if err != nil { return err } + sessionID = strings.TrimSpace(sessionID) + if sessionID == "" { + return fmt.Errorf("provider session ID cannot be empty") + } + if thread.ProviderSessionID != "" && thread.ProviderSessionID != sessionID { + return fmt.Errorf( + "provider session is already bound to %q, cannot replace it with %q", + thread.ProviderSessionID, + sessionID, + ) + } + if thread.ProviderSessionID == sessionID { + return nil + } thread.ProviderSessionID = sessionID thread.UpdatedAt = time.Now() return nil } +func (s *memoryThreadStore) SetTitle(_ context.Context, id string, update TitleUpdate) error { + s.mu.Lock() + defer s.mu.Unlock() + thread, err := s.thread(id) + if err != nil { + return err + } + title, err := normalizeTitle(update) + if err != nil { + return err + } + if !titleWins(thread.Title, s.titleSources[id], update.Source) { + return nil + } + thread.Title = title + s.titleSources[id] = update.Source + 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() @@ -180,5 +240,12 @@ func validateLastMessageReplacement(messages []UIMessage, replacement UIMessage) if !strings.EqualFold(messages[len(messages)-1].Role, "assistant") { return fmt.Errorf("last stored message must have assistant role") } + if messages[len(messages)-1].ID != "" && replacement.ID != messages[len(messages)-1].ID { + return fmt.Errorf( + "replacement message ID %q does not match stored message %q", + replacement.ID, + messages[len(messages)-1].ID, + ) + } return nil } diff --git a/pkg/aichat/threads_http.go b/pkg/aichat/threads_http.go index 9840cf1e..8914d766 100644 --- a/pkg/aichat/threads_http.go +++ b/pkg/aichat/threads_http.go @@ -1,29 +1,61 @@ package aichat import ( + "context" "encoding/json" + "errors" "fmt" "io" "net/http" + "strings" ) 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) + mux.HandleFunc("POST /api/chat/sessions", s.handleCreateThread) + mux.HandleFunc("GET /api/chat/sessions", s.handleListThreads) + mux.HandleFunc("GET /api/chat/sessions/{id}", s.handleGetThread) + mux.HandleFunc("PATCH /api/chat/sessions/{id}", s.handleRenameThread) + mux.HandleFunc("GET /api/chat/sessions/{id}/costs", s.handleThreadCosts) + mux.HandleFunc("DELETE /api/chat/sessions/{id}", s.handleDeleteThread) + mux.HandleFunc("POST /api/chat/sessions/{id}/approvals/{approvalID}", s.handleResolveToolApproval) + mux.HandleFunc("POST /api/chat/sessions/{id}/interrupt", s.handleInterrupt) } -func (s *Service) threadStore(w http.ResponseWriter) ThreadStore { +// threads resolves the thread store for one request. The store can differ per +// request when the application serves more than one database. +func (s *Service) threads(ctx context.Context) (ThreadStore, error) { if s.options.Threads == nil { - http.Error(w, "thread persistence is not configured", http.StatusNotImplemented) + return nil, errThreadsNotConfigured + } + store, err := s.options.Threads.ThreadStore(ctx) + if err != nil { + return nil, err + } + if store == nil { + return nil, errThreadsNotConfigured + } + return store, nil +} + +var errThreadsNotConfigured = errors.New("thread persistence is not configured") + +// threadStore resolves the request's thread store, writing the failure response +// itself and returning nil when it cannot. +func (s *Service) threadStore(w http.ResponseWriter, request *http.Request) ThreadStore { + store, err := s.threads(request.Context()) + if errors.Is(err, errThreadsNotConfigured) { + http.Error(w, err.Error(), http.StatusNotImplemented) + return nil + } + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) return nil } - return s.options.Threads + return store } func (s *Service) handleCreateThread(w http.ResponseWriter, request *http.Request) { - store := s.threadStore(w) + store := s.threadStore(w, request) if store == nil { return } @@ -36,9 +68,8 @@ func (s *Service) handleCreateThread(w http.ResponseWriter, request *http.Reques return } } - if body.Title == "" { - body.Title = "New conversation" - } + // An unnamed thread stays unnamed: it is named from its first message, or by + // the agent, or by a rename — never with a placeholder that reads like a title. thread, err := store.Create(request.Context(), body.Title) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) @@ -50,7 +81,7 @@ func (s *Service) handleCreateThread(w http.ResponseWriter, request *http.Reques } func (s *Service) handleListThreads(w http.ResponseWriter, request *http.Request) { - store := s.threadStore(w) + store := s.threadStore(w, request) if store == nil { return } @@ -65,10 +96,21 @@ func (s *Service) handleListThreads(w http.ResponseWriter, request *http.Request } func (s *Service) handleGetThread(w http.ResponseWriter, request *http.Request) { - store := s.threadStore(w) + store := s.threadStore(w, request) if store == nil { return } + if sessions, ok := store.(SessionReader); ok { + aggregate, err := sessions.GetSession(request.Context(), request.PathValue("id")) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + if err := writeJSON(w, http.StatusOK, aggregate); err != nil { + serviceLog.Errorf("write chat session %q: %v", request.PathValue("id"), err) + } + return + } thread, err := store.Get(request.Context(), request.PathValue("id")) if err != nil { http.Error(w, err.Error(), http.StatusNotFound) @@ -79,8 +121,59 @@ func (s *Service) handleGetThread(w http.ResponseWriter, request *http.Request) } } +func (s *Service) handleRenameThread(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w, request) + if store == nil { + return + } + body := struct { + Title string `json:"title"` + }{} + if err := json.NewDecoder(request.Body).Decode(&body); err != nil { + http.Error(w, fmt.Sprintf("invalid rename request: %v", err), http.StatusBadRequest) + return + } + if strings.TrimSpace(body.Title) == "" { + http.Error(w, "rename requires a title", http.StatusBadRequest) + return + } + id := request.PathValue("id") + if err := store.SetTitle(request.Context(), id, TitleUpdate{Title: body.Title, Source: TitleSourceUser}); err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + thread, err := store.Get(request.Context(), id) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + if err := writeJSON(w, http.StatusOK, thread); err != nil { + serviceLog.Errorf("write renamed chat thread %q: %v", id, err) + } +} + +func (s *Service) handleThreadCosts(w http.ResponseWriter, request *http.Request) { + store := s.threadStore(w, request) + if store == nil { + return + } + reader, ok := store.(ThreadCostReader) + if !ok { + http.Error(w, "thread cost breakdown requires a database-backed thread store", http.StatusNotImplemented) + return + } + costs, err := reader.GetThreadCosts(request.Context(), request.PathValue("id")) + if err != nil { + http.Error(w, err.Error(), http.StatusNotFound) + return + } + if err := writeJSON(w, http.StatusOK, costs); err != nil { + serviceLog.Errorf("write chat thread costs %q: %v", request.PathValue("id"), err) + } +} + func (s *Service) handleDeleteThread(w http.ResponseWriter, request *http.Request) { - store := s.threadStore(w) + store := s.threadStore(w, request) if store == nil { return } diff --git a/pkg/aichat/wire.go b/pkg/aichat/wire.go index 86ce1a95..9cd63e59 100644 --- a/pkg/aichat/wire.go +++ b/pkg/aichat/wire.go @@ -3,6 +3,7 @@ package aichat import ( "encoding/json" + "fmt" "strings" "github.com/flanksource/captain/pkg/ai" @@ -13,14 +14,17 @@ import ( // ChatRequest is the body posted by the AI SDK DefaultChatTransport. type ChatRequest struct { ID string `json:"id,omitempty"` + Trigger string `json:"trigger,omitempty"` + MessageID string `json:"messageId,omitempty"` Messages []UIMessage `json:"messages"` Model string `json:"model,omitempty"` + Runtime *api.Model `json:"runtime,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"` + ToolApproval *api.ToolApprovalResume `json:"-"` Context string `json:"context,omitempty"` ContextItems []ChatContextItem `json:"contextItems,omitempty"` @@ -29,6 +33,18 @@ type ChatRequest struct { ProviderSessionID string `json:"providerSessionId,omitempty"` } +func (r *ChatRequest) UnmarshalJSON(data []byte) error { + var fields map[string]json.RawMessage + if err := json.Unmarshal(data, &fields); err != nil { + return err + } + if _, exists := fields["toolApproval"]; exists { + return fmt.Errorf("toolApproval is server-owned; resolve approvals through the Captain session approval endpoint") + } + type wireChatRequest ChatRequest + return json.Unmarshal(data, (*wireChatRequest)(r)) +} + // ChatContextItem carries app-owned structured state alongside its readable label. type ChatContextItem struct { ID string `json:"id,omitempty"` @@ -43,6 +59,7 @@ type UIMessage struct { ID string `json:"id,omitempty"` Role string `json:"role"` Parts []UIPart `json:"parts"` + TurnID string `json:"turnId,omitempty"` Metadata *MessageMetadata `json:"metadata,omitempty"` } diff --git a/pkg/aichat/wire_ginkgo_test.go b/pkg/aichat/wire_ginkgo_test.go index b7a60901..31aa5104 100644 --- a/pkg/aichat/wire_ginkgo_test.go +++ b/pkg/aichat/wire_ginkgo_test.go @@ -14,6 +14,8 @@ var _ = Describe("AI SDK v6 wire types", func() { It("decodes the DefaultChatTransport request without losing UI parts", func() { const body = `{ "id":"chat-1", + "trigger":"regenerate-message", + "messageId":"message-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}}, @@ -34,6 +36,8 @@ var _ = Describe("AI SDK v6 wire types", func() { var request aichat.ChatRequest Expect(json.Unmarshal([]byte(body), &request)).To(Succeed()) Expect(request.ID).To(Equal("chat-1")) + Expect(request.Trigger).To(Equal("regenerate-message")) + Expect(request.MessageID).To(Equal("message-1")) Expect(request.Model).To(Equal("anthropic/claude-sonnet")) Expect(request.ReasoningEffort).To(Equal(api.EffortHigh)) Expect(request.Temperature).NotTo(BeNil()) @@ -62,10 +66,23 @@ var _ = Describe("AI SDK v6 wire types", func() { Expect(aichat.UIPart{Type: "text"}.EffectiveToolName()).To(BeEmpty()) }) - It("rejects the removed string tool approval policy", func() { + It("decodes an exact structured runtime", func() { var request aichat.ChatRequest - Expect(json.Unmarshal([]byte(`{"messages":[],"toolApproval":"manual"}`), &request)).To( - MatchError(ContainSubstring("cannot unmarshal string")), + Expect(json.Unmarshal([]byte(`{ + "runtime":{"model":"sonnet","backend":"claude-agent","effort":"high"}, + "messages":[{"role":"user","parts":[{"type":"text","text":"hello"}]}] + }`), &request)).To(Succeed()) + + Expect(request.Runtime).NotTo(BeNil()) + Expect(*request.Runtime).To(Equal(api.Model{ + Name: "sonnet", Backend: api.BackendClaudeAgent, Effort: api.EffortHigh, + })) + }) + + It("rejects client-owned tool approval state", func() { + var request aichat.ChatRequest + Expect(json.Unmarshal([]byte(`{"messages":[],"toolApproval":{"state":{}}}`), &request)).To( + MatchError(ContainSubstring("toolApproval is server-owned")), ) }) @@ -73,8 +90,9 @@ var _ = Describe("AI SDK v6 wire types", func() { strict := true models := aichat.ModelCatalogResponse{{ ID: "openai/gpt", Provider: "openai", Label: "GPT", Reasoning: true, - Temperature: true, Configured: true, ContextWindow: 128000, + Temperature: true, Configured: true, Availability: api.Available(), ContextWindow: 128000, InputMediaTypes: []string{"image/*"}, + Runtime: api.Model{Name: "gpt", Backend: api.BackendOpenAI}, }} tools := aichat.ToolCatalogResponse{Tools: []aichat.ToolCatalogEntry{{ Name: "invoice_get", Source: "custom", Group: "billing", @@ -85,7 +103,7 @@ var _ = Describe("AI SDK v6 wire types", func() { 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/*"]}]`)) + Expect(modelJSON).To(MatchJSON(`[{"id":"openai/gpt","provider":"openai","label":"GPT","runtime":{"model":"gpt","backend":"openai"},"reasoning":true,"temperature":true,"configured":true,"availability":{"state":"available"},"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/aimock/anthropicmock/anthropicmock_suite_test.go b/pkg/aimock/anthropicmock/anthropicmock_suite_test.go new file mode 100644 index 00000000..cd0e0e2a --- /dev/null +++ b/pkg/aimock/anthropicmock/anthropicmock_suite_test.go @@ -0,0 +1,13 @@ +package anthropicmock + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestAnthropicMock(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Anthropic Mock Suite") +} diff --git a/pkg/aimock/anthropicmock/health_ginkgo_test.go b/pkg/aimock/anthropicmock/health_ginkgo_test.go new file mode 100644 index 00000000..80f4d49d --- /dev/null +++ b/pkg/aimock/anthropicmock/health_ginkgo_test.go @@ -0,0 +1,32 @@ +package anthropicmock + +import ( + "net/http" + + "github.com/flanksource/captain/pkg/aimock" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Anthropic mock health probe", func() { + It("records the Claude SDK root HEAD probe without a miss", func() { + scenario, err := aimock.Parse([]byte(` +anthropic: + - respond: {text: unused} +`)) + Expect(err).NotTo(HaveOccurred()) + server, err := Start(Options{Scenario: scenario}) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(server.Close) + + request, err := http.NewRequest(http.MethodHead, server.URL()+"/", nil) + Expect(err).NotTo(HaveOccurred()) + response, err := http.DefaultClient.Do(request) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(response.Body.Close) + + Expect(response.StatusCode).To(Equal(http.StatusOK)) + Expect(server.Requests()).To(HaveLen(1)) + Expect(server.Requests()[0].Miss).To(BeEmpty()) + }) +}) diff --git a/pkg/aimock/anthropicmock/respond.go b/pkg/aimock/anthropicmock/respond.go index 7f091c6f..3714ddee 100644 --- a/pkg/aimock/anthropicmock/respond.go +++ b/pkg/aimock/anthropicmock/respond.go @@ -23,8 +23,9 @@ type Respond struct { Text string `json:"text,omitempty" yaml:"text,omitempty"` ToolUse *ToolUse `json:"tool_use,omitempty" yaml:"tool_use,omitempty"` - StopReason string `json:"stop_reason,omitempty" yaml:"stop_reason,omitempty"` - Usage Usage `json:"usage,omitempty" yaml:"usage,omitempty"` + StopReason string `json:"stop_reason,omitempty" yaml:"stop_reason,omitempty"` + Usage Usage `json:"usage,omitempty" yaml:"usage,omitempty"` + HoldOpenAfterContent bool `json:"hold_open_after_content,omitempty" yaml:"hold_open_after_content,omitempty"` // Error, when set, makes this rule return an API error instead of a reply — // for exercising the retry and error-mapping paths. diff --git a/pkg/aimock/anthropicmock/server.go b/pkg/aimock/anthropicmock/server.go index 9d9ef2cd..099eaa77 100644 --- a/pkg/aimock/anthropicmock/server.go +++ b/pkg/aimock/anthropicmock/server.go @@ -64,6 +64,7 @@ func Start(opts Options) (*Server, error) { mux.HandleFunc("POST /v1/messages", srv.handleMessages) mux.HandleFunc("POST /v1/messages/count_tokens", srv.handleCountTokens) mux.HandleFunc("GET /v1/models", srv.handleModels) + mux.HandleFunc("HEAD /{$}", srv.handleHealth) mux.HandleFunc("/", srv.handleUnknown) if err := srv.Listen(opts.Addr, mux, journal); err != nil { @@ -145,11 +146,15 @@ func (s *Server) handleMessages(w http.ResponseWriter, r *http.Request) { // stream can fail, so there is no status left to set — the note goes in // the journal, where a test asserting on Requests() will see it. note := "" - if err := streamMessage(w, model, respond); err != nil { - note = fmt.Sprintf("stream aborted: %v", err) - logger.Errorf("anthropicmock: %s", note) + cancelled := false + if err := streamMessage(r.Context(), w, model, respond); err != nil { + cancelled = aimock.IsClientCancellation(r.Context(), err) + if !cancelled { + note = fmt.Sprintf("stream aborted: %v", err) + logger.Errorf("anthropicmock: %s", note) + } } - s.record(r, norm, http.StatusOK, note) + s.recordOutcome(r, norm, http.StatusOK, note, cancelled) return } @@ -201,6 +206,11 @@ func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) { }) } +func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) { + s.record(r, aimock.Request{}, http.StatusOK, "") + w.WriteHeader(http.StatusOK) +} + // handleUnknown fails loudly on an unrouted path rather than 404-ing quietly, // so a client reaching for an endpoint the mock does not implement shows up as // a named gap instead of an opaque client-side error. @@ -216,14 +226,19 @@ func (s *Server) writeError(w http.ResponseWriter, r *http.Request, norm aimock. } func (s *Server) record(r *http.Request, norm aimock.Request, status int, miss string) { + s.recordOutcome(r, norm, status, miss, false) +} + +func (s *Server) recordOutcome(r *http.Request, norm aimock.Request, status int, miss string, cancelled bool) { s.Journal().Record(aimock.Recorded{ - Method: r.Method, - Path: r.URL.Path, - Status: status, - Stream: norm.Stream, - Model: norm.Model, - Request: norm, - Miss: miss, + Method: r.Method, + Path: r.URL.Path, + Status: status, + Stream: norm.Stream, + Model: norm.Model, + Request: norm, + Miss: miss, + Cancelled: cancelled, }) } diff --git a/pkg/aimock/anthropicmock/server_test.go b/pkg/aimock/anthropicmock/server_test.go index 6d824376..fb540d15 100644 --- a/pkg/aimock/anthropicmock/server_test.go +++ b/pkg/aimock/anthropicmock/server_test.go @@ -2,10 +2,13 @@ package anthropicmock import ( "bytes" + "context" "encoding/json" + "io" "net/http" "path/filepath" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -13,6 +16,27 @@ import ( "github.com/flanksource/captain/pkg/aimock" ) +func TestHeldStreamRecordsCancellationWithoutAMiss(t *testing.T) { + srv := startServer(t, "hold-open.yaml") + raw, err := json.Marshal(map[string]any{ + "model": "claude-sonnet-5", "stream": true, + "messages": []any{userTurn("wait for interruption")}, + }) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + request, err := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL()+"/v1/messages", bytes.NewReader(raw)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + require.NoError(t, err) + cancel() + _, _ = io.ReadAll(response.Body) + _ = response.Body.Close() + require.Eventually(t, func() bool { return len(srv.Requests()) == 1 }, time.Second, 10*time.Millisecond) + assert.True(t, srv.Requests()[0].Cancelled) + assert.Empty(t, srv.Requests()[0].Miss) +} + const scenarioDir = "../testdata/scenarios" func startServer(t *testing.T, scenarioFile string, opts ...func(*Options)) *Server { diff --git a/pkg/aimock/anthropicmock/stream.go b/pkg/aimock/anthropicmock/stream.go index 469dafca..7c6ba77d 100644 --- a/pkg/aimock/anthropicmock/stream.go +++ b/pkg/aimock/anthropicmock/stream.go @@ -4,6 +4,7 @@ package anthropicmock import ( + "context" "fmt" "net/http" @@ -21,7 +22,7 @@ const toolInputChunk = 24 // // An error here arrives after the 200 and the first frames are already on the // wire, so it cannot become a status — the caller journals it instead. -func streamMessage(w http.ResponseWriter, model string, respond Respond) error { +func streamMessage(ctx context.Context, w http.ResponseWriter, model string, respond Respond) error { sse, err := aimock.NewSSE(w) if err != nil { return err @@ -59,6 +60,9 @@ func streamMessage(w http.ResponseWriter, model string, respond Respond) error { return err } } + if err := aimock.WaitForCancellation(ctx, respond.HoldOpenAfterContent); err != nil { + return err + } if err := sse.Event("message_delta", messageDeltaFrame{ Type: "message_delta", diff --git a/pkg/aimock/anthropicmock/wire.go b/pkg/aimock/anthropicmock/wire.go index 71d1091e..d6509330 100644 --- a/pkg/aimock/anthropicmock/wire.go +++ b/pkg/aimock/anthropicmock/wire.go @@ -17,10 +17,16 @@ type messagesRequest struct { Model string `json:"model"` System json.RawMessage `json:"system,omitempty"` Messages []wireMessage `json:"messages"` + Tools []wireTool `json:"tools,omitempty"` Stream bool `json:"stream,omitempty"` MaxTokens int `json:"max_tokens,omitempty"` } +type wireTool struct { + Name string `json:"name"` + InputSchema json.RawMessage `json:"input_schema,omitempty"` +} + type wireMessage struct { Role string `json:"role"` Content json.RawMessage `json:"content"` @@ -52,6 +58,17 @@ func decodeRequest(r *http.Request, body []byte) (messagesRequest, aimock.Reques Stream: wire.Stream, Headers: headerMap(r), } + for _, tool := range wire.Tools { + if tool.Name != "" { + norm.ToolNames = append(norm.ToolNames, tool.Name) + if len(tool.InputSchema) > 0 { + if norm.ToolSchemas == nil { + norm.ToolSchemas = map[string]json.RawMessage{} + } + norm.ToolSchemas[tool.Name] = tool.InputSchema + } + } + } // tool_use ids seen so far, so a later tool_result can be resolved back to // the tool name the scenario matches on. diff --git a/pkg/aimock/journal.go b/pkg/aimock/journal.go index 6b595bd8..be9c7a0a 100644 --- a/pkg/aimock/journal.go +++ b/pkg/aimock/journal.go @@ -26,7 +26,8 @@ type Recorded struct { // Miss is the diagnostic for a request that produced no clean scripted // reply — no rule matched, the route is unimplemented, or the stream aborted // mid-flight. Empty on a normal request. - Miss string `json:"miss,omitempty"` + Miss string `json:"miss,omitempty"` + Cancelled bool `json:"cancelled,omitempty"` } // Journal records every served request in memory and, when opened with a path, diff --git a/pkg/aimock/openaimock/cancellation_test.go b/pkg/aimock/openaimock/cancellation_test.go new file mode 100644 index 00000000..3876444d --- /dev/null +++ b/pkg/aimock/openaimock/cancellation_test.go @@ -0,0 +1,48 @@ +package openaimock + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestHeldStreamsRecordCancellationWithoutAMiss(t *testing.T) { + for _, test := range []struct { + name string + path string + body map[string]any + }{ + {name: "responses", path: "/v1/responses", body: map[string]any{ + "model": "gpt-5", "stream": true, "input": userInput("wait for interruption"), + }}, + {name: "chat completions", path: "/v1/chat/completions", body: map[string]any{ + "model": "gpt-5", "stream": true, + "messages": []any{map[string]any{"role": "user", "content": "wait for interruption"}}, + }}, + } { + t.Run(test.name, func(t *testing.T) { + srv := startServer(t, "hold-open.yaml") + raw, err := json.Marshal(test.body) + require.NoError(t, err) + ctx, cancel := context.WithCancel(t.Context()) + request, err := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL()+test.path, bytes.NewReader(raw)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + require.NoError(t, err) + cancel() + _, _ = io.ReadAll(response.Body) + _ = response.Body.Close() + require.Eventually(t, func() bool { return len(srv.Requests()) == 1 }, time.Second, 10*time.Millisecond) + assert.True(t, srv.Requests()[0].Cancelled) + assert.Empty(t, srv.Requests()[0].Miss) + }) + } +} diff --git a/pkg/aimock/openaimock/chat.go b/pkg/aimock/openaimock/chat.go index d634d31a..f245fc34 100644 --- a/pkg/aimock/openaimock/chat.go +++ b/pkg/aimock/openaimock/chat.go @@ -4,6 +4,7 @@ package openaimock import ( + "context" "fmt" "io" "net/http" @@ -32,29 +33,34 @@ func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) { } model := modelOrDefault(wire.Model) + completionID := s.nextWireID("chatcmpl", model) if wire.Stream { note := "" - if err := streamChat(w, model, respond); err != nil { - note = fmt.Sprintf("stream aborted: %v", err) - logger.Errorf("openaimock: %s", note) + cancelled := false + if err := streamChat(r.Context(), w, completionID, model, respond); err != nil { + cancelled = aimock.IsClientCancellation(r.Context(), err) + if !cancelled { + note = fmt.Sprintf("stream aborted: %v", err) + logger.Errorf("openaimock: %s", note) + } } - s.record(r, norm, http.StatusOK, note) + s.recordOutcome(r, norm, http.StatusOK, note, cancelled) return } s.record(r, norm, http.StatusOK, "") - writeJSON(w, http.StatusOK, chatCompletion(model, respond)) + writeJSON(w, http.StatusOK, chatCompletion(completionID, model, respond)) } // chatCompletion renders the reply as a complete non-streaming completion. // Reasoning rides on the non-standard `reasoning_content` field, which is what // the deepseek-compatible endpoints captain talks to actually emit. -func chatCompletion(model string, respond Respond) map[string]any { +func chatCompletion(completionID, model string, respond Respond) map[string]any { message := map[string]any{"role": "assistant", "content": respond.Text} if respond.Reasoning != "" { message["reasoning_content"] = respond.Reasoning } - if call := chatToolCallPayload(respond); call != nil { + if call := chatToolCallPayload(respond, completionID); call != nil { message["tool_calls"] = []any{call} // A tool-calling choice carries no prose; content is explicitly null // rather than "" so a consumer distinguishing the two sees the right one. @@ -62,7 +68,7 @@ func chatCompletion(model string, respond Respond) map[string]any { } return map[string]any{ - "id": completionID(model), + "id": completionID, "object": "chat.completion", "created": 0, "model": model, @@ -77,8 +83,8 @@ func chatCompletion(model string, respond Respond) map[string]any { // chatToolCallPayload renders the scripted function call in the tool_calls shape, // or nil when the reply makes no call. -func chatToolCallPayload(respond Respond) map[string]any { - for _, it := range respond.items() { +func chatToolCallPayload(respond Respond, completionID string) map[string]any { + for _, it := range respond.items(completionID) { if it.Type != "function_call" { continue } @@ -95,7 +101,7 @@ func chatToolCallPayload(respond Respond) map[string]any { // streamChat renders respond as chat.completion.chunk frames: an opening role // delta, one delta per content chunk, a terminal finish_reason delta, a // usage-only chunk, then the [DONE] sentinel. -func streamChat(w http.ResponseWriter, model string, respond Respond) error { +func streamChat(ctx context.Context, w http.ResponseWriter, completionID, model string, respond Respond) error { sse, err := aimock.NewSSE(w) if err != nil { return err @@ -103,7 +109,7 @@ func streamChat(w http.ResponseWriter, model string, respond Respond) error { chunk := func(delta map[string]any, finish any) error { return sse.Data(map[string]any{ - "id": completionID(model), "object": "chat.completion.chunk", "created": 0, "model": model, + "id": completionID, "object": "chat.completion.chunk", "created": 0, "model": model, "choices": []any{map[string]any{"index": 0, "delta": delta, "finish_reason": finish}}, }) } @@ -118,7 +124,7 @@ func streamChat(w http.ResponseWriter, model string, respond Respond) error { } } - if call := chatToolCallPayload(respond); call != nil { + if call := chatToolCallPayload(respond, completionID); call != nil { if err := streamChatToolCall(chunk, call); err != nil { return err } @@ -129,6 +135,9 @@ func streamChat(w http.ResponseWriter, model string, respond Respond) error { return err } } + if err := aimock.WaitForCancellation(ctx, respond.HoldOpenAfterContent); err != nil { + return err + } if err := chunk(map[string]any{}, respond.resolvedFinishReason()); err != nil { return err @@ -137,7 +146,7 @@ func streamChat(w http.ResponseWriter, model string, respond Respond) error { // Usage arrives in its own choice-less chunk, matching what the API sends // under stream_options.include_usage. if err := sse.Data(map[string]any{ - "id": completionID(model), "object": "chat.completion.chunk", "created": 0, "model": model, + "id": completionID, "object": "chat.completion.chunk", "created": 0, "model": model, "choices": []any{}, "usage": respond.Usage.chat(), }); err != nil { return err @@ -168,5 +177,3 @@ func streamChatToolCall(chunk func(map[string]any, any) error, call map[string]a } return nil } - -func completionID(model string) string { return fmt.Sprintf("chatcmpl_mock_%s", model) } diff --git a/pkg/aimock/openaimock/namespace_ginkgo_test.go b/pkg/aimock/openaimock/namespace_ginkgo_test.go new file mode 100644 index 00000000..c373c246 --- /dev/null +++ b/pkg/aimock/openaimock/namespace_ginkgo_test.go @@ -0,0 +1,43 @@ +package openaimock + +import ( + "encoding/json" + "net/http/httptest" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Responses namespace tools", func() { + It("emits and normalizes the namespace separately from the function name", func() { + response := Respond{FunctionCall: &FunctionCall{ + Namespace: "mcp__captain", Name: "accounts_edit", CallID: "call_account", + Arguments: map[string]any{"id": "acc-1"}, + }} + item := response.items("resp_mock_namespace_1")[0].done() + Expect(item).To(HaveKeyWithValue("namespace", "mcp__captain")) + Expect(item).To(HaveKeyWithValue("name", "accounts_edit")) + + body, err := json.Marshal(map[string]any{ + "model": "gpt-5", "input": []any{ + item, + map[string]any{"type": "function_call_output", "call_id": "call_account", "output": "updated"}, + }, + }) + Expect(err).NotTo(HaveOccurred()) + request := httptest.NewRequest("POST", "/v1/responses", nil) + _, normalized, err := decodeResponses(request, body) + Expect(err).NotTo(HaveOccurred()) + Expect(normalized.ToolResultNames()).To(Equal([]string{"mcp__captain__accounts_edit"})) + }) + + It("assigns distinct response and item identities to successive requests", func() { + server := &Server{} + first := server.nextWireID("resp", "gpt-5") + second := server.nextWireID("resp", "gpt-5") + response := Respond{Text: "done"} + + Expect(second).NotTo(Equal(first)) + Expect(response.items(second)[0].ID).NotTo(Equal(response.items(first)[0].ID)) + }) +}) diff --git a/pkg/aimock/openaimock/openaimock_suite_test.go b/pkg/aimock/openaimock/openaimock_suite_test.go new file mode 100644 index 00000000..c536d44a --- /dev/null +++ b/pkg/aimock/openaimock/openaimock_suite_test.go @@ -0,0 +1,13 @@ +package openaimock + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestOpenAIMock(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "OpenAI Mock Suite") +} diff --git a/pkg/aimock/openaimock/respond.go b/pkg/aimock/openaimock/respond.go index 48ba1898..91fb5e5c 100644 --- a/pkg/aimock/openaimock/respond.go +++ b/pkg/aimock/openaimock/respond.go @@ -24,8 +24,9 @@ type Respond struct { Text string `json:"text,omitempty" yaml:"text,omitempty"` FunctionCall *FunctionCall `json:"function_call,omitempty" yaml:"function_call,omitempty"` - FinishReason string `json:"finish_reason,omitempty" yaml:"finish_reason,omitempty"` - Usage Usage `json:"usage,omitempty" yaml:"usage,omitempty"` + FinishReason string `json:"finish_reason,omitempty" yaml:"finish_reason,omitempty"` + Usage Usage `json:"usage,omitempty" yaml:"usage,omitempty"` + HoldOpenAfterContent bool `json:"hold_open_after_content,omitempty" yaml:"hold_open_after_content,omitempty"` // Error, when set, makes this rule return an API error instead of a reply — // for exercising the retry and error-mapping paths. @@ -35,6 +36,7 @@ type Respond struct { // FunctionCall is a scripted tool call. Arguments are written as YAML and // marshalled to the JSON string the wire carries. type FunctionCall struct { + Namespace string `json:"namespace,omitempty" yaml:"namespace,omitempty"` Name string `json:"name" yaml:"name"` CallID string `json:"call_id,omitempty" yaml:"call_id,omitempty"` Arguments map[string]any `json:"arguments,omitempty" yaml:"arguments,omitempty"` @@ -76,6 +78,7 @@ type item struct { Type string ID string Text string + Namespace string Name string CallID string Arguments string @@ -84,10 +87,10 @@ type item struct { // items renders the reply into ordered output items. A reply with neither text // nor a tool call still produces an empty message, because every Responses reply // has at least one output item. -func (r Respond) items() []item { +func (r Respond) items(responseID string) []item { var out []item if r.Reasoning != "" { - out = append(out, item{Type: "reasoning", ID: fmt.Sprintf("rs_mock_%d", len(out)), Text: r.Reasoning}) + out = append(out, item{Type: "reasoning", ID: fmt.Sprintf("rs_%s_%d", responseID, len(out)), Text: r.Reasoning}) } if r.FunctionCall != nil { callID := r.FunctionCall.CallID @@ -107,14 +110,15 @@ func (r Respond) items() []item { } out = append(out, item{ Type: "function_call", - ID: fmt.Sprintf("fc_mock_%d", len(out)), + ID: fmt.Sprintf("fc_%s_%d", responseID, len(out)), + Namespace: r.FunctionCall.Namespace, Name: r.FunctionCall.Name, CallID: callID, Arguments: string(raw), }) } if r.Text != "" || len(out) == 0 { - out = append(out, item{Type: "message", ID: fmt.Sprintf("msg_mock_%d", len(out)), Text: r.Text}) + out = append(out, item{Type: "message", ID: fmt.Sprintf("msg_%s_%d", responseID, len(out)), Text: r.Text}) } return out } @@ -126,7 +130,11 @@ func (i item) added() map[string]any { case "reasoning": return map[string]any{"id": i.ID, "type": "reasoning", "summary": []any{}} case "function_call": - return map[string]any{"id": i.ID, "type": "function_call", "status": "in_progress", "name": i.Name, "call_id": i.CallID, "arguments": ""} + payload := map[string]any{"id": i.ID, "type": "function_call", "status": "in_progress", "name": i.Name, "call_id": i.CallID, "arguments": ""} + if i.Namespace != "" { + payload["namespace"] = i.Namespace + } + return payload default: return map[string]any{"id": i.ID, "type": "message", "status": "in_progress", "role": "assistant", "content": []any{}} } @@ -142,7 +150,11 @@ func (i item) done() map[string]any { "summary": []any{map[string]any{"type": "summary_text", "text": i.Text}}, } case "function_call": - return map[string]any{"id": i.ID, "type": "function_call", "status": "completed", "name": i.Name, "call_id": i.CallID, "arguments": i.Arguments} + payload := map[string]any{"id": i.ID, "type": "function_call", "status": "completed", "name": i.Name, "call_id": i.CallID, "arguments": i.Arguments} + if i.Namespace != "" { + payload["namespace"] = i.Namespace + } + return payload default: return map[string]any{ "id": i.ID, "type": "message", "status": "completed", "role": "assistant", diff --git a/pkg/aimock/openaimock/responses.go b/pkg/aimock/openaimock/responses.go index 61cf1798..68863fb5 100644 --- a/pkg/aimock/openaimock/responses.go +++ b/pkg/aimock/openaimock/responses.go @@ -4,6 +4,7 @@ package openaimock import ( + "context" "fmt" "io" "net/http" @@ -37,33 +38,38 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) { } model := modelOrDefault(wire.Model) + responseID := s.nextWireID("resp", model) if wire.Stream { // The 200 and the first frames are already on the wire by the time a // stream can fail, so there is no status left to set — the note goes in // the journal, where a test asserting on Requests() will see it. note := "" - if err := streamResponses(w, model, respond); err != nil { - note = fmt.Sprintf("stream aborted: %v", err) - logger.Errorf("openaimock: %s", note) + cancelled := false + if err := streamResponses(r.Context(), w, responseID, model, respond); err != nil { + cancelled = aimock.IsClientCancellation(r.Context(), err) + if !cancelled { + note = fmt.Sprintf("stream aborted: %v", err) + logger.Errorf("openaimock: %s", note) + } } - s.record(r, norm, http.StatusOK, note) + s.recordOutcome(r, norm, http.StatusOK, note, cancelled) return } s.record(r, norm, http.StatusOK, "") - writeJSON(w, http.StatusOK, completedResponse(model, respond)) + writeJSON(w, http.StatusOK, completedResponse(responseID, model, respond)) } // completedResponse is the whole reply, used both as the non-streaming body and // as the payload of the terminal response.completed frame. -func completedResponse(model string, respond Respond) map[string]any { - items := respond.items() +func completedResponse(responseID, model string, respond Respond) map[string]any { + items := respond.items(responseID) output := make([]map[string]any, 0, len(items)) for _, it := range items { output = append(output, it.done()) } return map[string]any{ - "id": responseID(model), + "id": responseID, "object": "response", "status": "completed", "model": model, @@ -76,9 +82,9 @@ func completedResponse(model string, respond Respond) map[string]any { // inProgressResponse is the envelope carried by response.created and // response.in_progress: identity only, with no output and no usage yet. -func inProgressResponse(model string) map[string]any { +func inProgressResponse(responseID, model string) map[string]any { return map[string]any{ - "id": responseID(model), + "id": responseID, "object": "response", "status": "in_progress", "model": model, @@ -92,13 +98,13 @@ func inProgressResponse(model string) map[string]any { // streamResponses renders respond as the documented Responses API event // sequence: response.created / .in_progress → per item (output_item.added, the // item's own delta and done frames, output_item.done) → response.completed. -func streamResponses(w http.ResponseWriter, model string, respond Respond) error { +func streamResponses(ctx context.Context, w http.ResponseWriter, responseID, model string, respond Respond) error { sse, err := aimock.NewSSE(w) if err != nil { return err } - envelope := inProgressResponse(model) + envelope := inProgressResponse(responseID, model) if err := sse.Event("response.created", map[string]any{"type": "response.created", "response": envelope}); err != nil { return err } @@ -106,15 +112,18 @@ func streamResponses(w http.ResponseWriter, model string, respond Respond) error return err } - for index, it := range respond.items() { + for index, it := range respond.items(responseID) { if err := streamItem(sse, index, it); err != nil { return err } } + if err := aimock.WaitForCancellation(ctx, respond.HoldOpenAfterContent); err != nil { + return err + } return sse.Event("response.completed", map[string]any{ "type": "response.completed", - "response": completedResponse(model, respond), + "response": completedResponse(responseID, model, respond), }) } @@ -239,8 +248,6 @@ func chunkRunes(raw string, size int) []string { return chunks } -func responseID(model string) string { return fmt.Sprintf("resp_mock_%s", model) } - func modelOrDefault(model string) string { if model == "" { return "gpt-mock" diff --git a/pkg/aimock/openaimock/server.go b/pkg/aimock/openaimock/server.go index b14abe92..634ec4aa 100644 --- a/pkg/aimock/openaimock/server.go +++ b/pkg/aimock/openaimock/server.go @@ -7,6 +7,7 @@ import ( "encoding/json" "fmt" "net/http" + "sync/atomic" "github.com/flanksource/captain/pkg/aimock" ) @@ -29,7 +30,8 @@ type Options struct { // Server is a mock OpenAI API serving both wire APIs from one scenario section. type Server struct { aimock.Base - rules *aimock.Rules[Respond] + rules *aimock.Rules[Respond] + sequence atomic.Uint64 } var _ aimock.Server = (*Server)(nil) @@ -100,6 +102,10 @@ func (s *Server) Env() []string { return Env(s.URL()) } // a run played the whole scenario. func (s *Server) Remaining() []string { return s.rules.Remaining() } +func (s *Server) nextWireID(kind, model string) string { + return fmt.Sprintf("%s_mock_%s_%d", kind, model, s.sequence.Add(1)) +} + // resolve picks the scripted reply for a request, writing the miss diagnostic or // the scripted error itself and reporting false when there is nothing to serve. func (s *Server) resolve(w http.ResponseWriter, r *http.Request, norm aimock.Request) (Respond, bool) { @@ -157,14 +163,19 @@ func (s *Server) writeError(w http.ResponseWriter, r *http.Request, norm aimock. } func (s *Server) record(r *http.Request, norm aimock.Request, status int, miss string) { + s.recordOutcome(r, norm, status, miss, false) +} + +func (s *Server) recordOutcome(r *http.Request, norm aimock.Request, status int, miss string, cancelled bool) { s.Journal().Record(aimock.Recorded{ - Method: r.Method, - Path: r.URL.Path, - Status: status, - Stream: norm.Stream, - Model: norm.Model, - Request: norm, - Miss: miss, + Method: r.Method, + Path: r.URL.Path, + Status: status, + Stream: norm.Stream, + Model: norm.Model, + Request: norm, + Miss: miss, + Cancelled: cancelled, }) } diff --git a/pkg/aimock/openaimock/wire.go b/pkg/aimock/openaimock/wire.go index eebf8d61..cdd1795b 100644 --- a/pkg/aimock/openaimock/wire.go +++ b/pkg/aimock/openaimock/wire.go @@ -14,10 +14,18 @@ import ( // responsesRequest is the subset of a /v1/responses body worth matching on. type responsesRequest struct { - Model string `json:"model"` - Instructions string `json:"instructions,omitempty"` - Input json.RawMessage `json:"input"` - Stream bool `json:"stream,omitempty"` + Model string `json:"model"` + Instructions string `json:"instructions,omitempty"` + Input json.RawMessage `json:"input"` + Tools []json.RawMessage `json:"tools,omitempty"` + Stream bool `json:"stream,omitempty"` +} + +type wireTool struct { + Type string `json:"type,omitempty"` + Name string `json:"name"` + Parameters json.RawMessage `json:"parameters,omitempty"` + InputSchema json.RawMessage `json:"input_schema,omitempty"` } // inputItem covers every entry shape the Responses API accepts in `input`: @@ -27,6 +35,7 @@ type inputItem struct { Type string `json:"type,omitempty"` Role string `json:"role,omitempty"` Content json.RawMessage `json:"content,omitempty"` + Namespace string `json:"namespace,omitempty"` Name string `json:"name,omitempty"` CallID string `json:"call_id,omitempty"` Arguments string `json:"arguments,omitempty"` @@ -53,6 +62,29 @@ func decodeResponses(r *http.Request, body []byte) (responsesRequest, aimock.Req Stream: wire.Stream, Headers: headerMap(r), } + for _, definition := range wire.Tools { + var tool wireTool + if err := json.Unmarshal(definition, &tool); err != nil { + return wire, norm, fmt.Errorf("decode responses tool: %w", err) + } + if tool.Name != "" { + norm.ToolNames = append(norm.ToolNames, tool.Name) + if norm.ToolDefinitions == nil { + norm.ToolDefinitions = map[string]json.RawMessage{} + } + norm.ToolDefinitions[tool.Name] = definition + schema := tool.Parameters + if len(schema) == 0 { + schema = tool.InputSchema + } + if len(schema) > 0 { + if norm.ToolSchemas == nil { + norm.ToolSchemas = map[string]json.RawMessage{} + } + norm.ToolSchemas[tool.Name] = schema + } + } + } // A bare string input is the single-user-turn shorthand. var text string @@ -74,6 +106,9 @@ func decodeResponses(r *http.Request, body []byte) (responsesRequest, aimock.Req case "function_call": if in.CallID != "" && in.Name != "" { callNames[in.CallID] = in.Name + if in.Namespace != "" { + callNames[in.CallID] = in.Namespace + "__" + in.Name + } } norm.Messages = append(norm.Messages, aimock.Message{Role: aimock.RoleAssistant, Content: in.Arguments}) case "function_call_output": @@ -102,9 +137,17 @@ func decodeResponses(r *http.Request, body []byte) (responsesRequest, aimock.Req type chatRequest struct { Model string `json:"model"` Messages []chatMessage `json:"messages"` + Tools []chatTool `json:"tools,omitempty"` Stream bool `json:"stream,omitempty"` } +type chatTool struct { + Function struct { + Name string `json:"name"` + Parameters json.RawMessage `json:"parameters,omitempty"` + } `json:"function"` +} + type chatMessage struct { Role string `json:"role"` Content json.RawMessage `json:"content,omitempty"` @@ -129,6 +172,17 @@ func decodeChat(r *http.Request, body []byte) (chatRequest, aimock.Request, erro } norm := aimock.Request{Model: wire.Model, Stream: wire.Stream, Headers: headerMap(r)} + for _, tool := range wire.Tools { + if tool.Function.Name != "" { + norm.ToolNames = append(norm.ToolNames, tool.Function.Name) + if len(tool.Function.Parameters) > 0 { + if norm.ToolSchemas == nil { + norm.ToolSchemas = map[string]json.RawMessage{} + } + norm.ToolSchemas[tool.Function.Name] = tool.Function.Parameters + } + } + } callNames := map[string]string{} var systems []string diff --git a/pkg/aimock/request.go b/pkg/aimock/request.go index a69bab38..9bba4196 100644 --- a/pkg/aimock/request.go +++ b/pkg/aimock/request.go @@ -3,7 +3,10 @@ package aimock -import "strings" +import ( + "encoding/json" + "strings" +) // Role values on a normalized Message. Both wire protocols collapse onto these. const ( @@ -25,11 +28,14 @@ type Message struct { // the fields worth matching on. Each server builds one of these from its own // request type before consulting the rules. type Request struct { - Model string `json:"model,omitempty"` - System string `json:"system,omitempty"` - Messages []Message `json:"messages,omitempty"` - Headers map[string]string `json:"headers,omitempty"` - Stream bool `json:"stream,omitempty"` + Model string `json:"model,omitempty"` + System string `json:"system,omitempty"` + Messages []Message `json:"messages,omitempty"` + ToolNames []string `json:"toolNames,omitempty"` + ToolSchemas map[string]json.RawMessage `json:"toolSchemas,omitempty"` + ToolDefinitions map[string]json.RawMessage `json:"toolDefinitions,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + Stream bool `json:"stream,omitempty"` } // LastUserText is the content of the most recent user turn — the "prompt" that diff --git a/pkg/aimock/sse.go b/pkg/aimock/sse.go index fe46dcf9..554306db 100644 --- a/pkg/aimock/sse.go +++ b/pkg/aimock/sse.go @@ -4,12 +4,33 @@ package aimock import ( + "context" "encoding/json" + "errors" "fmt" + "net" "net/http" "strings" + "syscall" ) +func WaitForCancellation(ctx context.Context, hold bool) error { + if !hold { + return nil + } + <-ctx.Done() + return ctx.Err() +} + +func IsClientCancellation(ctx context.Context, err error) bool { + return errors.Is(err, context.Canceled) || + errors.Is(err, context.DeadlineExceeded) || + errors.Is(err, net.ErrClosed) || + errors.Is(err, syscall.EPIPE) || + errors.Is(err, syscall.ECONNRESET) || + ctx.Err() != nil +} + // SSE writes server-sent-event frames to an http.ResponseWriter. type SSE struct { w http.ResponseWriter diff --git a/pkg/aimock/testdata/scenarios/chat-agent-flows.yaml b/pkg/aimock/testdata/scenarios/chat-agent-flows.yaml new file mode 100644 index 00000000..58d74ee1 --- /dev/null +++ b/pkg/aimock/testdata/scenarios/chat-agent-flows.yaml @@ -0,0 +1,78 @@ +name: chat-agent-flows +description: Complete agent-provider chat lifecycle through request, approval, interruption, and resume. + +anthropic: + - match: {prompt_contains: "Return the lifecycle greeting"} + respond: + text: "Lifecycle response complete." + usage: {input: 20, output: 5} + - match: {prompt_contains: "Approve the account update"} + respond: + tool_use: + name: mcp__captain__accounts_edit + id: toolu_approve_account + input: {id: acc-1, name: Draft Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: mcp__captain__accounts_edit} + respond: + text: "Approved account update completed." + usage: {input: 40, output: 6} + - match: {prompt_contains: "Reject the account update"} + respond: + tool_use: + name: mcp__captain__accounts_edit + id: toolu_reject_account + input: {id: acc-2, name: Rejected Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: mcp__captain__accounts_edit} + respond: + text: "Rejected account update was not applied." + usage: {input: 40, output: 7} + - match: {prompt_contains: "Wait for the lifecycle interrupt"} + respond: + text: "Partial lifecycle response." + hold_open_after_content: true + usage: {input: 15, output: 4} + - match: {prompt_contains: "Resume after the lifecycle interrupt"} + respond: + text: "Lifecycle resume complete." + usage: {input: 25, output: 5} + +openai: + - match: {prompt_contains: "Return the lifecycle greeting"} + respond: + text: "Lifecycle response complete." + usage: {input: 20, output: 5} + - match: {prompt_contains: "Approve the account update"} + respond: + function_call: + namespace: mcp__captain + name: accounts_edit + call_id: call_approve_account + arguments: {id: acc-1, name: Draft Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: mcp__captain__accounts_edit} + respond: + text: "Approved account update completed." + usage: {input: 40, output: 6} + - match: {prompt_contains: "Reject the account update"} + respond: + function_call: + namespace: mcp__captain + name: accounts_edit + call_id: call_reject_account + arguments: {id: acc-2, name: Rejected Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: mcp__captain__accounts_edit} + respond: + text: "Rejected account update was not applied." + usage: {input: 40, output: 7} + - match: {prompt_contains: "Wait for the lifecycle interrupt"} + respond: + text: "Partial lifecycle response." + hold_open_after_content: true + usage: {input: 15, output: 4} + - match: {prompt_contains: "Resume after the lifecycle interrupt"} + respond: + text: "Lifecycle resume complete." + usage: {input: 25, output: 5} diff --git a/pkg/aimock/testdata/scenarios/chat-api-flows.yaml b/pkg/aimock/testdata/scenarios/chat-api-flows.yaml new file mode 100644 index 00000000..c5774928 --- /dev/null +++ b/pkg/aimock/testdata/scenarios/chat-api-flows.yaml @@ -0,0 +1,76 @@ +name: chat-api-flows +description: Complete API-provider chat lifecycle through request, approval, interruption, and resume. + +anthropic: + - match: {prompt_contains: "Return the lifecycle greeting"} + respond: + text: "Lifecycle response complete." + usage: {input: 20, output: 5} + - match: {prompt_contains: "Approve the account update"} + respond: + tool_use: + name: accounts_edit + id: toolu_approve_account + input: {id: acc-1, name: Draft Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: accounts_edit} + respond: + text: "Approved account update completed." + usage: {input: 40, output: 6} + - match: {prompt_contains: "Reject the account update"} + respond: + tool_use: + name: accounts_edit + id: toolu_reject_account + input: {id: acc-2, name: Rejected Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: accounts_edit} + respond: + text: "Rejected account update was not applied." + usage: {input: 40, output: 7} + - match: {prompt_contains: "Wait for the lifecycle interrupt"} + respond: + text: "Partial lifecycle response." + hold_open_after_content: true + usage: {input: 15, output: 4} + - match: {prompt_contains: "Resume after the lifecycle interrupt"} + respond: + text: "Lifecycle resume complete." + usage: {input: 25, output: 5} + +openai: + - match: {prompt_contains: "Return the lifecycle greeting"} + respond: + text: "Lifecycle response complete." + usage: {input: 20, output: 5} + - match: {prompt_contains: "Approve the account update"} + respond: + function_call: + name: accounts_edit + call_id: call_approve_account + arguments: {id: acc-1, name: Draft Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: accounts_edit} + respond: + text: "Approved account update completed." + usage: {input: 40, output: 6} + - match: {prompt_contains: "Reject the account update"} + respond: + function_call: + name: accounts_edit + call_id: call_reject_account + arguments: {id: acc-2, name: Rejected Account} + usage: {input: 30, output: 8} + - match: {tool_result_for: accounts_edit} + respond: + text: "Rejected account update was not applied." + usage: {input: 40, output: 7} + - match: {prompt_contains: "Wait for the lifecycle interrupt"} + respond: + text: "Partial lifecycle response." + hold_open_after_content: true + usage: {input: 15, output: 4} + - match: {prompt_contains: "Resume after the lifecycle interrupt"} + respond: + text: "Lifecycle resume complete." + usage: {input: 25, output: 5} diff --git a/pkg/aimock/testdata/scenarios/hold-open.yaml b/pkg/aimock/testdata/scenarios/hold-open.yaml new file mode 100644 index 00000000..9b0464d0 --- /dev/null +++ b/pkg/aimock/testdata/scenarios/hold-open.yaml @@ -0,0 +1,18 @@ +name: hold-open +description: Streams partial content and waits for the caller to interrupt it. + +anthropic: + - match: + prompt_contains: "wait for interruption" + respond: + text: "Partial response before interruption." + hold_open_after_content: true + usage: {input: 12, output: 4} + +openai: + - match: + prompt_contains: "wait for interruption" + respond: + text: "Partial response before interruption." + hold_open_after_content: true + usage: {input: 12, output: 4} diff --git a/pkg/api/availability.go b/pkg/api/availability.go new file mode 100644 index 00000000..cc2bf605 --- /dev/null +++ b/pkg/api/availability.go @@ -0,0 +1,32 @@ +package api + +// AvailabilityState is a stable machine-readable reason a runtime or model +// cannot currently be selected. +type AvailabilityState string + +const ( + AvailabilityAvailable AvailabilityState = "available" + AvailabilityDisabled AvailabilityState = "disabled" + AvailabilityMissingCredential AvailabilityState = "missing_credentials" + AvailabilityNotAuthenticated AvailabilityState = "not_authenticated" + AvailabilityMissingExecutable AvailabilityState = "missing_executable" + AvailabilityMissingDependency AvailabilityState = "missing_dependency" + AvailabilityUnsupported AvailabilityState = "unsupported" + AvailabilityUnavailable AvailabilityState = "unavailable" +) + +// Availability carries presentation-safe readiness details. Reason explains +// the current state; Remediation tells the user how to make it selectable. +type Availability struct { + State AvailabilityState `json:"state"` + Reason string `json:"reason,omitempty"` + Remediation string `json:"remediation,omitempty"` +} + +func Available() Availability { + return Availability{State: AvailabilityAvailable} +} + +func (a Availability) IsAvailable() bool { + return a.State == AvailabilityAvailable +} diff --git a/pkg/api/registry/model.go b/pkg/api/registry/model.go index 7df4be45..97880e10 100644 --- a/pkg/api/registry/model.go +++ b/pkg/api/registry/model.go @@ -70,6 +70,8 @@ type Model struct { 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"` + // CallerTools reports that the runtime can expose caller-supplied tools. + CallerTools bool `json:"callerTools,omitempty" yaml:"callerTools,omitempty" jsonschema:"readOnly" pretty:"label=Caller Tools"` // 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. @@ -95,6 +97,7 @@ func (m Model) Capabilities() Model { m.Resume = caps.Resume m.Interrupt = caps.Interrupt m.Steer = caps.Steer + m.CallerTools = caps.CallerTools m.MediaTypes = p.MediaTypesFor(mode, m.Name) return m } diff --git a/pkg/api/registry/provider.go b/pkg/api/registry/provider.go index ff9515ed..69376eab 100644 --- a/pkg/api/registry/provider.go +++ b/pkg/api/registry/provider.go @@ -26,6 +26,9 @@ type ModeCapabilities struct { Interrupt bool // Steer: the adapter implements SteerableProvider. Steer bool + // CallerTools reports that the adapter can expose caller-supplied + // api.Config.Tools rather than only its built-in tool ecosystem. + CallerTools 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 diff --git a/pkg/api/registry/providers.go b/pkg/api/registry/providers.go index b2dce7ec..62b8aeba 100644 --- a/pkg/api/registry/providers.go +++ b/pkg/api/registry/providers.go @@ -19,9 +19,9 @@ var ( PricingPrefix: "anthropic", EnvVars: []string{"ANTHROPIC_API_KEY"}, modes: map[RuntimeMode]ModeCapabilities{ - ModeAPI: {Backend: BackendAnthropic, Streaming: true, MediaTypes: []string{"image/*"}}, + ModeAPI: {Backend: BackendAnthropic, Streaming: true, CallerTools: 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"}}, + ModeAgent: {Backend: BackendClaudeAgent, Streaming: true, Resume: true, Interrupt: true, Steer: true, CallerTools: 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{ @@ -43,9 +43,9 @@ var ( PricingPrefix: "openai", EnvVars: []string{"OPENAI_API_KEY"}, modes: map[RuntimeMode]ModeCapabilities{ - ModeAPI: {Backend: BackendOpenAI, Streaming: true, MediaTypes: []string{"image/*"}}, + ModeAPI: {Backend: BackendOpenAI, Streaming: true, CallerTools: 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/*"}}, + ModeAgent: {Backend: BackendCodexAgent, Streaming: true, Resume: true, Interrupt: true, CallerTools: 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" @@ -77,7 +77,7 @@ var ( 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"}}, + ModeAPI: {Backend: BackendGemini, Streaming: true, CallerTools: true, MediaTypes: []string{"image/*", "audio/*", "video/*", "application/pdf"}}, ModeCLI: {Backend: BackendGeminiCLI, Streaming: true}, }, modeTokens: sortModeTokens([]modeToken{ @@ -98,7 +98,7 @@ var ( 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}, + ModeAPI: {Backend: BackendDeepSeek, Streaming: true, CallerTools: true}, }, claimPrefixes: []string{"deepseek"}, families: []string{"deepseek"}, diff --git a/pkg/api/response_set.go b/pkg/api/response_set.go new file mode 100644 index 00000000..fa76865d --- /dev/null +++ b/pkg/api/response_set.go @@ -0,0 +1,36 @@ +package api + +// ResponseSet tracks which provider responses have already contributed usage, +// so a transcript can be accumulated without double-counting. +// +// Agent transcript formats commonly write one record per content block — a +// thinking block, a text block and a tool call each get their own line — and +// repeat the whole usage object on every one of them. Accumulating per record +// therefore counts a single response two or three times: a claude session that +// actually read 3.1M cached tokens reports 4.9M, and its cost inflates in step. +// +// This is the FALLBACK for sources that carry no result record. Where the +// provider reports a result — an invocation summary such as claude's +// stream-json `result` line, or a running total such as codex's +// total_token_usage — read that instead. A reported total is exact; a total +// reconstructed from per-record accumulation can only approach it. +type ResponseSet struct { + seen map[string]bool +} + +// First reports whether a response id has not yet contributed, and marks it as +// having done so. An empty id cannot be correlated to a response, so it counts +// once rather than being dropped. +func (r *ResponseSet) First(id string) bool { + if id == "" { + return true + } + if r.seen[id] { + return false + } + if r.seen == nil { + r.seen = map[string]bool{} + } + r.seen[id] = true + return true +} diff --git a/pkg/api/runtime_catalog.go b/pkg/api/runtime_catalog.go index df3b6ec4..40e619c0 100644 --- a/pkg/api/runtime_catalog.go +++ b/pkg/api/runtime_catalog.go @@ -45,20 +45,26 @@ type RuntimeModeEntry struct { // is what lets a client stop shipping its own "claude-sonnet-5" literal, // which went stale on every model release. DefaultModel string `json:"defaultModel,omitempty"` + // CatalogProvider is the provider key used by /api/chat/models for this + // mode. Local Claude/Codex modes share their agent catalog. It is always + // served: a client joining the model list to a mode falls back to the + // family's CatalogPrefix when it is absent, which puts every local Claude + // mode on "anthropic" and hands the Agent picker the Anthropic API rows. + CatalogProvider string `json:"catalogProvider"` // Disabled reports the user's opt-out. The entry is still served so a UI can - // explain the absence instead of silently shrinking; every picker drops it. + // explain the absence instead of silently shrinking the picker. Disabled bool `json:"disabled"` // DisabledReason names the switch that turned it off — "mode cmux", // "provider deepseek", "backend claude-agent" — or "" when enabled. - DisabledReason string `json:"disabledReason,omitempty"` + DisabledReason string `json:"disabledReason,omitempty"` + Availability Availability `json:"availability"` } // RuntimeCatalog projects the provider registry into the picker descriptor, // annotated with the installed opt-out set. // -// Annotated, not filtered: whoami renders the disabled entries as switched-off -// rows, and every other consumer drops them. Filtering here would take that -// choice away from the one surface that needs to show them. +// Annotated, not filtered: whoami and runtime pickers render disabled entries +// with an explanation. Filtering here would make the missing choice opaque. func RuntimeCatalog() []RuntimeFamily { disabled := Disabled() out := make([]RuntimeFamily, 0, len(registry.Providers())) @@ -74,14 +80,25 @@ func RuntimeCatalog() []RuntimeFamily { if !ok { continue } + reason := disabled.Reason(caps.Backend) + availability := Available() + if reason != "" { + availability = Availability{ + State: AvailabilityDisabled, + Reason: "Disabled by " + reason + " in Captain configuration.", + Remediation: "Enable " + reason + " on the Whoami page, then refresh.", + } + } family.Modes = append(family.Modes, RuntimeModeEntry{ - Mode: string(mode), - Backend: string(caps.Backend), - Kind: caps.Backend.Kind(), - Keyless: caps.Keyless, - DefaultModel: DefaultModelFor(caps.Backend), - Disabled: disabled.Backend(caps.Backend), - DisabledReason: disabled.Reason(caps.Backend), + Mode: string(mode), + Backend: string(caps.Backend), + Kind: caps.Backend.Kind(), + Keyless: caps.Keyless, + DefaultModel: DefaultModelFor(caps.Backend), + CatalogProvider: CatalogProviderFor(caps.Backend), + Disabled: disabled.Backend(caps.Backend), + DisabledReason: reason, + Availability: availability, }) } out = append(out, family) @@ -92,6 +109,30 @@ func RuntimeCatalog() []RuntimeFamily { // DefaultModelFor is the model a picker should seed for one backend. func DefaultModelFor(b Backend) string { return registry.DefaultModelFor(b) } +// CatalogProviderFor is the `provider` value /api/chat/models stamps on the +// menu row that serves this backend — the key a picker filters the flat model +// list by once the user has chosen a family and a mode. +// +// It is not CatalogPrefixFor: the menu lists one row per model per *menu* +// backend, and the local modes of a family collapse onto its agent backend +// (claude-cli and claude-cmux models are listed as claude-agent rows). A family +// with no agent mode has nothing to collapse onto — Gemini's CLI models already +// appear under the googleai API rows — so it keeps the catalog prefix. +func CatalogProviderFor(b Backend) string { + p, mode, ok := registry.ProviderFor(b) + if !ok { + return "" + } + if mode == registry.ModeAPI { + return p.CatalogPrefix + } + agent, err := p.BackendFor(registry.ModeAgent) + if err != nil { + return p.CatalogPrefix + } + return string(agent) +} + // CatalogPrefixFor is the namespace a backend's model ids live under: // "anthropic" for every Claude mode, "googleai" — not "google" — for Gemini. // diff --git a/pkg/api/runtime_catalog_ginkgo_test.go b/pkg/api/runtime_catalog_ginkgo_test.go index 3634a1ab..74582f9f 100644 --- a/pkg/api/runtime_catalog_ginkgo_test.go +++ b/pkg/api/runtime_catalog_ginkgo_test.go @@ -78,6 +78,36 @@ var _ = Describe("RuntimeCatalog", func() { Expect(kinds).To(Equal([]string{"api", "cli", "cli", "cli"})) }) + It("names the model-menu catalog provider of every mode", func() { + // The menu serves one row per model per *menu* backend: the three local + // Claude modes collapse onto claude-agent, while the API mode keeps the + // provider namespace. Leaving these empty made a picker fall back to the + // family's CatalogPrefix, so "Claude Agent" listed the Anthropic API rows + // and clicking one switched the backend to anthropic behind the user. + claude := familyNamed(api.RuntimeCatalog(), "claude") + Expect(modeNamed(claude, "api").CatalogProvider).To(Equal("anthropic")) + Expect(modeNamed(claude, "agent").CatalogProvider).To(Equal("claude-agent")) + Expect(modeNamed(claude, "cli").CatalogProvider).To(Equal("claude-agent")) + Expect(modeNamed(claude, "cmux").CatalogProvider).To(Equal("claude-agent")) + + codex := familyNamed(api.RuntimeCatalog(), "codex") + Expect(modeNamed(codex, "api").CatalogProvider).To(Equal("openai")) + Expect(modeNamed(codex, "agent").CatalogProvider).To(Equal("codex-agent")) + Expect(modeNamed(codex, "cli").CatalogProvider).To(Equal("codex-agent")) + Expect(modeNamed(codex, "cmux").CatalogProvider).To(Equal("codex-agent")) + }) + + It("keeps a family with no agent mode on its catalog prefix", func() { + // Gemini's CLI models are already listed under the googleai API rows, so + // there is no separate agent catalog for them to collapse onto. + gemini := familyNamed(api.RuntimeCatalog(), "gemini") + Expect(modeNamed(gemini, "api").CatalogProvider).To(Equal("googleai")) + Expect(modeNamed(gemini, "cli").CatalogProvider).To(Equal("googleai")) + + deepseek := familyNamed(api.RuntimeCatalog(), "deepseek") + Expect(modeNamed(deepseek, "api").CatalogProvider).To(Equal("deepseek")) + }) + It("marks only cmux modes keyless", func() { for _, f := range api.RuntimeCatalog() { for _, m := range f.Modes { @@ -91,6 +121,7 @@ var _ = Describe("RuntimeCatalog", func() { for _, m := range f.Modes { Expect(m.Disabled).To(BeFalse()) Expect(m.DisabledReason).To(BeEmpty()) + Expect(m.Availability).To(Equal(api.Available())) } } }) @@ -104,6 +135,9 @@ var _ = Describe("RuntimeCatalog", func() { cmux := modeNamed(familyNamed(families, "claude"), "cmux") Expect(cmux.Disabled).To(BeTrue()) Expect(cmux.DisabledReason).To(Equal("mode cmux")) + Expect(cmux.Availability.State).To(Equal(api.AvailabilityDisabled)) + Expect(cmux.Availability.Reason).To(ContainSubstring("mode cmux")) + Expect(cmux.Availability.Remediation).NotTo(BeEmpty()) Expect(modeNamed(familyNamed(families, "claude"), "api").Disabled).To(BeFalse()) }) diff --git a/pkg/api/runtime_config.go b/pkg/api/runtime_config.go index c5bfde1d..e63fc7c0 100644 --- a/pkg/api/runtime_config.go +++ b/pkg/api/runtime_config.go @@ -2,6 +2,10 @@ package api import ( "context" + "fmt" + "net" + "net/url" + "strings" "time" ) @@ -13,10 +17,11 @@ type PermissionFunc func(ctx context.Context, req PermissionRequest) (Permission // PermissionRequest describes the tool an agent wants to run. SessionID is filled // in by the provider from the live session so a caller can key approvals by it. type PermissionRequest struct { - Tool string - Input map[string]any - ToolUseID string - SessionID string + Tool string + Input map[string]any + ToolUseID string + ToolUseIDGenerated bool + SessionID string } // PermissionDecision is the answer to a PermissionRequest. On Allow the tool runs @@ -36,6 +41,59 @@ type SchemaRepairConfig struct { Prompt string // optional .prompt file path; empty means embedded default } +// CallerToolEndpoint is an authenticated, request-scoped MCP endpoint exposing +// caller-owned tools. Headers are transport credentials and must never be +// serialized into specs, command arguments, events, or logs. +type CallerToolEndpoint struct { + Name string + URL string + Headers map[string]string +} + +func (endpoint CallerToolEndpoint) Validate() error { + if endpoint.Name == "" { + return fmt.Errorf("caller-tool endpoint name is required") + } + for _, value := range endpoint.Name { + if !isCallerToolNameRune(value) { + return fmt.Errorf("caller-tool endpoint name %q contains unsupported characters", endpoint.Name) + } + } + parsed, err := url.Parse(endpoint.URL) + if err != nil || parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return fmt.Errorf("caller-tool endpoint URL must be an absolute HTTP or HTTPS URL") + } + if parsed.User != nil { + return fmt.Errorf("caller-tool endpoint URL must not contain credentials") + } + if parsed.Scheme == "http" && !isLoopbackHost(parsed.Hostname()) { + return fmt.Errorf("caller-tool endpoint requires HTTPS outside loopback") + } + authorization := "" + for name, value := range endpoint.Headers { + if strings.EqualFold(name, "Authorization") { + authorization = strings.TrimSpace(value) + break + } + } + if !strings.HasPrefix(authorization, "Bearer ") || strings.TrimSpace(strings.TrimPrefix(authorization, "Bearer ")) == "" { + return fmt.Errorf("caller-tool endpoint requires a bearer credential") + } + return nil +} + +func isCallerToolNameRune(value rune) bool { + return value >= 'a' && value <= 'z' || + value >= 'A' && value <= 'Z' || + value >= '0' && value <= '9' || + value == '-' || + value == '_' +} + +func isLoopbackHost(host string) bool { + return strings.EqualFold(host, "localhost") || net.ParseIP(host).IsLoopback() +} + // 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 @@ -48,8 +106,9 @@ type Config struct { // APIURL overrides the backend's endpoint (empty = the provider default). // Anthropic/OpenAI/DeepSeek honour it; Gemini rejects it, because genkit's // googlegenai plugin exposes no override and silently calling the real API - // would be worse. codex-cli honours it by declaring a model_providers entry, - // since it ignores OPENAI_BASE_URL once account auth is stored. + // would be worse. Claude Agent passes it through to the SDK child; Codex CLI + // and Codex Agent declare a model_providers entry because stored account auth + // otherwise takes precedence over OPENAI_BASE_URL. APIURL string // Sandbox runs local agent CLI processes through sandbox-runtime. Provider // selection must resolve to CLI mode so the flag cannot be silently ignored. @@ -64,6 +123,11 @@ type Config struct { NoCache bool MaxConcurrent int SessionID string + // CaptainSessionID is Captain's own session/thread UUID, as distinct from + // SessionID (the provider's id for the same conversation). Caller-tool MCP + // endpoints are scoped by it so an approval brokered for one Captain thread + // cannot be replayed against another that happens to share a provider id. + CaptainSessionID string ProjectName string SchemaRepair SchemaRepairConfig @@ -76,10 +140,16 @@ type Config struct { 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). + // in-process. Tool-capable API providers invoke the handlers directly; + // out-of-process agent providers expose them through a private Captain MCP + // endpoint. Never serialized (Go closures). Tools []ToolDefinition `json:"-"` + + // CallerTools supplies a pre-issued Captain MCP endpoint. When nil, an + // out-of-process tool-capable provider creates a private loopback endpoint + // from Tools. It is runtime-only because Headers contain a short-lived + // credential. + CallerTools *CallerToolEndpoint `json:"-"` } // ResolvedSandbox returns the sandbox selection for the run, folding the legacy diff --git a/pkg/api/runtime_config_ginkgo_test.go b/pkg/api/runtime_config_ginkgo_test.go new file mode 100644 index 00000000..ee0961bc --- /dev/null +++ b/pkg/api/runtime_config_ginkgo_test.go @@ -0,0 +1,39 @@ +package api_test + +import ( + "github.com/flanksource/captain/pkg/api" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Caller-tool endpoints", func() { + It("accepts authenticated loopback HTTP and remote HTTPS endpoints", func() { + for _, endpoint := range []api.CallerToolEndpoint{ + { + Name: "captain", URL: "http://127.0.0.1:43210/mcp", + Headers: map[string]string{"Authorization": "Bearer loopback-secret"}, + }, + { + Name: "captain-remote", URL: "https://captain.example.com/mcp", + Headers: map[string]string{"authorization": "Bearer remote-secret"}, + }, + } { + Expect(endpoint.Validate()).To(Succeed()) + } + }) + + It("rejects invalid names, unauthenticated endpoints, and remote plaintext HTTP", func() { + Expect((api.CallerToolEndpoint{ + Name: "captain tools", URL: "https://captain.example.com/mcp", + Headers: map[string]string{"Authorization": "Bearer secret"}, + }).Validate()).To(MatchError(ContainSubstring("name"))) + Expect((api.CallerToolEndpoint{ + Name: "captain", URL: "https://captain.example.com/mcp", + }).Validate()).To(MatchError(ContainSubstring("bearer"))) + Expect((api.CallerToolEndpoint{ + Name: "captain", URL: "http://captain.example.com/mcp", + Headers: map[string]string{"Authorization": "Bearer secret"}, + }).Validate()).To(MatchError(ContainSubstring("HTTPS"))) + }) +}) diff --git a/pkg/api/runtime_event.go b/pkg/api/runtime_event.go index f54c49ac..e3a29b23 100644 --- a/pkg/api/runtime_event.go +++ b/pkg/api/runtime_event.go @@ -33,13 +33,14 @@ type Response struct { type EventKind string const ( - EventText EventKind = "text" - EventThinking EventKind = "thinking" - EventToolUse EventKind = "tool_use" - EventToolResult EventKind = "tool_result" - EventResult EventKind = "result" - EventError EventKind = "error" - EventSystem EventKind = "system" + EventText EventKind = "text" + EventThinking EventKind = "thinking" + EventToolUse EventKind = "tool_use" + EventToolResult EventKind = "tool_result" + EventResult EventKind = "result" + EventError EventKind = "error" + EventInterrupted EventKind = "interrupted" + EventSystem EventKind = "system" // EventPermission surfaces a tool-permission request brokered via CanUseTool // so callers can observe what is awaiting approval. Tool/Input/ToolCallID carry // the requested tool; the decision itself flows back through the CanUseTool @@ -59,6 +60,9 @@ type Event struct { // (the call) and EventToolResult (its complete output). Backends that stream // output incrementally accumulate it and emit a single EventToolResult. ToolCallID string + // ApprovalID is the durable captain_turn_requests UUID associated with an + // EventPermission. It is distinct from the provider's tool-call ID. + ApprovalID string Usage *Usage // when Kind == EventResult CostUSD float64 // when Kind == EventResult @@ -66,6 +70,7 @@ type Event struct { SessionID string // when Kind == EventSystem Model string Error string // when Kind == EventError + Reason string // when Kind == EventInterrupted // StructuredData is the validated structured output (raw JSON) carried on an // EventResult when the request supplied a schema; nil for text-mode runs. It diff --git a/pkg/api/spec_merge_differential_test.go b/pkg/api/spec_merge_differential_test.go index f735523f..a8284a18 100644 --- a/pkg/api/spec_merge_differential_test.go +++ b/pkg/api/spec_merge_differential_test.go @@ -199,7 +199,7 @@ func neutralize(s Spec) Spec { // Fields the hand-written mergers simply forgot; the structural engine cannot // forget one, so these now carry through. s.Prompt.Attachments = nil - s.Model.Streaming, s.Model.Resume, s.Model.Interrupt, s.Model.Steer = false, false, false, false + s.Model.Streaming, s.Model.Resume, s.Model.Interrupt, s.Model.Steer, s.Model.CallerTools = false, false, false, false, false s.Model.MediaTypes = nil s.Model.Provider = nil return s diff --git a/pkg/api/tool_approval.go b/pkg/api/tool_approval.go index 9b9a383c..27262e97 100644 --- a/pkg/api/tool_approval.go +++ b/pkg/api/tool_approval.go @@ -32,17 +32,28 @@ type ToolApprovalCall struct { Result *ToolResult `json:"result,omitempty" yaml:"result,omitempty"` } +// ProviderCheckpoint is opaque provider-native conversation state. It is +// persisted beside the prompt run and is deliberately excluded from every +// public transcript and session response. +type ProviderCheckpoint struct { + Codec string + Version int + Payload []byte +} + // 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"` + Messages []Message `json:"messages" yaml:"messages"` + Calls []ToolApprovalCall `json:"calls" yaml:"calls"` + ProviderCheckpoint *ProviderCheckpoint `json:"-" yaml:"-"` } // ToolApprovalDecision resolves one pending call. Approve may replace Input; // Deny may carry a Message; Respond supplies an already-computed Result. type ToolApprovalDecision struct { + ApprovalID string `json:"approvalId,omitempty" yaml:"approvalId,omitempty"` ToolCallID string `json:"toolCallId" yaml:"toolCallId"` Tool string `json:"tool" yaml:"tool"` Action ToolApprovalAction `json:"action" yaml:"action"` @@ -86,6 +97,11 @@ func (s ToolApprovalState) Pending() []ToolApprovalRequest { } func (s ToolApprovalState) Validate() error { + if s.ProviderCheckpoint != nil { + if strings.TrimSpace(s.ProviderCheckpoint.Codec) == "" || s.ProviderCheckpoint.Version <= 0 || len(s.ProviderCheckpoint.Payload) == 0 { + return fmt.Errorf("provider checkpoint requires a codec, positive version, and payload") + } + } if err := ValidateMessages(s.Messages); err != nil { return fmt.Errorf("approval messages: %w", err) } diff --git a/pkg/api/workspace.go b/pkg/api/workspace.go index c2ed4426..2621769b 100644 --- a/pkg/api/workspace.go +++ b/pkg/api/workspace.go @@ -1,5 +1,7 @@ package api +import "time" + // Workspace is the runtime state of a run's working directory — the output // counterpart to Spec.Setup (the input checkout/worktree config). It records // where the run executed, the git details, and what it changed / committed / @@ -13,6 +15,7 @@ type Workspace struct { Base string `json:"base,omitempty" yaml:"base,omitempty"` // worktree base ref Changed []string `json:"changed,omitempty" yaml:"changed,omitempty"` // agent-changed files (repo-relative) Commits []CommitRecord `json:"commits,omitempty" yaml:"commits,omitempty"` // commits made during the run + Notices []Notice `json:"notices,omitempty" yaml:"notices,omitempty"` // lifecycle lines hooks reported Diff string `json:"diff,omitempty" yaml:"diff,omitempty"` // working diff Plan string `json:"plan,omitempty" yaml:"plan,omitempty"` // plan the agent produced (path or content) SessionID string `json:"sessionId,omitempty" yaml:"sessionId,omitempty"` // agent session @@ -33,3 +36,24 @@ func (w *Workspace) AddCommit(sha, message string) { } w.Commits = append(w.Commits, CommitRecord{SHA: sha, Message: message}) } + +// Notice is one thing a lifecycle hook did, reported in the run's own voice — +// "committed abc1234", "nothing to stage". Hooks act between the model's turns, +// where the provider transcript has nothing to say, so without these a run's +// commits, pushes and teardowns are invisible to anyone reading it back. +// +// At is the moment it happened, which is what lets a notice be sorted back into +// its place among the turns it sits between rather than clumping at the end. +type Notice struct { + At time.Time `json:"at" yaml:"at"` + Phase string `json:"phase,omitempty" yaml:"phase,omitempty"` + Text string `json:"text" yaml:"text"` +} + +// AddNotice appends a notice; nil-safe convenience for hooks. +func (w *Workspace) AddNotice(at time.Time, phase, text string) { + if w == nil { + return + } + w.Notices = append(w.Notices, Notice{At: at, Phase: phase, Text: text}) +} diff --git a/pkg/bash/shell_transform.go b/pkg/bash/shell_transform.go new file mode 100644 index 00000000..01264161 --- /dev/null +++ b/pkg/bash/shell_transform.go @@ -0,0 +1,126 @@ +package bash + +import ( + "maps" + "path/filepath" + "strings" + + "mvdan.cc/sh/v3/syntax" +) + +type ShellCommand struct { + Command string + Shell string + Flags []string + Args []string +} + +func TransformShellCommand(command string) (ShellCommand, bool) { + file, err := syntax.NewParser().Parse(strings.NewReader(command), "") + if err != nil || len(file.Stmts) != 1 || len(file.Stmts[0].Redirs) > 0 { + return ShellCommand{}, false + } + call, ok := file.Stmts[0].Cmd.(*syntax.CallExpr) + if !ok || len(call.Assigns) > 0 || len(call.Args) < 3 { + return ShellCommand{}, false + } + args, ok := staticWords(call.Args) + if !ok { + return ShellCommand{}, false + } + shell := filepath.Base(args[0]) + if shell != "sh" && shell != "bash" && shell != "zsh" { + return ShellCommand{}, false + } + + flags, commandIndex, ok := shellCommandFlag(args[1:]) + if !ok || commandIndex+2 >= len(args) { + return ShellCommand{}, false + } + return ShellCommand{ + Command: args[commandIndex+2], + Shell: shell, + Flags: flags, + Args: append([]string(nil), args[commandIndex+3:]...), + }, true +} + +func TransformBashInput(input map[string]any) map[string]any { + if input == nil { + return nil + } + transformed := maps.Clone(input) + if shell, _ := transformed["shell"].(string); shell != "" { + return transformed + } + command, _ := transformed["command"].(string) + wrapped, ok := TransformShellCommand(command) + if !ok { + return transformed + } + transformed["command"] = wrapped.Command + transformed["shell"] = wrapped.Shell + if len(wrapped.Flags) > 0 { + transformed["shellFlags"] = wrapped.Flags + } + if len(wrapped.Args) > 0 { + transformed["shellArgs"] = wrapped.Args + } + return transformed +} + +func staticWords(words []*syntax.Word) ([]string, bool) { + values := make([]string, len(words)) + for i, word := range words { + if !isStaticWord(word) { + return nil, false + } + values[i] = wordToString(word) + } + return values, true +} + +func isStaticWord(word *syntax.Word) bool { + if word == nil { + return false + } + for _, part := range word.Parts { + switch value := part.(type) { + case *syntax.Lit, *syntax.SglQuoted: + case *syntax.DblQuoted: + if !isStaticWord(&syntax.Word{Parts: value.Parts}) { + return false + } + default: + return false + } + } + return true +} + +func shellCommandFlag(args []string) ([]string, int, bool) { + var flags []string + for i, arg := range args { + if arg == "--" || !strings.HasPrefix(arg, "-") || arg == "-" { + return nil, 0, false + } + if strings.HasPrefix(arg, "--") { + flags = append(flags, arg) + continue + } + options := strings.TrimPrefix(arg, "-") + commandOption := strings.IndexByte(options, 'c') + if commandOption < 0 { + flags = append(flags, arg) + continue + } + if commandOption != len(options)-1 { + return nil, 0, false + } + if remaining := options[:commandOption]; remaining != "" { + flags = append(flags, "-"+remaining) + } + return flags, i, true + } + return nil, 0, false +} diff --git a/pkg/bash/shell_transform_ginkgo_test.go b/pkg/bash/shell_transform_ginkgo_test.go new file mode 100644 index 00000000..71b40d7f --- /dev/null +++ b/pkg/bash/shell_transform_ginkgo_test.go @@ -0,0 +1,56 @@ +package bash + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestShellTransform(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Shell Transform Suite") +} + +var _ = Describe("TransformShellCommand", func() { + It("unwraps a login zsh command without losing shell metadata", func() { + transformed, ok := TransformShellCommand(`/bin/zsh -lc 'gavel pr status 50 --logs'`) + + Expect(ok).To(BeTrue()) + Expect(transformed).To(Equal(ShellCommand{ + Command: "gavel pr status 50 --logs", + Shell: "zsh", + Flags: []string{"-l"}, + })) + }) + + It("retains positional arguments used by the command body", func() { + transformed, ok := TransformShellCommand(`/bin/bash -c 'printf "%s" "$1"' command-name value`) + + Expect(ok).To(BeTrue()) + Expect(transformed.Command).To(Equal(`printf "%s" "$1"`)) + Expect(transformed.Shell).To(Equal("bash")) + Expect(transformed.Args).To(Equal([]string{"command-name", "value"})) + }) + + It("does not transform a dynamically expanded wrapper", func() { + _, ok := TransformShellCommand(`/bin/zsh -lc "echo $HOME"`) + Expect(ok).To(BeFalse()) + }) + + It("normalizes Bash input idempotently", func() { + input := map[string]any{"command": `/bin/zsh -lc 'pnpm test'`, "timeout": float64(1000)} + + first := TransformBashInput(input) + second := TransformBashInput(first) + + Expect(first).To(Equal(map[string]any{ + "command": "pnpm test", + "shell": "zsh", + "shellFlags": []string{"-l"}, + "timeout": float64(1000), + })) + Expect(second).To(Equal(first)) + Expect(input["command"]).To(Equal(`/bin/zsh -lc 'pnpm test'`)) + }) +}) diff --git a/pkg/claude/cost.go b/pkg/claude/cost.go index fc976b47..d86500b9 100644 --- a/pkg/claude/cost.go +++ b/pkg/claude/cost.go @@ -58,8 +58,17 @@ type TokenSummary struct { CacheWriteTokens int `json:"cacheWriteTokens" pretty:"label=Cache Write"` CacheReadTokens int `json:"cacheReadTokens" pretty:"label=Cache Read"` TotalCost float64 `json:"totalCost"` + // ProviderCostUSD is the model provider's own billed total, when one was + // recorded. Zero means TotalCost is a list-price reconstruction from token + // counts rather than a figure the provider reported — the two must not be + // presented as the same thing. + ProviderCostUSD float64 `json:"providerCostUsd,omitempty"` } +// Estimated reports whether TotalCost was recomputed from token counts rather +// than taken from a result the provider reported. +func (s TokenSummary) Estimated() bool { return s.ProviderCostUSD == 0 } + func (s *TokenSummary) Add(usage *Usage, model string) { if usage == nil { return diff --git a/pkg/claude/history.go b/pkg/claude/history.go index 5b6229e7..3327b588 100644 --- a/pkg/claude/history.go +++ b/pkg/claude/history.go @@ -44,6 +44,11 @@ type TranscriptEvent struct { // Message represents a conversation message type Message struct { + // ID is the provider's message id (e.g. "msg_011Cdhc1..."), identifying one + // API response. Claude Code writes a separate transcript line per content + // block, all sharing this id and repeating the same Usage, so consumers that + // aggregate usage must deduplicate on it. + ID string `json:"id,omitempty"` Model string `json:"model,omitempty"` Role MessageRole `json:"role"` Content []ContentBlock `json:"-"` @@ -54,6 +59,7 @@ type Message struct { // UnmarshalJSON handles polymorphic content field (string, array, or null) func (m *Message) UnmarshalJSON(data []byte) error { type messageAlias struct { + ID string `json:"id,omitempty"` Model string `json:"model,omitempty"` Role MessageRole `json:"role"` Content json.RawMessage `json:"content"` @@ -66,6 +72,7 @@ func (m *Message) UnmarshalJSON(data []byte) error { return err } + m.ID = alias.ID m.Model = alias.Model m.Role = alias.Role m.StopReason = alias.StopReason diff --git a/pkg/claude/session.go b/pkg/claude/session.go index 4aad22a4..691f0826 100644 --- a/pkg/claude/session.go +++ b/pkg/claude/session.go @@ -7,6 +7,7 @@ import ( "strings" "time" + "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/claude/tools" "github.com/segmentio/encoding/json" ) @@ -490,6 +491,20 @@ type SessionCost struct { Files []string `json:"files,omitempty"` Context *ContextBreakdown `json:"context,omitempty"` ToolCosts []ToolTokenSummary `json:"toolCosts,omitempty"` + + // responses deduplicates the per-content-block lines a single API response + // is written across. This is the no-result fallback — claude transcripts + // carry no result record to read instead. See api.ResponseSet. + responses api.ResponseSet +} + +// addUsage accumulates one assistant line's tokens, skipping repeated +// content-block lines of a response already counted. +func (sc *SessionCost) addUsage(entry HistoryEntry) { + if !sc.responses.First(entry.Message.ID) { + return + } + sc.Tokens.Add(entry.Message.Usage, entry.Message.Model) } func ParseCosts(currentDir string, searchAll bool, since *time.Time) ([]SessionCost, error) { @@ -581,7 +596,7 @@ func costsFromEntries(sessionFile string, entries []HistoryEntry, projectRoot st sc.Tier = tier } - sc.Tokens.Add(entry.Message.Usage, model) + sc.addUsage(entry) sc.Messages++ } @@ -682,7 +697,7 @@ func ParseCostsDetailedWithFilter(currentDir string, searchAll bool, since *time sc.Tier = tier } - sc.Tokens.Add(entry.Message.Usage, model) + sc.addUsage(entry) sc.Messages++ } } diff --git a/pkg/claude/shell_transform_ginkgo_test.go b/pkg/claude/shell_transform_ginkgo_test.go new file mode 100644 index 00000000..1bc6f921 --- /dev/null +++ b/pkg/claude/shell_transform_ginkgo_test.go @@ -0,0 +1,30 @@ +package claude + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/segmentio/encoding/json" +) + +var _ = Describe("Claude Bash normalization", func() { + It("transforms shell wrappers while extracting history tools", func() { + uses := ExtractToolUses([]HistoryEntry{{ + Message: Message{ + Role: MessageRoleAssistant, + Content: []ContentBlock{{ + Type: ContentTypeToolUse, + ID: "tool-1", + Name: "Bash", + Input: json.RawMessage(`{"command":"/bin/zsh -lc 'pnpm test'"}`), + }}, + }, + }}) + + Expect(uses).To(HaveLen(1)) + Expect(uses[0].Input).To(Equal(map[string]any{ + "command": "pnpm test", + "shell": "zsh", + "shellFlags": []string{"-l"}, + })) + }) +}) diff --git a/pkg/claude/tools/bash.go b/pkg/claude/tools/bash.go index d05f75ca..a6e38279 100644 --- a/pkg/claude/tools/bash.go +++ b/pkg/claude/tools/bash.go @@ -26,7 +26,11 @@ func (t *BashTool) Category() string { return "" } func (t *BashTool) Pretty() api.Text { cmd := t.command() color := "text-green-400 font-medium" - text := t.header(bashIcon, strings.ToLower(t.Name()), color) + label := strings.ToLower(t.Name()) + if shell := t.Str("shell"); shell != "" { + label = shell + } + text := t.header(bashIcon, label, color) if timeout := t.Float("timeout"); timeout > 0 { text = text.Append(fmt.Sprintf(" (%ds)", int(timeout/1000)), "text-gray-500") @@ -87,6 +91,9 @@ func (t *BashTool) command() string { } func (t *BashTool) interpreter() string { + if t.Str("shell") != "" { + return "" + } cmd := t.Str("command") if cmd == "" { return "" diff --git a/pkg/claude/tools/bash_shell_ginkgo_test.go b/pkg/claude/tools/bash_shell_ginkgo_test.go new file mode 100644 index 00000000..67f3d1b7 --- /dev/null +++ b/pkg/claude/tools/bash_shell_ginkgo_test.go @@ -0,0 +1,23 @@ +package tools + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Bash shell metadata", func() { + It("renders the transformed shell and actual command in every format", func() { + tool := &BashTool{BaseTool: BaseTool{Input: map[string]any{ + "command": `python -c 'print(42)'`, + "shell": "zsh", + }}} + + Expect(tool.Name()).To(Equal("Bash")) + Expect(tool.Pretty().String()).To(And( + ContainSubstring("zsh"), + ContainSubstring(`python -c 'print(42)'`), + Not(ContainSubstring("/bin/zsh -lc")), + )) + Expect(tool.Detail().Markdown()).To(ContainSubstring(`python -c 'print(42)'`)) + }) +}) diff --git a/pkg/claude/tools/generic.go b/pkg/claude/tools/generic.go index 5f47e17b..a0a4d1d4 100644 --- a/pkg/claude/tools/generic.go +++ b/pkg/claude/tools/generic.go @@ -1,9 +1,11 @@ package tools import ( - "github.com/segmentio/encoding/json" + "bytes" "strings" + "github.com/segmentio/encoding/json" + "github.com/flanksource/clicky" "github.com/flanksource/clicky/api" "github.com/flanksource/clicky/api/icons" @@ -31,8 +33,8 @@ func (t *GenericTool) Pretty() api.Text { color = "text-blue-400 font-medium" } text := t.header(icon, strings.ToLower(t.RawTool), color) - if b, err := json.Marshal(t.Input); err == nil { - text = text.Append(" " + messagePreview(string(b))) + if preview := genericPreview(t.Input); preview != "" { + text = text.Append(" " + preview) } return text } @@ -44,9 +46,68 @@ func (t *GenericTool) Detail() api.Textable { } // The preview is a truncated one-liner; the full input belongs somewhere a // non-terminal format can still reach it. - b, err := json.Marshal(t.Input) - if err != nil || len(t.Input) == 0 { + if len(t.Input) == 0 { + return nil + } + b, err := encodeJSON(t.Input, true) + if err != nil { return nil } return api.NewCode(string(b), "json") } + +// genericPreview renders an unmapped tool's input as a one-line JSON preview. +// +// String values are whitespace-collapsed and HTML escaping is switched off +// before marshalling. Marshalling a raw input map encodes every newline as the +// two characters backslash-n and every "<" as backslash-u-0-0-3-c, so a value +// holding a multi-line body -- a system prompt, a heredoc -- renders as a wall +// of escape sequences that consumes the whole preview budget before any of the +// body is reached. +func genericPreview(input map[string]any) string { + if len(input) == 0 { + return "" + } + b, err := encodeJSON(compactStrings(input), false) + if err != nil { + return "" + } + return messagePreview(string(b)) +} + +// compactStrings collapses whitespace in every string the value tree holds. +func compactStrings(v any) any { + switch v := v.(type) { + case string: + return compactText(v) + case map[string]any: + out := make(map[string]any, len(v)) + for k, val := range v { + out[k] = compactStrings(val) + } + return out + case []any: + out := make([]any, len(v)) + for i, val := range v { + out[i] = compactStrings(val) + } + return out + default: + return v + } +} + +// encodeJSON marshals without the HTML escaping the package applies by +// default, so "<", ">" and "&" reach rendered output as themselves. +func encodeJSON(v any, indent bool) ([]byte, error) { + var buf bytes.Buffer + enc := json.NewEncoder(&buf) + enc.SetEscapeHTML(false) + if indent { + enc.SetIndent("", " ") + } + if err := enc.Encode(v); err != nil { + return nil, err + } + return bytes.TrimRight(buf.Bytes(), "\n"), nil +} diff --git a/pkg/claude/tools/muted_styles_test.go b/pkg/claude/tools/muted_styles_test.go index 4aeebc45..bbadef3f 100644 --- a/pkg/claude/tools/muted_styles_test.go +++ b/pkg/claude/tools/muted_styles_test.go @@ -29,6 +29,9 @@ var _ = Describe("semantic muted pretty styles", func() { Entry("reasoning body", func() api.Textable { return (&ReasoningTool{BaseTool: BaseTool{Input: map[string]any{"text": "reasoning payload"}}}).Pretty() }, " reasoning payload"), + Entry("system body", func() api.Textable { + return (&SystemTool{BaseTool: BaseTool{Input: map[string]any{"text": "system payload"}}}).Pretty() + }, " system payload"), Entry("session model", func() api.Textable { return (&SystemInitTool{BaseTool: BaseTool{Input: map[string]any{"model": "model-x"}}}).Pretty() }, " model-x"), diff --git a/pkg/claude/tools/system.go b/pkg/claude/tools/system.go new file mode 100644 index 00000000..40fcf507 --- /dev/null +++ b/pkg/claude/tools/system.go @@ -0,0 +1,37 @@ +package tools + +import ( + "github.com/flanksource/clicky/api" + "github.com/flanksource/clicky/api/icons" +) + +// SystemTool surfaces a system-role message -- the Codex session prompt, plugin +// and skill instruction blocks -- as a history row, the same way UserTool and +// AssistantTool surface the other two conversational roles. +// +// Without it the row falls through to GenericTool, which marshals the whole +// input map. A system prompt carries hundreds of newlines, so the preview +// budget is spent on JSON escape sequences and the reader never sees the +// prompt at all. +type SystemTool struct{ BaseTool } + +func (t *SystemTool) Name() string { return "System" } +func (t *SystemTool) Category() string { return "chat" } +func (t *SystemTool) FilePath() string { return "" } +func (t *SystemTool) ExtractPath() string { return "" } + +func (t *SystemTool) Pretty() api.Text { + icon := icons.Icon{Unicode: "⚙️", Iconify: "mdi:cog", Style: "muted"} + text := t.header(icon, "system", "text-slate-500 font-medium") + if body := t.Str("text"); body != "" { + text = text.Append(" "+messagePreview(body), "text-muted") + } + return text +} + +func (t *SystemTool) Detail() api.Textable { + if denied := t.BaseTool.Detail(); denied != nil { + return denied + } + return messageDetail(t.Str("text")) +} diff --git a/pkg/claude/tools/system_test.go b/pkg/claude/tools/system_test.go new file mode 100644 index 00000000..23b44388 --- /dev/null +++ b/pkg/claude/tools/system_test.go @@ -0,0 +1,76 @@ +package tools + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// systemPrompt is a miniature of a real Codex system row: a multi-line body +// holding an XML-ish tag, the two shapes that JSON marshalling mangles. +const systemPrompt = "\n## Plugins\nA plugin is a local bundle of skills.\n" + +// escapedNewline and escapedLT are what json.Marshal emits for a real newline +// and a real "<". Both are two- and six-character sequences in the rendered +// output, not control characters, which is exactly why they are unreadable. +const ( + escapedNewline = "\\n" + escapedLT = "\\u003c" +) + +func TestNewTool_SystemRoleDispatchesToSystemTool(t *testing.T) { + tool := NewTool(BaseTool{RawTool: "System", Input: map[string]any{"text": systemPrompt}}) + + require.IsType(t, &SystemTool{}, tool) + assert.Equal(t, "System", tool.Name()) +} + +func TestSystemTool_PrettyShowsBodyWithoutJSONEscapes(t *testing.T) { + pretty := (&SystemTool{BaseTool: BaseTool{ + RawTool: "System", + Input: map[string]any{"text": systemPrompt}, + }}).Pretty().String() + + assert.Contains(t, pretty, "system") + assert.Contains(t, pretty, " ## Plugins A plugin is a local bundle of skills.") + assert.NotContains(t, pretty, escapedNewline, "newlines must be collapsed, not JSON-escaped") + assert.NotContains(t, pretty, escapedLT, "angle brackets must not be HTML-escaped") + assert.NotContains(t, pretty, `{"text"`, "the row must show the body, not its JSON envelope") +} + +func TestSystemTool_DetailKeepsFullBody(t *testing.T) { + detail := (&SystemTool{BaseTool: BaseTool{ + RawTool: "System", + Input: map[string]any{"text": systemPrompt}, + }}).Detail() + + require.NotNil(t, detail) + assert.Contains(t, detail.String(), systemPrompt) +} + +func TestGenericTool_PreviewCollapsesNewlinesAndKeepsAngleBrackets(t *testing.T) { + tool := NewTool(BaseTool{RawTool: "write_stdin", Input: map[string]any{ + "chars": "echo \nexit\n", + "tags": []any{"a\nb"}, + }}) + require.IsType(t, &GenericTool{}, tool) + + pretty := tool.Pretty().String() + assert.Contains(t, pretty, "write_stdin") + assert.Contains(t, pretty, "echo exit") + assert.Contains(t, pretty, `"a b"`, "nested string values are collapsed too") + assert.NotContains(t, pretty, escapedNewline) + assert.NotContains(t, pretty, escapedLT) +} + +func TestGenericTool_DetailIsReadableJSON(t *testing.T) { + detail := (&GenericTool{BaseTool: BaseTool{ + RawTool: "write_stdin", + Input: map[string]any{"chars": "echo "}, + }}).Detail() + + require.NotNil(t, detail) + assert.Contains(t, detail.String(), "echo ") + assert.NotContains(t, detail.String(), escapedLT) +} diff --git a/pkg/claude/tools/tool.go b/pkg/claude/tools/tool.go index 496830dc..849da032 100644 --- a/pkg/claude/tools/tool.go +++ b/pkg/claude/tools/tool.go @@ -239,6 +239,8 @@ func NewTool(base BaseTool) Tool { return &PlanTool{BaseTool: base} case "User": return &UserTool{BaseTool: base} + case "System": + return &SystemTool{BaseTool: base} case "Assistant": return &AssistantTool{BaseTool: base} case "Reasoning": diff --git a/pkg/claude/tooluse.go b/pkg/claude/tooluse.go index 45e670d4..b1e52e51 100644 --- a/pkg/claude/tooluse.go +++ b/pkg/claude/tooluse.go @@ -161,6 +161,9 @@ func ExtractToolUses(entries []HistoryEntry) []ToolUse { if content.Input != nil { _ = json.Unmarshal(content.Input, &inputMap) } + if content.Name == "Bash" { + inputMap = bash.TransformBashInput(inputMap) + } var cwd string if inputMap != nil { diff --git a/pkg/cli/ai.go b/pkg/cli/ai.go index 6b5e56df..de6c8ed8 100644 --- a/pkg/cli/ai.go +++ b/pkg/cli/ai.go @@ -2,6 +2,7 @@ package cli import ( "context" + "errors" "fmt" "os" "strconv" @@ -16,8 +17,6 @@ import ( "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/api/registry" "github.com/flanksource/captain/pkg/captainconfig" - "github.com/flanksource/captain/pkg/claude" - "github.com/flanksource/captain/pkg/claude/tools" "github.com/flanksource/captain/pkg/collections" ) @@ -523,7 +522,7 @@ func runStreaming(ctx context.Context, sp ai.StreamingProvider, req ai.Request) structuredOutput map[string]any structuredErr error ) - renderer := newLineRenderer(os.Stderr, 8) + renderer := NewEventRenderer(os.Stderr) loop, err := ai.RunUntil(ctx, ai.LoopOptions{ Provider: sp, MaxIterations: 1, @@ -534,14 +533,14 @@ func runStreaming(ctx context.Context, sp ai.StreamingProvider, req ai.Request) } return req, true }, - OnEvent: func(_ int, ev ai.Event) { + OnEvent: func(iteration int, ev ai.Event) { if ev.Model != "" { model = ev.Model } if ev.SessionID != "" { session = ev.SessionID } - renderEvent(os.Stderr, renderer, ev) + renderer.Handle(iteration, ev) if ev.Kind == ai.EventText { text += ev.Text } @@ -557,6 +556,9 @@ func runStreaming(ctx context.Context, sp ai.StreamingProvider, req ai.Request) } }, }) + if renderErr := renderer.Flush(); renderErr != nil { + return nil, errors.Join(err, renderErr) + } if err != nil { return nil, err } @@ -615,148 +617,6 @@ func actualRunDir(req ai.Request) string { 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 -// `captain history` for the same tools (including session-start banners). -func renderEvent(w *os.File, renderer *lineRenderer, ev ai.Event) { - if entry, ok := ev.Raw.(claude.HistoryEntry); ok { - if renderClaudeEntry(renderer, ev, entry) { - return - } - } - if tu, ok := ev.Raw.(claude.ToolUse); ok { - if renderCodexEntry(renderer, ev, tu) { - return - } - } - - switch ev.Kind { - case ai.EventText: - fmt.Fprintf(w, "%s", ev.Text) - case ai.EventThinking: - if log.IsDebugEnabled() { - fmt.Fprintf(w, "[thinking] %s\n", truncForStderr(ev.Text, 200)) - } - case ai.EventToolUse: - fmt.Fprintf(w, "\n[tool] %s %s\n", ev.Tool, summariseInput(ev.Input)) - case ai.EventPermission: - fmt.Fprintf(w, "\n[permission] %s %s awaiting approval\n", ev.Tool, summariseInput(ev.Input)) - case ai.EventToolResult: - if ev.Text != "" { - label := "tool-result" - if !ev.Success { - label = "tool-error" - } - fmt.Fprintf(w, "[%s] %s\n", label, truncForStderr(ev.Text, 500)) - } - case ai.EventResult: - renderResultEvent(renderer, ev) - case ai.EventError: - fmt.Fprintf(w, "\n[error] %s\n", ev.Error) - log.Errorf("%s", ev.Error) - case ai.EventSystem: - if ev.SessionID != "" { - fmt.Fprintf(w, "[session] %s\n", ev.SessionID) - } - } -} - -// renderResultEvent synthesizes a Result tools.Tool from the ai.Event so -// streaming output renders end-of-session result lines with the same -// "🏁 result turns=N $X 1.2s" formatting as `captain history`. -func renderResultEvent(renderer *lineRenderer, ev ai.Event) { - input := map[string]any{} - for k, v := range ev.Input { - input[k] = v - } - if ev.CostUSD > 0 { - if _, ok := input["total_cost_usd"]; !ok { - input["total_cost_usd"] = ev.CostUSD - } - } - if !ev.Success { - input["is_error"] = true - if _, ok := input["result"]; !ok && ev.Error != "" { - input["result"] = ev.Error - } - } - base := tools.BaseTool{ - RawTool: "Result", - Input: input, - Timestamp: nil, - } - if ev.Usage != nil && (ev.Usage.InputTokens > 0 || ev.Usage.OutputTokens > 0) { - base.Models = tools.Models{{ - Model: ev.Model, - InputTokens: ev.Usage.InputTokens, - OutputTokens: ev.Usage.OutputTokens, - }} - } - renderer.Render(tools.NewTool(base), true) -} - -// renderClaudeEntry feeds a claude HistoryEntry through the shared lineRenderer -// so live streaming output uses the same row format and session-start banners -// as `captain history`. Both real tool uses and synthetic Result/SessionInit -// entries flow through the same rendering path. Returns false when there is -// nothing renderable so the caller can fall back to generic event handling. -func renderClaudeEntry(renderer *lineRenderer, ev ai.Event, entry claude.HistoryEntry) bool { - switch ev.Kind { - case ai.EventToolUse, ai.EventResult, ai.EventSystem: - default: - return false - } - tl := claude.ExtractToolsWithTokens([]claude.HistoryEntry{entry}) - if len(tl) == 0 { - return false - } - for _, t := range tl { - renderer.Render(t, true) - } - return true -} - -// renderCodexEntry mirrors renderClaudeEntry for codex live events, which -// stash a synthesized claude.ToolUse on ev.Raw rather than a HistoryEntry -// (codex's stream schema does not match Claude's message-shaped envelope). -// Routing the codex tool use through ToolUsesToTools keeps the rendering -// path identical to `captain history` for codex JSONL. -func renderCodexEntry(renderer *lineRenderer, ev ai.Event, tu claude.ToolUse) bool { - switch ev.Kind { - case ai.EventToolUse, ai.EventResult, ai.EventSystem: - default: - return false - } - tl := claude.ToolUsesToTools([]claude.ToolUse{tu}) - if len(tl) == 0 { - return false - } - for _, t := range tl { - renderer.Render(t, true) - } - return true -} - -func summariseInput(input map[string]any) string { - if len(input) == 0 { - return "" - } - for _, key := range []string{"file_path", "path", "command", "pattern", "url"} { - if v, ok := input[key].(string); ok && v != "" { - return truncForStderr(v, 80) - } - } - return "" -} - -func truncForStderr(s string, max int) string { - if len(s) <= max { - return s - } - return s[:max] + "…" -} - type AITestOptions struct { AIProviderOptions Timeout string `flag:"timeout" help:"Request timeout" default:"60s"` diff --git a/pkg/cli/ai_agent.go b/pkg/cli/ai_agent.go index de215bdf..c8a35586 100644 --- a/pkg/cli/ai_agent.go +++ b/pkg/cli/ai_agent.go @@ -2,6 +2,7 @@ package cli import ( "context" + "errors" "fmt" "os" "strings" @@ -195,7 +196,7 @@ func RunAIAgent(opts AIAgentOptions) (any, error) { return nil, err } - renderer := newLineRenderer(os.Stderr, 8) + renderer := NewEventRenderer(os.Stderr) runner := &agent.Runner[string]{ Provider: sp, Request: baseReq, @@ -204,7 +205,7 @@ func RunAIAgent(opts AIAgentOptions) (any, error) { Repo: cwd, Cwd: cwd, Scope: scope, - OnEvent: func(_ int, ev ai.Event) { renderEvent(os.Stderr, renderer, ev) }, + OnEvent: renderer.Handle, } timeout, _ := time.ParseDuration(opts.Timeout) @@ -216,6 +217,7 @@ func RunAIAgent(opts AIAgentOptions) (any, error) { start := time.Now() result, runErr := runner.Run(ctx) + renderErr := renderer.Flush() ws := result.Response.Workspace res := AIAgentResult{ @@ -237,8 +239,8 @@ func RunAIAgent(opts AIAgentOptions) (any, error) { // A failed loop/verify is surfaced through the result (Passed=false), not as // a command error, so --format output is still rendered. A genuine provider // or plugin error is returned. - if runErr != nil && len(result.Verdicts) == 0 { - return res, runErr + if renderErr != nil || (runErr != nil && len(result.Verdicts) == 0) { + return res, errors.Join(runErr, renderErr) } return res, nil } diff --git a/pkg/cli/ai_render_codex_test.go b/pkg/cli/ai_render_codex_test.go deleted file mode 100644 index 8df8469a..00000000 --- a/pkg/cli/ai_render_codex_test.go +++ /dev/null @@ -1,87 +0,0 @@ -package cli - -import ( - "bytes" - "strings" - "testing" - - "github.com/flanksource/captain/pkg/ai" - "github.com/flanksource/captain/pkg/claude" -) - -// TestRenderEvent_CodexLiveUsesLineRenderer verifies that codex live events -// flow through the same shared lineRenderer as `captain history` does for -// codex JSONL — emitting a session-start banner, a tool row, and a result row -// with cost/usage. Without unification, renderEvent falls back to the bare -// "[tool] name" / "[result] ..." printer. -func TestRenderEvent_CodexLiveUsesLineRenderer(t *testing.T) { - var buf bytes.Buffer - renderer := newLineRenderer(&buf, 8) - - session := claude.ToolUse{ - Tool: "SessionInit", - SessionID: "019e0365-dc2a-7ad0-a5a8-78936481a928", - Source: "codex", - Model: "gpt-5", - } - exec := claude.ToolUse{ - Tool: "Bash", - Input: map[string]any{"command": "ls"}, - SessionID: "019e0365-dc2a-7ad0-a5a8-78936481a928", - Source: "codex", - Model: "gpt-5", - } - result := claude.ToolUse{ - Tool: "Result", - Input: map[string]any{"total_cost_usd": 0.5}, - SessionID: "019e0365-dc2a-7ad0-a5a8-78936481a928", - Source: "codex", - Model: "gpt-5", - InputTokens: 100, - OutputTokens: 50, - } - - renderEvent(nil, renderer, ai.Event{ - Kind: ai.EventSystem, - Tool: "SessionInit", - SessionID: session.SessionID, - Model: session.Model, - Raw: session, - }) - renderEvent(nil, renderer, ai.Event{ - Kind: ai.EventToolUse, - Tool: exec.Tool, - Input: exec.Input, - Model: exec.Model, - Raw: exec, - }) - renderEvent(nil, renderer, ai.Event{ - Kind: ai.EventResult, - Tool: "Result", - Model: result.Model, - Success: true, - CostUSD: 0.5, - Usage: &ai.Usage{InputTokens: 100, OutputTokens: 50}, - Raw: result, - }) - - out := buf.String() - for _, want := range []string{ - "Codex", // session header capitalises the source name - "gpt-5", // model in the header - "019e0365", // shortened session id - "Bash", // tool row label - "$0.5", // cost - } { - if !strings.Contains(out, want) { - t.Errorf("rendered output missing %q\nfull output:\n%s", want, out) - } - } - - // Negative: the bare fallback printer must NOT be used when Raw is set. - for _, forbidden := range []string{"[tool]", "[result]", "[session]"} { - if strings.Contains(out, forbidden) { - t.Errorf("output should not contain bare fallback %q\nfull output:\n%s", forbidden, out) - } - } -} diff --git a/pkg/cli/analysis.go b/pkg/cli/analysis.go index 28305a5b..b3686c65 100644 --- a/pkg/cli/analysis.go +++ b/pkg/cli/analysis.go @@ -47,39 +47,21 @@ func AnalyzeToolUse(t tools.Tool) ToolAnalysis { return claude.AbsolutePath(path, base.CWD, base.ProjectRoot) } - // File-writing tools are resolved through history's canonical tool→input-key - // table rather than a second list here, so the write set reported to callers - // that stage from it cannot drift from the one the agent runner records. - // It is also the only place that knows NotebookEdit names its file - // notebook_path, which a plain FilePath() lookup misses entirely. - for _, path := range history.ModifiedFiles([]history.ToolUse{{Tool: base.RawTool, Input: base.Input}}) { + // Which files a tool touched is resolved through history's canonical + // footprint rather than a second set of rules here, so the paths reported to + // callers that stage from them cannot drift from the ones the session + // builders persist and the agent runner records. + footprint := history.ToolFootprint(history.ToolUse{Tool: base.RawTool, Input: base.Input, CWD: base.CWD}) + for _, path := range footprint.Written { a.WritePaths = appendUnique(a.WritePaths, abs(path)) } + for _, path := range footprint.Read { + a.ReadPaths = appendUnique(a.ReadPaths, abs(path)) + } switch base.RawTool { - case "Read": - if path := t.FilePath(); path != "" { - a.ReadPaths = append(a.ReadPaths, abs(path)) - } - case "Grep": - if path, ok := base.Input["path"].(string); ok && path != "" { - a.ReadPaths = append(a.ReadPaths, abs(path)) - } - case "Glob": - if path, ok := base.Input["path"].(string); ok && path != "" { - a.ReadPaths = append(a.ReadPaths, abs(path)) - } case "Bash": - a.analyzeBash(base.Input, abs) - command, _ := base.Input["command"].(string) - for _, path := range tools.ExtractApplyPatchPaths(command) { - a.WritePaths = appendUnique(a.WritePaths, abs(path)) - } - case "exec", "apply_patch": - input, _ := base.Input["input"].(string) - for _, path := range tools.ExtractApplyPatchPaths(input) { - a.WritePaths = appendUnique(a.WritePaths, abs(path)) - } + a.analyzeBash(base.Input) case "WebFetch": if urlStr, ok := base.Input["url"].(string); ok { if u, err := url.Parse(urlStr); err == nil && u.Host != "" { @@ -107,25 +89,15 @@ func AnalyzeToolUse(t tools.Tool) ToolAnalysis { return a } -func (a *ToolAnalysis) analyzeBash(input map[string]any, abs func(string) string) { +// analyzeBash covers the parts of a shell command that are not a file +// footprint: which binaries it invokes and which hosts it can reach. The read +// and write paths come from history.ToolFootprint with every other tool's. +func (a *ToolAnalysis) analyzeBash(input map[string]any) { cmd, _ := input["command"].(string) if cmd == "" { return } - result, err := bash.Analyze(cmd) - if result == nil { - return - } - _ = err - - for _, op := range result.Operations { - a.WritePaths = appendUnique(a.WritePaths, abs(op.Path)) - } - for _, path := range result.ReferencedPaths { - a.ReadPaths = appendUnique(a.ReadPaths, abs(path)) - } - binaries := make(map[string]bool) extractBinariesFromInput(input, binaries) a.Binaries = sortedKeys(binaries) diff --git a/pkg/cli/analysis_test.go b/pkg/cli/analysis_test.go index 1871472d..d613d385 100644 --- a/pkg/cli/analysis_test.go +++ b/pkg/cli/analysis_test.go @@ -135,9 +135,11 @@ func TestAnalyzeToolUse_Bash(t *testing.T) { domains: []string{}, }, { + // A file the command creates is a write, not a read. It used to be + // reported as both here while the Codex session builder reported it + // as a write only; history.ToolFootprint now decides once. name: "touch creates file", cmd: "touch /home/user/project/output.txt", - readPaths: []string{"/home/user/project/output.txt"}, writePaths: []string{"/home/user/project/output.txt"}, binaries: []string{"touch"}, domains: []string{}, diff --git a/pkg/cli/attachments_gc.go b/pkg/cli/attachments_gc.go index 12b185de..83a3f45d 100644 --- a/pkg/cli/attachments_gc.go +++ b/pkg/cli/attachments_gc.go @@ -94,7 +94,9 @@ func collectAttachmentReferences(root, storeDirectory string) (map[string]struct } func collectDatabaseAttachmentReferences(ctx context.Context) (map[string]struct{}, error) { - db, err := captainDB(ctx) + // GC deletes attachments captain owns, so it only ever consults the + // database captain writes. + db, err := captainDefaultDB(ctx) if err != nil { return nil, fmt.Errorf("open attachment reference database: %w", err) } diff --git a/pkg/cli/chat_thread_store.go b/pkg/cli/chat_thread_store.go deleted file mode 100644 index b825c572..00000000 --- a/pkg/cli/chat_thread_store.go +++ /dev/null @@ -1,260 +0,0 @@ -package cli - -import ( - "context" - "crypto/rand" - "encoding/hex" - "encoding/json" - "errors" - "fmt" - "os" - "path/filepath" - "sort" - "sync" - "time" - - "github.com/flanksource/captain/pkg/aichat" -) - -type fileThreadStore struct { - path string - mu sync.Mutex -} - -type threadStoreFile struct { - Threads []*aichat.Thread `json:"threads"` -} - -func newFileThreadStore(path string) *fileThreadStore { - return &fileThreadStore{path: path} -} - -func (s *fileThreadStore) Create(_ context.Context, title string) (*aichat.Thread, error) { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return nil, err - } - if title == "" { - title = "New conversation" - } - now := time.Now() - thread := &aichat.Thread{ - ID: newThreadID(), - Title: title, - CreatedAt: now, - UpdatedAt: now, - } - state.Threads = append(state.Threads, thread) - if err := s.saveLocked(state); err != nil { - return nil, err - } - return cloneThread(thread), nil -} - -func (s *fileThreadStore) List(_ context.Context) ([]*aichat.Thread, error) { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return nil, err - } - threads := make([]*aichat.Thread, 0, len(state.Threads)) - for _, thread := range state.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 *fileThreadStore) Get(_ context.Context, id string) (*aichat.Thread, error) { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return nil, err - } - thread := findThread(state, id) - if thread == nil { - return nil, fmt.Errorf("thread %q not found", id) - } - return cloneThread(thread), nil -} - -func (s *fileThreadStore) AppendMessage(_ context.Context, id string, msg aichat.UIMessage) error { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return err - } - thread := findThread(state, id) - if thread == nil { - return fmt.Errorf("thread %q not found", id) - } - thread.Messages = append(thread.Messages, msg) - thread.UpdatedAt = time.Now() - return s.saveLocked(state) -} - -func (s *fileThreadStore) ReplaceLastMessage(_ context.Context, id string, msg aichat.UIMessage) error { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return err - } - thread := findThread(state, id) - if thread == nil { - return fmt.Errorf("thread %q not found", id) - } - if msg.Role != "assistant" { - return fmt.Errorf("thread %q replacement message must have assistant role", id) - } - if len(thread.Messages) == 0 { - return fmt.Errorf("thread %q cannot replace a message in an empty thread", id) - } - if thread.Messages[len(thread.Messages)-1].Role != "assistant" { - return fmt.Errorf("thread %q last stored message must have assistant role", id) - } - thread.Messages[len(thread.Messages)-1] = msg - thread.UpdatedAt = time.Now() - return s.saveLocked(state) -} - -func (s *fileThreadStore) Delete(_ context.Context, id string) error { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return err - } - for i, thread := range state.Threads { - if thread.ID == id { - state.Threads = append(state.Threads[:i], state.Threads[i+1:]...) - return s.saveLocked(state) - } - } - return fmt.Errorf("thread %q not found", id) -} - -func (s *fileThreadStore) SetProviderSession(_ context.Context, id, providerSessionID string) error { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return err - } - thread := findThread(state, id) - if thread == nil { - return fmt.Errorf("thread %q not found", id) - } - thread.ProviderSessionID = providerSessionID - thread.UpdatedAt = time.Now() - return s.saveLocked(state) -} - -func (s *fileThreadStore) AddUsage(_ context.Context, id string, usage aichat.TurnUsage) (*aichat.Thread, error) { - s.mu.Lock() - defer s.mu.Unlock() - - state, err := s.loadLocked() - if err != nil { - return nil, err - } - thread := findThread(state, id) - if thread == nil { - return nil, fmt.Errorf("thread %q not found", id) - } - 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() - if err := s.saveLocked(state); err != nil { - return nil, err - } - return cloneThread(thread), nil -} - -func (s *fileThreadStore) loadLocked() (*threadStoreFile, error) { - data, err := os.ReadFile(s.path) - if errors.Is(err, os.ErrNotExist) { - return &threadStoreFile{}, nil - } - if err != nil { - return nil, err - } - if len(data) == 0 { - return &threadStoreFile{}, nil - } - var state threadStoreFile - if err := json.Unmarshal(data, &state); err != nil { - return nil, fmt.Errorf("read chat threads %s: %w", s.path, err) - } - return &state, nil -} - -func (s *fileThreadStore) saveLocked(state *threadStoreFile) error { - if err := os.MkdirAll(filepath.Dir(s.path), 0o755); err != nil { - return err - } - data, err := json.MarshalIndent(state, "", " ") - if err != nil { - return err - } - tmp := s.path + ".tmp" - if err := os.WriteFile(tmp, append(data, '\n'), 0o600); err != nil { - return err - } - return os.Rename(tmp, s.path) -} - -func findThread(state *threadStoreFile, id string) *aichat.Thread { - for _, thread := range state.Threads { - if thread.ID == id { - return thread - } - } - return nil -} - -func cloneThread(thread *aichat.Thread) *aichat.Thread { - if thread == nil { - return nil - } - data, err := json.Marshal(thread) - if err != nil { - copy := *thread - copy.Messages = append([]aichat.UIMessage(nil), thread.Messages...) - return © - } - var out aichat.Thread - if err := json.Unmarshal(data, &out); err != nil { - copy := *thread - copy.Messages = append([]aichat.UIMessage(nil), thread.Messages...) - return © - } - return &out -} - -func newThreadID() string { - var raw [8]byte - if _, err := rand.Read(raw[:]); err != nil { - return fmt.Sprintf("thread-%d", time.Now().UnixNano()) - } - return "thread-" + hex.EncodeToString(raw[:]) -} diff --git a/pkg/cli/chat_thread_store_test.go b/pkg/cli/chat_thread_store_test.go deleted file mode 100644 index 72a83d65..00000000 --- a/pkg/cli/chat_thread_store_test.go +++ /dev/null @@ -1,91 +0,0 @@ -package cli - -import ( - "context" - "path/filepath" - "testing" - - "github.com/flanksource/captain/pkg/aichat" -) - -func TestFileThreadStorePersistsThreads(t *testing.T) { - ctx := context.Background() - path := filepath.Join(t.TempDir(), "threads.json") - store := newFileThreadStore(path) - - thread, err := store.Create(ctx, "Launch cleanup") - if err != nil { - t.Fatalf("Create: %v", err) - } - if thread.ID == "" { - t.Fatal("Create returned empty thread id") - } - - if err := store.SetProviderSession(ctx, thread.ID, "provider-session-1"); err != nil { - t.Fatalf("SetProviderSession: %v", err) - } - msg := aichat.UIMessage{ - Role: "user", - Parts: []aichat.UIPart{{Type: "text", Text: "continue"}}, - } - if err := store.AppendMessage(ctx, thread.ID, msg); err != nil { - t.Fatalf("AppendMessage: %v", err) - } - assistant := aichat.UIMessage{ - Role: "assistant", - Parts: []aichat.UIPart{{Type: "text", Text: "pending"}}, - } - if err := store.AppendMessage(ctx, thread.ID, assistant); err != nil { - t.Fatalf("AppendMessage assistant: %v", err) - } - assistant.Parts[0].Text = "completed" - if err := store.ReplaceLastMessage(ctx, thread.ID, assistant); err != nil { - t.Fatalf("ReplaceLastMessage: %v", err) - } - updated, err := store.AddUsage(ctx, thread.ID, aichat.TurnUsage{ - InputTokens: 10, - OutputTokens: 5, - CostUSD: 0.25, - }) - if err != nil { - t.Fatalf("AddUsage: %v", err) - } - if updated.ProviderSessionID != "provider-session-1" { - t.Fatalf("ProviderSessionID = %q", updated.ProviderSessionID) - } - - reloaded := newFileThreadStore(path) - got, err := reloaded.Get(ctx, thread.ID) - if err != nil { - t.Fatalf("Get reloaded: %v", err) - } - if got.Title != "Launch cleanup" { - t.Errorf("Title = %q", got.Title) - } - if got.ProviderSessionID != "provider-session-1" { - t.Errorf("ProviderSessionID = %q", got.ProviderSessionID) - } - if len(got.Messages) != 2 || - got.Messages[0].Parts[0].Text != "continue" || - got.Messages[1].Parts[0].Text != "completed" { - 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) - } - - list, err := reloaded.List(ctx) - if err != nil { - t.Fatalf("List: %v", err) - } - if len(list) != 1 || list[0].ID != thread.ID { - t.Fatalf("List = %+v", list) - } - - if err := reloaded.Delete(ctx, thread.ID); err != nil { - t.Fatalf("Delete: %v", err) - } - if _, err := reloaded.Get(ctx, thread.ID); err == nil { - t.Fatal("Get deleted thread returned nil error") - } -} diff --git a/pkg/cli/cmux_info.go b/pkg/cli/cmux_info.go new file mode 100644 index 00000000..521da296 --- /dev/null +++ b/pkg/cli/cmux_info.go @@ -0,0 +1,191 @@ +package cli + +import ( + "context" + "fmt" + "sort" + "strconv" + "strings" + "time" + + "github.com/flanksource/captain/pkg/cmux" +) + +type CmuxInfoOptions struct { + Input []string `args:"true" stdin:"true" help:"cmux Copy IDs lines or a process ID"` + Stack bool `flag:"stack" help:"Capture goroutine stacks from Go processes with gops"` +} + +type cmuxInfoSelector struct { + PID int + Selector cmux.Selector +} + +type CmuxInfoResult struct { + Target CmuxInfoTarget `json:"target"` + Processes []CmuxProcess `json:"processes"` +} + +type CmuxInfoTarget struct { + Kind string `json:"kind"` + PID int `json:"pid,omitempty"` + Selector *cmux.Selector `json:"selector,omitempty"` +} + +type CmuxProcess struct { + PID int `json:"pid"` + PPID int `json:"ppid,omitempty"` + Name string `json:"name,omitempty"` + Executable string `json:"executable,omitempty"` + Command string `json:"command,omitempty"` + Runtime string `json:"runtime,omitempty"` + CPUPercent float64 `json:"cpu_percent"` + RSSBytes uint64 `json:"rss_bytes"` + Listeners []CmuxListener `json:"listeners"` + Locations []cmux.ProcessLocation `json:"locations"` + Stack string `json:"stack,omitempty"` + StackError string `json:"stack_error,omitempty"` + InspectionError string `json:"inspection_error,omitempty"` +} + +type CmuxListener struct { + Protocol string `json:"protocol"` + Address string `json:"address"` + Port uint32 `json:"port"` +} + +var ( + loadCmuxTop = cmux.Top + inspectCmuxPID = inspectProcess + loadGoStack = captureGoStack +) + +func parseCmuxInfoSelector(input []string) (cmuxInfoSelector, error) { + lines := splitCmuxInfoInput(input) + if len(lines) == 0 { + return cmuxInfoSelector{}, fmt.Errorf("cmux info requires cmux Copy IDs lines or a PID") + } + + parsed := cmuxInfoSelector{} + values := map[string]*string{ + "workspace_id": &parsed.Selector.WorkspaceID, "workspace_ref": &parsed.Selector.WorkspaceRef, + "pane_id": &parsed.Selector.PaneID, "pane_ref": &parsed.Selector.PaneRef, + "surface_id": &parsed.Selector.SurfaceID, "surface_ref": &parsed.Selector.SurfaceRef, + } + for _, line := range lines { + if !strings.Contains(line, "=") { + if err := parseCmuxInfoPID(line, &parsed); err != nil { + return cmuxInfoSelector{}, err + } + continue + } + parts := strings.SplitN(line, "=", 2) + destination, ok := values[strings.TrimSpace(parts[0])] + if !ok { + return cmuxInfoSelector{}, fmt.Errorf("unknown cmux selector key %q", strings.TrimSpace(parts[0])) + } + value := strings.TrimSpace(parts[1]) + if value == "" { + return cmuxInfoSelector{}, fmt.Errorf("cmux selector %s is empty", strings.TrimSpace(parts[0])) + } + if *destination != "" && *destination != value { + return cmuxInfoSelector{}, fmt.Errorf("conflicting %s values %q and %q", strings.TrimSpace(parts[0]), *destination, value) + } + *destination = value + } + if parsed.PID > 0 && parsed.Selector.Kind() != "" { + return cmuxInfoSelector{}, fmt.Errorf("cannot mix a PID with cmux selector lines") + } + if parsed.PID == 0 && parsed.Selector.Kind() == "" { + return cmuxInfoSelector{}, fmt.Errorf("cmux info requires cmux Copy IDs lines or a PID") + } + return parsed, nil +} + +func splitCmuxInfoInput(input []string) []string { + var lines []string + for _, value := range input { + for _, line := range strings.Split(value, "\n") { + if line = strings.TrimSpace(line); line != "" { + lines = append(lines, line) + } + } + } + return lines +} + +func parseCmuxInfoPID(value string, parsed *cmuxInfoSelector) error { + pid, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return fmt.Errorf("invalid cmux selector line %q", value) + } + if pid <= 0 { + return fmt.Errorf("PID must be positive, got %d", pid) + } + if parsed.PID != 0 && parsed.PID != pid { + return fmt.Errorf("multiple PIDs are not supported: %d and %d", parsed.PID, pid) + } + parsed.PID = pid + return nil +} + +func RunCmuxInfo(ctx context.Context, opts CmuxInfoOptions) (CmuxInfoResult, error) { + selector, err := parseCmuxInfoSelector(opts.Input) + if err != nil { + return CmuxInfoResult{}, err + } + target, pids, locations, err := resolveCmuxInfoTarget(selector) + if err != nil { + return CmuxInfoResult{}, err + } + + result := CmuxInfoResult{Target: target, Processes: make([]CmuxProcess, 0, len(pids))} + for _, pid := range pids { + if err := ctx.Err(); err != nil { + return CmuxInfoResult{}, err + } + row, inspectErr := inspectCmuxPID(ctx, pid) + row.PID = pid + row.Locations = locations[pid] + if row.Locations == nil { + row.Locations = []cmux.ProcessLocation{} + } + if inspectErr != nil { + row.InspectionError = inspectErr.Error() + } + if opts.Stack && row.Runtime == "go" { + stackCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + row.Stack, err = loadGoStack(stackCtx, pid) + cancel() + if err != nil { + row.StackError = err.Error() + } + } + result.Processes = append(result.Processes, row) + } + return result, nil +} + +func resolveCmuxInfoTarget(selector cmuxInfoSelector) (CmuxInfoTarget, []int, map[int][]cmux.ProcessLocation, error) { + if selector.PID > 0 { + return CmuxInfoTarget{Kind: "pid", PID: selector.PID}, []int{selector.PID}, nil, nil + } + snapshot, err := loadCmuxTop() + if err != nil { + return CmuxInfoTarget{}, nil, nil, err + } + resolved, err := snapshot.Resolve(selector.Selector) + if err != nil { + return CmuxInfoTarget{}, nil, nil, err + } + if len(resolved.PIDs) == 0 { + return CmuxInfoTarget{}, nil, nil, fmt.Errorf("cmux %s target has no running processes", resolved.Kind) + } + for _, pid := range resolved.PIDs { + if pid <= 0 { + return CmuxInfoTarget{}, nil, nil, fmt.Errorf("cmux %s target returned invalid PID %d", resolved.Kind, pid) + } + } + sort.Ints(resolved.PIDs) + return CmuxInfoTarget{Kind: resolved.Kind, Selector: &selector.Selector}, resolved.PIDs, resolved.Locations, nil +} diff --git a/pkg/cli/cmux_info_ginkgo_test.go b/pkg/cli/cmux_info_ginkgo_test.go new file mode 100644 index 00000000..3413c304 --- /dev/null +++ b/pkg/cli/cmux_info_ginkgo_test.go @@ -0,0 +1,225 @@ +package cli + +import ( + "context" + "encoding/json" + "errors" + "os" + "syscall" + "time" + + "github.com/flanksource/captain/pkg/cmux" + gopsnet "github.com/shirou/gopsutil/v3/net" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("cmux info selectors", func() { + DescribeTable("parses valid input", + func(input []string, expected cmuxInfoSelector) { + actual, err := parseCmuxInfoSelector(input) + Expect(err).NotTo(HaveOccurred()) + Expect(actual).To(Equal(expected)) + }, + Entry("bare PID", []string{"33745"}, cmuxInfoSelector{PID: 33745}), + Entry("separate Copy IDs lines", []string{ + "workspace_ref=workspace:3", + "pane_id=pane-id", + "surface_ref=surface:21", + }, cmuxInfoSelector{Selector: cmux.Selector{ + WorkspaceRef: "workspace:3", + PaneID: "pane-id", + SurfaceRef: "surface:21", + }}), + Entry("quoted multiline Copy IDs", []string{"pane_ref=pane:5\nsurface_id=surface-id"}, cmuxInfoSelector{Selector: cmux.Selector{ + PaneRef: "pane:5", + SurfaceID: "surface-id", + }}), + ) + + DescribeTable("rejects invalid input", + func(input []string, message string) { + _, err := parseCmuxInfoSelector(input) + Expect(err).To(MatchError(ContainSubstring(message))) + }, + Entry("empty input", nil, "requires cmux Copy IDs lines or a PID"), + Entry("unknown field", []string{"tab_id=tab-id"}, "unknown cmux selector key"), + Entry("conflicting fields", []string{"pane_ref=pane:5", "pane_ref=pane:6"}, "conflicting pane_ref"), + Entry("PID mixed with selector", []string{"33745", "pane_ref=pane:5"}, "cannot mix a PID"), + Entry("non-positive PID", []string{"0"}, "PID must be positive"), + ) +}) + +var _ = Describe("process runtime detection", func() { + DescribeTable("classifies executables", + func(executable, name, expected string) { + Expect(detectProcessRuntime(executable, name)).To(Equal(expected)) + }, + Entry("Node", "/usr/local/bin/node", "node", "node"), + Entry("Bun", "/usr/local/bin/bun", "bun", "bun"), + Entry("Deno", "/usr/local/bin/deno", "deno", "deno"), + Entry("Python", "/usr/bin/python3.13", "Python", "python"), + Entry("Java", "/usr/bin/java", "java", "java"), + Entry("Ruby", "/usr/bin/ruby", "ruby", "ruby"), + Entry("shell", "/bin/zsh", "zsh", "shell"), + Entry("native", "/usr/bin/git", "git", "native"), + Entry("unknown", "", "", "unknown"), + ) + + It("detects Go build information before executable-name heuristics", func() { + executable, err := os.Executable() + Expect(err).NotTo(HaveOccurred()) + Expect(detectProcessRuntime(executable, "node")).To(Equal("go")) + }) + + It("reports only listening TCP endpoints in stable order", func() { + connections := []gopsnet.ConnectionStat{ + {Type: syscall.SOCK_DGRAM, Laddr: gopsnet.Addr{IP: "127.0.0.1", Port: 5353}}, + {Type: syscall.SOCK_STREAM, Status: "ESTABLISHED", Laddr: gopsnet.Addr{IP: "127.0.0.1", Port: 443}}, + {Type: syscall.SOCK_STREAM, Status: "LISTEN", Laddr: gopsnet.Addr{IP: "::1", Port: 9090}}, + {Type: syscall.SOCK_STREAM, Status: "LISTEN", Laddr: gopsnet.Addr{IP: "127.0.0.1", Port: 8080}}, + } + + Expect(listeningTCPEndpoints(connections)).To(Equal([]CmuxListener{ + {Protocol: "tcp4", Address: "127.0.0.1", Port: 8080}, + {Protocol: "tcp6", Address: "::1", Port: 9090}, + })) + }) + + It("inspects the current Go process", func() { + row, err := inspectProcess(context.Background(), os.Getpid()) + + Expect(err).NotTo(HaveOccurred()) + Expect(row).To(SatisfyAll( + HaveField("PID", os.Getpid()), + HaveField("PPID", BeNumerically(">", 0)), + HaveField("Executable", Not(BeEmpty())), + HaveField("Runtime", "go"), + HaveField("RSSBytes", BeNumerically(">", 0)), + )) + }) +}) + +var _ = Describe("RunCmuxInfo", func() { + BeforeEach(func() { + originalTop := loadCmuxTop + originalInspect := inspectCmuxPID + originalStack := loadGoStack + DeferCleanup(func() { + loadCmuxTop = originalTop + inspectCmuxPID = originalInspect + loadGoStack = originalStack + }) + }) + + It("inspects every resolved process in PID order and preserves locations", func() { + loadCmuxTop = func() (cmux.TopSnapshot, error) { + return cmux.TopSnapshot{Windows: []cmux.TopWindow{{Workspaces: []cmux.TopWorkspace{{ + ID: "workspace-id", Ref: "workspace:3", Panes: []cmux.TopPane{{ + ID: "pane-id", Ref: "pane:5", Surfaces: []cmux.TopSurface{{ + ID: "surface-id", Ref: "surface:21", Resources: cmux.ProcessResources{PIDs: []int{202, 101}}, + }}, + }}, + }}}}}, nil + } + inspectCmuxPID = func(_ context.Context, pid int) (CmuxProcess, error) { + return CmuxProcess{PID: pid, Runtime: "native"}, nil + } + + result, err := RunCmuxInfo(context.Background(), CmuxInfoOptions{Input: []string{"surface_ref=surface:21"}}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.Target.Kind).To(Equal("surface")) + Expect(result.Processes).To(HaveLen(2)) + Expect([]int{result.Processes[0].PID, result.Processes[1].PID}).To(Equal([]int{101, 202})) + Expect(result.Processes[0].Locations).To(HaveLen(1)) + }) + + It("captures stacks only for Go processes and keeps gops failures on the row", func() { + loadCmuxTop = func() (cmux.TopSnapshot, error) { + return cmux.TopSnapshot{Windows: []cmux.TopWindow{{Workspaces: []cmux.TopWorkspace{{ + ID: "workspace-id", Ref: "workspace:3", Panes: []cmux.TopPane{{ + ID: "pane-id", Ref: "pane:5", Surfaces: []cmux.TopSurface{{ + ID: "surface-id", Ref: "surface:21", Resources: cmux.ProcessResources{PIDs: []int{101, 202}}, + }}, + }}, + }}}}}, nil + } + inspectCmuxPID = func(_ context.Context, pid int) (CmuxProcess, error) { + runtime := "node" + if pid == 101 { + runtime = "go" + } + return CmuxProcess{PID: pid, Runtime: runtime}, nil + } + loadGoStack = func(_ context.Context, pid int) (string, error) { + Expect(pid).To(Equal(101)) + return "", errors.New("gops agent is unavailable") + } + + result, err := RunCmuxInfo(context.Background(), CmuxInfoOptions{Input: []string{"surface_ref=surface:21"}, Stack: true}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.Processes[0].StackError).To(Equal("gops agent is unavailable")) + Expect(result.Processes[1].StackError).To(BeEmpty()) + }) + + It("keeps processes that disappear during inspection", func() { + inspectCmuxPID = func(_ context.Context, pid int) (CmuxProcess, error) { + return CmuxProcess{PID: pid}, errors.New("process no longer exists") + } + + result, err := RunCmuxInfo(context.Background(), CmuxInfoOptions{Input: []string{"101"}}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.Processes[0].InspectionError).To(Equal("process no longer exists")) + }) + + It("fails when a resolved cmux target has no processes", func() { + loadCmuxTop = func() (cmux.TopSnapshot, error) { + return cmux.TopSnapshot{Windows: []cmux.TopWindow{{Workspaces: []cmux.TopWorkspace{{ + ID: "workspace-id", Ref: "workspace:3", + }}}}}, nil + } + + _, err := RunCmuxInfo(context.Background(), CmuxInfoOptions{Input: []string{"workspace_ref=workspace:3"}}) + + Expect(err).To(MatchError(ContainSubstring("has no running processes"))) + }) + + It("renders typed rows with detail-only diagnostics", func() { + row := CmuxProcess{ + PID: 101, PPID: 50, Name: "captain", Runtime: "go", CPUPercent: 1.25, RSSBytes: 64 * 1024 * 1024, + Executable: "/opt/bin/captain", Command: "captain serve", Stack: "goroutine 1 [running]", StackError: "", + Listeners: []CmuxListener{{Protocol: "tcp4", Address: "127.0.0.1", Port: 8080}}, + Locations: []cmux.ProcessLocation{{SurfaceRef: "surface:21", SurfaceTitle: "API"}}, + } + + Expect(row.Row()).To(HaveKeyWithValue("pid", "101")) + Expect(row.Row()).To(HaveKey("listeners")) + Expect(row.RowDetail().String()).To(ContainSubstring("/opt/bin/captain")) + Expect(row.RowDetail().String()).To(ContainSubstring("goroutine 1")) + Expect(CmuxInfoResult{Processes: []CmuxProcess{row}}.Pretty().String()).To(ContainSubstring("captain")) + }) + + It("keeps zero-valued resource metrics and omits PID selectors in JSON", func() { + payload, err := json.Marshal(CmuxInfoResult{ + Target: CmuxInfoTarget{Kind: "pid", PID: 101}, + Processes: []CmuxProcess{{PID: 101, Listeners: []CmuxListener{}, Locations: []cmux.ProcessLocation{}}}, + }) + + Expect(err).NotTo(HaveOccurred()) + Expect(payload).To(MatchJSON(`{ + "target":{"kind":"pid","pid":101}, + "processes":[{"pid":101,"cpu_percent":0,"rss_bytes":0,"listeners":[],"locations":[]}] + }`)) + }) + + It("bounds gops stack collection", func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancel() + _, err := captureGoStack(ctx, 101) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/pkg/cli/cmux_info_process.go b/pkg/cli/cmux_info_process.go new file mode 100644 index 00000000..afedd32a --- /dev/null +++ b/pkg/cli/cmux_info_process.go @@ -0,0 +1,158 @@ +package cli + +import ( + "context" + "debug/buildinfo" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "sort" + "strconv" + "strings" + "syscall" + + gopsnet "github.com/shirou/gopsutil/v3/net" + "github.com/shirou/gopsutil/v3/process" +) + +func inspectProcess(ctx context.Context, pid int) (CmuxProcess, error) { + row := CmuxProcess{PID: pid, Listeners: []CmuxListener{}} + proc, err := process.NewProcessWithContext(ctx, int32(pid)) + if err != nil { + return row, fmt.Errorf("open process %d: %w", pid, err) + } + + var issues []error + if value, err := proc.PpidWithContext(ctx); err != nil { + issues = append(issues, fmt.Errorf("read parent PID: %w", err)) + } else { + row.PPID = int(value) + } + if row.Name, err = proc.NameWithContext(ctx); err != nil { + issues = append(issues, fmt.Errorf("read process name: %w", err)) + } + if row.Executable, err = proc.ExeWithContext(ctx); err != nil { + issues = append(issues, fmt.Errorf("read executable: %w", err)) + } + if row.Command, err = proc.CmdlineWithContext(ctx); err != nil { + issues = append(issues, fmt.Errorf("read command: %w", err)) + } + if row.CPUPercent, err = proc.CPUPercentWithContext(ctx); err != nil { + issues = append(issues, fmt.Errorf("read CPU usage: %w", err)) + } + if memory, memoryErr := proc.MemoryInfoWithContext(ctx); memoryErr != nil { + issues = append(issues, fmt.Errorf("read memory usage: %w", memoryErr)) + } else { + row.RSSBytes = memory.RSS + } + if connections, connectionsErr := proc.ConnectionsWithContext(ctx); connectionsErr != nil { + issues = append(issues, fmt.Errorf("read network listeners: %w", connectionsErr)) + } else { + row.Listeners = listeningTCPEndpoints(connections) + } + row.Runtime = detectProcessRuntime(row.Executable, row.Name) + return row, errors.Join(issues...) +} + +func detectProcessRuntime(executable, name string) string { + if executable != "" { + if _, err := buildinfo.ReadFile(executable); err == nil { + return "go" + } + } + if executable == "" && name == "" { + return "unknown" + } + binary := strings.ToLower(filepath.Base(executable)) + if binary == "" { + binary = strings.ToLower(name) + } + binary = strings.TrimSuffix(binary, ".exe") + switch { + case binary == "node" || strings.HasPrefix(binary, "nodejs"): + return "node" + case binary == "bun": + return "bun" + case binary == "deno": + return "deno" + case strings.HasPrefix(binary, "python"): + return "python" + case binary == "java": + return "java" + case strings.HasPrefix(binary, "ruby"): + return "ruby" + case isShell(binary): + return "shell" + default: + return "native" + } +} + +func captureGoStack(ctx context.Context, pid int) (string, error) { + gops, err := gopsBinary() + if err != nil { + return "", err + } + output, err := exec.CommandContext(ctx, gops, "stack", strconv.Itoa(pid)).CombinedOutput() + if err == nil { + return strings.TrimSpace(string(output)), nil + } + if ctx.Err() != nil { + return "", fmt.Errorf("gops stack %d: %w", pid, ctx.Err()) + } + if detail := strings.TrimSpace(string(output)); detail != "" { + return "", fmt.Errorf("gops stack %d: %w: %s", pid, err, detail) + } + return "", fmt.Errorf("gops stack %d: %w", pid, err) +} + +func isShell(binary string) bool { + switch binary { + case "sh", "bash", "zsh", "fish", "dash", "ksh", "csh", "tcsh": + return true + default: + return false + } +} + +func listeningTCPEndpoints(connections []gopsnet.ConnectionStat) []CmuxListener { + listeners := make([]CmuxListener, 0) + for _, connection := range connections { + if connection.Type != syscall.SOCK_STREAM || !strings.EqualFold(connection.Status, "LISTEN") { + continue + } + protocol := "tcp4" + if strings.Contains(connection.Laddr.IP, ":") { + protocol = "tcp6" + } + listeners = append(listeners, CmuxListener{Protocol: protocol, Address: connection.Laddr.IP, Port: connection.Laddr.Port}) + } + sort.Slice(listeners, func(i, j int) bool { + return listeners[i].String() < listeners[j].String() + }) + return listeners +} + +func gopsBinary() (string, error) { + if configured := os.Getenv("GOPS_BIN"); configured != "" { + path, err := exec.LookPath(configured) + if err != nil { + return "", fmt.Errorf("find GOPS_BIN %q: %w", configured, err) + } + return path, nil + } + if path, err := exec.LookPath("gops"); err == nil { + return path, nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("find gops: %w", err) + } + path, err := exec.LookPath(filepath.Join(home, "go", "bin", "gops")) + if err != nil { + return "", fmt.Errorf("find gops: %w", err) + } + return path, nil +} diff --git a/pkg/cli/cmux_info_render.go b/pkg/cli/cmux_info_render.go new file mode 100644 index 00000000..1b8bb93a --- /dev/null +++ b/pkg/cli/cmux_info_render.go @@ -0,0 +1,109 @@ +package cli + +import ( + "fmt" + "net" + "strconv" + "strings" + + "github.com/flanksource/clicky/api" +) + +func (r CmuxInfoResult) Pretty() api.Text { + return api.Text{}.Add(api.NewTableFrom(r.Processes)) +} + +func (r CmuxProcess) Columns() []api.ColumnDef { + return []api.ColumnDef{ + api.Column("pid").Label("PID").Build(), + api.Column("ppid").Label("PPID").Build(), + api.Column("runtime").Label("Runtime").Build(), + api.Column("process").Label("Process").MaxWidth(24).Build(), + api.Column("cpu").Label("CPU").Build(), + api.Column("rss").Label("RSS").Build(), + api.Column("listeners").Label("Listeners").MaxWidth(30).Build(), + api.Column("surface").Label("Surface").MaxWidth(24).Build(), + } +} + +func (r CmuxProcess) Row() map[string]any { + row := map[string]any{ + "pid": strconv.Itoa(r.PID), + "runtime": r.Runtime, + "process": r.Name, + "cpu": fmt.Sprintf("%.1f%%", r.CPUPercent), + "rss": api.HumanizeBytes(int64(r.RSSBytes)), + } + if r.PPID > 0 { + row["ppid"] = strconv.Itoa(r.PPID) + } + if listeners := r.listenerLabels(); len(listeners) > 0 { + row["listeners"] = api.CompactList(listeners) + } + if surfaces := r.surfaceLabels(); len(surfaces) > 0 { + row["surface"] = api.CompactList(surfaces) + } + return row +} + +func (r CmuxProcess) RowDetail() api.Textable { + items := []api.KeyValuePair{ + api.KeyValue("Executable", r.Executable), + api.KeyValue("Command", r.Command), + api.KeyValue("Inspection error", r.InspectionError), + api.KeyValue("Stack error", r.StackError), + } + detail := api.Text{}.Add(api.DescriptionList{Items: items}) + if locations := r.locationLabels(); len(locations) > 0 { + detail = detail.NewLine().Append("Locations: ", "text-muted").Add(api.CompactList(locations)) + } + if r.Stack != "" { + detail = detail.NewLine().Add(api.CodeBlock("text/plain", r.Stack)) + } + return detail +} + +func (l CmuxListener) String() string { + host := l.Address + if host == "" { + host = "*" + } + return l.Protocol + " " + net.JoinHostPort(host, strconv.FormatUint(uint64(l.Port), 10)) +} + +func (r CmuxProcess) listenerLabels() []string { + labels := make([]string, 0, len(r.Listeners)) + for _, listener := range r.Listeners { + labels = append(labels, listener.String()) + } + return labels +} + +func (r CmuxProcess) surfaceLabels() []string { + labels := make([]string, 0, len(r.Locations)) + for _, location := range r.Locations { + label := location.SurfaceRef + if label == "" { + label = location.SurfaceID + } + if location.SurfaceTitle != "" { + label = strings.TrimSpace(label + " " + location.SurfaceTitle) + } + labels = append(labels, label) + } + return labels +} + +func (r CmuxProcess) locationLabels() []string { + labels := make([]string, 0, len(r.Locations)) + for _, location := range r.Locations { + labels = append(labels, fmt.Sprintf( + "%s (%s) / %s (%s) / %s (%s) %s", + location.WorkspaceRef, location.WorkspaceID, + location.PaneRef, location.PaneID, + location.SurfaceRef, location.SurfaceID, + location.TTY, + )) + } + return labels +} diff --git a/pkg/cli/contexts.go b/pkg/cli/contexts.go new file mode 100644 index 00000000..5eab65da --- /dev/null +++ b/pkg/cli/contexts.go @@ -0,0 +1,46 @@ +package cli + +import ( + "context" + "fmt" + + "github.com/flanksource/captain/pkg/database" +) + +// ContextsOptions configures the database context listing. +type ContextsOptions struct { + // Check opens each configured context and reports whether it answers a + // trivial query. Without it nothing is connected to. + Check bool `json:"check" flag:"check" description:"Connect to each context and report whether it is reachable"` +} + +// RunContexts lists the databases captain can read. Only the default context is +// monitored and written to; every other context is read-only. +func RunContexts(ctx context.Context, opts ContextsOptions) (ContextsResult, error) { + result, err := describeDatabaseContexts(activeDatabaseContextName(ctx)) + if err != nil { + return ContextsResult{}, err + } + if !opts.Check { + return result, nil + } + for i, row := range result.Contexts { + result.Contexts[i].Status = checkDatabaseContext(ctx, row.Name) + if dsn, source := contextDatabaseIdentity(row.Name); source != "" { + result.Contexts[i].Source = source + result.Contexts[i].DSN = database.MaskDSN(dsn) + } + } + return result, nil +} + +func checkDatabaseContext(ctx context.Context, name string) string { + db, err := openContextDB(ctx, name, captainDatabaseNoMigrations) + if err != nil { + return fmt.Sprintf("unreachable: %v", err) + } + if err := db.Gorm().WithContext(ctx).Exec("SELECT 1").Error; err != nil { + return fmt.Sprintf("unreachable: %v", err) + } + return "ok" +} diff --git a/pkg/cli/cost.go b/pkg/cli/cost.go index 7dea21cd..48ed021a 100644 --- a/pkg/cli/cost.go +++ b/pkg/cli/cost.go @@ -1,6 +1,7 @@ package cli import ( + "context" "fmt" "os" "path/filepath" @@ -61,7 +62,7 @@ type CategoryCostResult struct { Rows []CategoryCostRow `json:"rows"` } -func RunCost(opts CostOptions) (any, error) { +func RunCost(ctx context.Context, opts CostOptions) (any, error) { cwd, err := os.Getwd() if err != nil { return nil, err @@ -77,6 +78,9 @@ func RunCost(opts CostOptions) (any, error) { return nil, err } sessions = filterCostsBySessionID(sessions, sessionIDs) + // Transcripts hold no result record, so prefer the figures captain recorded + // from the provider's own results for any session it ran. + applyResultCosts(ctx, sessions) grouped := groupSessions(sessions, opts.GroupBy) @@ -92,6 +96,7 @@ func RunCost(opts CostOptions) (any, error) { total.CacheWriteTokens += s.Tokens.CacheWriteTokens total.CacheReadTokens += s.Tokens.CacheReadTokens total.TotalCost += s.Tokens.TotalCost + total.ProviderCostUSD += s.Tokens.ProviderCostUSD rows = append(rows, CostRow{ Project: s.Project, @@ -102,13 +107,13 @@ func RunCost(opts CostOptions) (any, error) { CacheRead: session.FormatTokens(s.Tokens.CacheReadTokens), CacheWrite: session.FormatTokens(s.Tokens.CacheWriteTokens), Msgs: s.Messages, - APICost: session.FormatCost(s.Tokens.TotalCost), + APICost: session.FormatCostEstimated(s.Tokens.TotalCost, s.Tokens.Estimated()), Time: claude.FormatTimeAgo(&s.End), }) } return CostResult{ - TotalAPICost: session.FormatCost(total.TotalCost), + TotalAPICost: session.FormatCostEstimated(total.TotalCost, total.Estimated()), TotalTokens: session.FormatTokens(total.TotalTokens()), Rows: rows, }, nil @@ -325,6 +330,7 @@ func mergeInto(g *claude.SessionCost, s claude.SessionCost) { g.Tokens.CacheWriteTokens += s.Tokens.CacheWriteTokens g.Tokens.CacheReadTokens += s.Tokens.CacheReadTokens g.Tokens.TotalCost += s.Tokens.TotalCost + g.Tokens.ProviderCostUSD += s.Tokens.ProviderCostUSD g.Messages += s.Messages if s.Start.Before(g.Start) { g.Start = s.Start diff --git a/pkg/cli/database.go b/pkg/cli/database.go index 13a2a6ca..57d3d7e4 100644 --- a/pkg/cli/database.go +++ b/pkg/cli/database.go @@ -2,7 +2,6 @@ package cli import ( "context" - "errors" "fmt" "os" "strings" @@ -12,126 +11,45 @@ import ( "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 + databaseURLFlag = "db-url" + databaseContextFlag = "context" ) -// 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 - } +var ( + databaseURLs []string + databaseContextFlagValue string +) - 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 +// BindDatabaseFlags exposes database selection as root persistent flags. A bare +// --db-url overrides the monitored database; name=URL declares an additional +// read-only context, which --context then selects. +func BindDatabaseFlags(flags *pflag.FlagSet) { + flags.StringArrayVar(&databaseURLs, databaseURLFlag, nil, + "PostgreSQL database URL, or name=URL to declare an additional read-only context (repeatable)") + flags.StringVar(&databaseContextFlagValue, databaseContextFlag, "", + "Database context to read from (default: the monitored database)") } -// 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 +// captainDB opens the database for the active context. Reads use this; writes +// must use captainDefaultDB. +func captainDB(ctx context.Context) (*database.DB, error) { + return openContextDB(ctx, activeDatabaseContextName(ctx), captainDatabaseNoMigrations) } -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 +// captainDefaultDB opens the monitored database regardless of the active +// context. Every write goes through it, because captain only ever owns the +// default context's data. +func captainDefaultDB(ctx context.Context) (*database.DB, error) { + return openContextDB(ctx, defaultDatabaseContextName, captainDatabaseNoMigrations) } -func captainDatabaseIdentity() (dsn, source string) { - captainDBState.mu.Lock() - defer captainDBState.mu.Unlock() - return captainDBState.dsn, captainDBState.source +// captainServeDB opens and migrates the monitored database. It is the only +// migrating path in the process. +func captainServeDB(ctx context.Context) (*database.DB, error) { + return openContextDB(ctx, defaultDatabaseContextName, captainDatabaseWithMigrations) } // serveMonitorState holds the serve process's live monitor so prompt-run code @@ -164,11 +82,19 @@ 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. +// +// Monitoring is a property of this machine and the database it writes, so a +// non-default context is returned as read: a read of another machine's +// database must never write to it. func freshenSessionDB(ctx context.Context) (*database.DB, error) { - db, err := captainDB(ctx) + name := activeDatabaseContextName(ctx) + db, err := openContextDB(ctx, name, captainDatabaseNoMigrations) if err != nil { return nil, err } + if name != defaultDatabaseContextName { + return db, nil + } 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) @@ -176,11 +102,16 @@ func freshenSessionDB(ctx context.Context) (*database.DB, error) { return db, nil } -// captainDSN resolves the database connection: explicit env DSNs first, then a -// gavel-shared database, finally captain's own shared embedded postgres. +// captainDSN resolves the default context's 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 + override, err := defaultDatabaseURLOverride() + if err != nil { + return "", "", err + } + if override != "" { + return override, "--" + databaseURLFlag, nil } for _, env := range []string{gavelDBEnvDSN, gavelCacheEnvDSN, captainSessionEnvDSN} { if dsn := strings.TrimSpace(os.Getenv(env)); dsn != "" { diff --git a/pkg/cli/database_mode_ginkgo_test.go b/pkg/cli/database_mode_ginkgo_test.go index eb09717b..05d4bd85 100644 --- a/pkg/cli/database_mode_ginkgo_test.go +++ b/pkg/cli/database_mode_ginkgo_test.go @@ -16,11 +16,9 @@ var _ = Describe("Captain database migration mode", Serial, func() { 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() + setCaptainContextDBForTest(testDatabaseHandle{ + Name: defaultDatabaseContextName, DB: db, Unmigrated: true, + }) _, err = captainServeDB(ctx) diff --git a/pkg/cli/database_test.go b/pkg/cli/database_test.go index 0bd0d465..923772c9 100644 --- a/pkg/cli/database_test.go +++ b/pkg/cli/database_test.go @@ -40,14 +40,14 @@ func TestConfigureNativeDatabaseRejectsNil(t *testing.T) { func TestCaptainDSNPrecedence(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) - databaseURL = "" - t.Cleanup(func() { databaseURL = "" }) + databaseURLs = nil + t.Cleanup(func() { databaseURLs = nil }) t.Run("db-url flag wins", func(t *testing.T) { flags := pflag.NewFlagSet("captain", pflag.ContinueOnError) - BindDatabaseURLFlag(flags) + BindDatabaseFlags(flags) require.NoError(t, flags.Parse([]string{"--db-url", "postgres://flag/captain"})) - t.Cleanup(func() { databaseURL = "" }) + t.Cleanup(func() { databaseURLs = nil }) t.Setenv(gavelDBEnvDSN, "postgres://primary/gavel") t.Setenv(gavelCacheEnvDSN, "postgres://cache/gavel") diff --git a/pkg/cli/db_context.go b/pkg/cli/db_context.go new file mode 100644 index 00000000..54594974 --- /dev/null +++ b/pkg/cli/db_context.go @@ -0,0 +1,284 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "os" + "regexp" + "sort" + "strings" + "sync" +) + +const ( + // defaultDatabaseContextName is the reserved name of the database captain + // monitors, migrates, and writes to. Every other context is read-only. + defaultDatabaseContextName = "default" + // databaseContextEnv selects the active context for a shell. + databaseContextEnv = "CAPTAIN_DB_CONTEXT" + // databaseContextsEnv defines ad-hoc contexts as ";"- or newline-separated + // name=dsn entries. + databaseContextsEnv = "CAPTAIN_DB_CONTEXTS" +) + +var errUnknownDatabaseContext = errors.New("unknown database context") + +// databaseContextNamePattern keeps context names usable as cookie values, flag +// values, and map keys without quoting. +var databaseContextNamePattern = regexp.MustCompile(`^[a-z0-9][a-z0-9_.-]{0,62}$`) + +// DatabaseContext is one readable Captain database. Exactly one context is +// active per CLI invocation or HTTP request; only the default is monitored, +// migrated, and written to. +type DatabaseContext struct { + Name string + Label string + DSN string + Source string + Default bool + ReadOnly bool +} + +type databaseContextKey struct{} + +// ContextWithDatabaseContext binds a database context name to ctx. The HTTP +// middleware and the root command's persistent hook are the only writers. +func ContextWithDatabaseContext(ctx context.Context, name string) context.Context { + return context.WithValue(ctx, databaseContextKey{}, name) +} + +func databaseContextNameFromContext(ctx context.Context) (string, bool) { + if ctx == nil { + return "", false + } + name, ok := ctx.Value(databaseContextKey{}).(string) + if !ok || strings.TrimSpace(name) == "" { + return "", false + } + return name, true +} + +// activeDatabaseContextName resolves the context a read should target: the ctx +// value set by the HTTP middleware or the root persistent hook, then the +// --context flag (also the /api/v1 executor's flags path), then the +// environment, then the default. +func activeDatabaseContextName(ctx context.Context) string { + if name, ok := databaseContextNameFromContext(ctx); ok { + return name + } + if name := strings.TrimSpace(databaseContextFlagValue); name != "" { + return name + } + if name := strings.TrimSpace(os.Getenv(databaseContextEnv)); name != "" { + return name + } + return defaultDatabaseContextName +} + +// ResolveDatabaseContextName resolves the active context name and verifies it +// is configured, so an unknown --context is rejected before any command runs +// rather than at the first query. +func ResolveDatabaseContextName(ctx context.Context) (string, error) { + name := activeDatabaseContextName(ctx) + if _, err := lookupDatabaseContext(name); err != nil { + return "", err + } + return name, nil +} + +// databaseContextCache memoizes config resolution so the HTTP middleware can +// look up a context on every request without re-reading db.json. +var databaseContextCache struct { + mu sync.Mutex + resolved bool + contexts []DatabaseContext + err error +} + +// databaseContexts returns every configured context, the default first and the +// rest sorted by name. It opens no connections: the default context's DSN is +// resolved lazily by the registry, because resolving it can start captain's +// embedded postgres. +func databaseContexts() ([]DatabaseContext, error) { + databaseContextCache.mu.Lock() + defer databaseContextCache.mu.Unlock() + if !databaseContextCache.resolved { + databaseContextCache.contexts, databaseContextCache.err = resolveDatabaseContexts() + databaseContextCache.resolved = true + } + return databaseContextCache.contexts, databaseContextCache.err +} + +func lookupDatabaseContext(name string) (DatabaseContext, error) { + contexts, err := databaseContexts() + if err != nil { + return DatabaseContext{}, err + } + for _, ctx := range contexts { + if ctx.Name == name { + return ctx, nil + } + } + return DatabaseContext{}, fmt.Errorf("%w %q (configured: %s)", errUnknownDatabaseContext, name, + strings.Join(databaseContextNames(contexts), ", ")) +} + +func databaseContextNames(contexts []DatabaseContext) []string { + names := make([]string, 0, len(contexts)) + for _, ctx := range contexts { + names = append(names, ctx.Name) + } + return names +} + +// resetDatabaseContextCache drops memoized config so tests (and a future +// config-reload command) observe a changed environment. +func resetDatabaseContextCache() { + databaseContextCache.mu.Lock() + databaseContextCache.resolved = false + databaseContextCache.contexts = nil + databaseContextCache.err = nil + databaseContextCache.mu.Unlock() +} + +// databaseContextSpec is one context's definition before merging. The default +// context has no spec: it is resolved by captainDSN. +type databaseContextSpec struct { + Name string + DSN string + Label string + Source string + ReadOnly *bool // nil => read-only +} + +func resolveDatabaseContexts() ([]DatabaseContext, error) { + specs, err := configContextSpecs() + if err != nil { + return nil, err + } + envSpecs, err := envContextSpecs() + if err != nil { + return nil, err + } + flagSpecs, err := flagContextSpecs() + if err != nil { + return nil, err + } + // Precedence: --db-url over CAPTAIN_DB_CONTEXTS over db.json. + for _, overlay := range []map[string]databaseContextSpec{envSpecs, flagSpecs} { + for name, spec := range overlay { + specs[name] = spec + } + } + + // The default context's DSN and source are filled in once it is opened; + // resolving them here could start captain's embedded postgres. + contexts := []DatabaseContext{{ + Name: defaultDatabaseContextName, Label: "Monitored database", Default: true, + }} + names := make([]string, 0, len(specs)) + for name := range specs { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + spec := specs[name] + label := spec.Label + if label == "" { + label = name + } + contexts = append(contexts, DatabaseContext{ + Name: name, + Label: label, + DSN: spec.DSN, + Source: spec.Source, + ReadOnly: spec.ReadOnly == nil || *spec.ReadOnly, + }) + } + return contexts, nil +} + +func flagContextSpecs() (map[string]databaseContextSpec, error) { + specs := map[string]databaseContextSpec{} + for _, value := range databaseURLs { + name, dsn, named := splitContextSpec(value) + if !named { + continue + } + if err := validateContextSpec(name, dsn, "--"+databaseURLFlag); err != nil { + return nil, err + } + specs[name] = databaseContextSpec{Name: name, DSN: dsn, Source: "--" + databaseURLFlag} + } + return specs, nil +} + +func envContextSpecs() (map[string]databaseContextSpec, error) { + specs := map[string]databaseContextSpec{} + raw := strings.TrimSpace(os.Getenv(databaseContextsEnv)) + if raw == "" { + return specs, nil + } + for _, entry := range strings.FieldsFunc(raw, func(r rune) bool { return r == ';' || r == '\n' }) { + entry = strings.TrimSpace(entry) + if entry == "" { + continue + } + name, dsn, named := splitContextSpec(entry) + if !named { + return nil, fmt.Errorf("%s entry %q must be name=dsn, with a context name matching %s and a URL-form DSN", + databaseContextsEnv, entry, databaseContextNamePattern) + } + if err := validateContextSpec(name, dsn, databaseContextsEnv); err != nil { + return nil, err + } + specs[name] = databaseContextSpec{Name: name, DSN: dsn, Source: databaseContextsEnv} + } + return specs, nil +} + +// splitContextSpec splits "name=dsn". A value is named only when the text +// before the first "=" is a plausible context name and the remainder is a URL, +// so libpq keyword DSNs ("host=localhost dbname=gavel") stay unnamed. +func splitContextSpec(value string) (name, dsn string, named bool) { + value = strings.TrimSpace(value) + prefix, rest, found := strings.Cut(value, "=") + if !found { + return "", value, false + } + prefix = strings.TrimSpace(prefix) + rest = strings.TrimSpace(rest) + if !databaseContextNamePattern.MatchString(prefix) || !strings.Contains(rest, "://") { + return "", value, false + } + return prefix, rest, true +} + +func validateContextSpec(name, dsn, source string) error { + if name == defaultDatabaseContextName { + return fmt.Errorf("%s: %q is a reserved context name; it always refers to the monitored database", source, defaultDatabaseContextName) + } + if !databaseContextNamePattern.MatchString(name) { + return fmt.Errorf("%s: invalid context name %q, expected %s", source, name, databaseContextNamePattern) + } + if strings.TrimSpace(dsn) == "" { + return fmt.Errorf("%s: context %q has an empty dsn", source, name) + } + return nil +} + +// defaultDatabaseURLOverride returns the unnamed --db-url value, which +// overrides the default context's DSN. +func defaultDatabaseURLOverride() (string, error) { + override := "" + for _, value := range databaseURLs { + if _, dsn, named := splitContextSpec(value); !named { + if override != "" { + return "", fmt.Errorf("--%s may specify at most one unnamed DSN (the default context); use name=dsn for additional contexts", databaseURLFlag) + } + override = dsn + } + } + return override, nil +} diff --git a/pkg/cli/db_context_dsn.go b/pkg/cli/db_context_dsn.go new file mode 100644 index 00000000..3fee2df7 --- /dev/null +++ b/pkg/cli/db_context_dsn.go @@ -0,0 +1,55 @@ +package cli + +import ( + "fmt" + "net/url" + "strings" +) + +// readOnlySessionOption forces every session on a secondary context's pool into +// read-only transactions, so a read path handed a writing query fails at the +// database rather than mutating somebody else's session store. +const readOnlySessionOption = "-c default_transaction_read_only=on" + +// readOnlyDSN returns dsn with the read-only session option merged in. URL-form +// DSNs keep any existing options; keyword-form DSNs that already set options +// are rejected rather than silently mangled — use "readOnly": false in db.json +// for those. +func readOnlyDSN(dsn string) (string, error) { + trimmed := strings.TrimSpace(dsn) + if trimmed == "" { + return "", fmt.Errorf("cannot make an empty DSN read-only") + } + if !strings.Contains(trimmed, "://") { + if keywordDSNHasOptions(trimmed) { + return "", fmt.Errorf("cannot add %q to a keyword DSN that already sets options; set \"readOnly\": false for this context", readOnlySessionOption) + } + return trimmed + " options='" + readOnlySessionOption + "'", nil + } + + parsed, err := url.Parse(trimmed) + if err != nil { + return "", fmt.Errorf("parse DSN: %w", err) + } + query := parsed.Query() + existing := strings.TrimSpace(query.Get("options")) + if strings.Contains(existing, "default_transaction_read_only") { + return trimmed, nil + } + if existing == "" { + query.Set("options", readOnlySessionOption) + } else { + query.Set("options", existing+" "+readOnlySessionOption) + } + parsed.RawQuery = query.Encode() + return parsed.String(), nil +} + +func keywordDSNHasOptions(dsn string) bool { + for _, field := range strings.Fields(dsn) { + if key, _, found := strings.Cut(field, "="); found && strings.EqualFold(strings.TrimSpace(key), "options") { + return true + } + } + return false +} diff --git a/pkg/cli/db_context_dsn_ginkgo_test.go b/pkg/cli/db_context_dsn_ginkgo_test.go new file mode 100644 index 00000000..fd7e562f --- /dev/null +++ b/pkg/cli/db_context_dsn_ginkgo_test.go @@ -0,0 +1,45 @@ +package cli + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Read-only DSN rewriting", func() { + DescribeTable("adds the read-only session option", + func(dsn, expected string) { + Expect(readOnlyDSN(dsn)).To(Equal(expected)) + }, + Entry("URL without a query", + "postgres://reader@prod:5432/gavel", + "postgres://reader@prod:5432/gavel?options=-c+default_transaction_read_only%3Don"), + Entry("URL with an existing query", + "postgres://reader@prod:5432/gavel?sslmode=disable", + "postgres://reader@prod:5432/gavel?options=-c+default_transaction_read_only%3Don&sslmode=disable"), + Entry("URL with existing options preserves them", + "postgres://reader@prod:5432/gavel?options=-c+statement_timeout%3D5000", + "postgres://reader@prod:5432/gavel?options=-c+statement_timeout%3D5000+-c+default_transaction_read_only%3Don"), + Entry("keyword DSN", + "host=prod dbname=gavel user=reader", + "host=prod dbname=gavel user=reader options='-c default_transaction_read_only=on'"), + ) + + It("is idempotent for a DSN that is already read-only", func() { + once, err := readOnlyDSN("postgres://reader@prod:5432/gavel") + Expect(err).NotTo(HaveOccurred()) + + Expect(readOnlyDSN(once)).To(Equal(once)) + }) + + It("rejects a keyword DSN that already sets options", func() { + _, err := readOnlyDSN("host=prod dbname=gavel options='-c statement_timeout=5000'") + + Expect(err).To(MatchError(ContainSubstring(`set "readOnly": false`))) + }) + + It("rejects an empty DSN", func() { + _, err := readOnlyDSN(" ") + + Expect(err).To(MatchError(ContainSubstring("empty DSN"))) + }) +}) diff --git a/pkg/cli/db_context_ginkgo_test.go b/pkg/cli/db_context_ginkgo_test.go new file mode 100644 index 00000000..417a75e4 --- /dev/null +++ b/pkg/cli/db_context_ginkgo_test.go @@ -0,0 +1,226 @@ +package cli + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + . "github.com/onsi/gomega/gstruct" +) + +var _ = Describe("Database contexts", Serial, func() { + // writeGavelDBConfig points HOME at a scratch dir and writes db.json there, + // so context resolution reads a known configuration. + writeGavelDBConfig := func(body string) { + home := GinkgoT().TempDir() + GinkgoT().Setenv("HOME", home) + dir := filepath.Join(home, ".config", "gavel") + Expect(os.MkdirAll(dir, 0o755)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(dir, "db.json"), []byte(body), 0o644)).To(Succeed()) + resetDatabaseContextCache() + } + + BeforeEach(func() { + databaseURLs = nil + databaseContextFlagValue = "" + GinkgoT().Setenv(databaseContextsEnv, "") + GinkgoT().Setenv(databaseContextEnv, "") + GinkgoT().Setenv("HOME", GinkgoT().TempDir()) + resetDatabaseContextCache() + }) + + AfterEach(func() { + databaseURLs = nil + databaseContextFlagValue = "" + resetDatabaseContextCache() + }) + + Describe("configuration sources", func() { + It("declares only the default context when nothing is configured", func() { + contexts, err := databaseContexts() + + Expect(err).NotTo(HaveOccurred()) + Expect(contexts).To(HaveLen(1)) + Expect(contexts[0].Name).To(Equal(defaultDatabaseContextName)) + Expect(contexts[0].Default).To(BeTrue()) + }) + + It("declares only the default context for a legacy db.json", func() { + writeGavelDBConfig(`{"mode":"dsn","dsn":"postgres://legacy/gavel"}`) + + contexts, err := databaseContexts() + + Expect(err).NotTo(HaveOccurred()) + Expect(databaseContextNames(contexts)).To(Equal([]string{defaultDatabaseContextName})) + }) + + It("reads named contexts from db.json", func() { + writeGavelDBConfig(`{"mode":"dsn","dsn":"postgres://local/gavel","contexts":{ + "prod":{"dsn":"postgres://reader@prod/gavel","label":"Production"}, + "box2":{"dsn":"postgres://moshe@box2/gavel","readOnly":false}}}`) + + contexts, err := databaseContexts() + + Expect(err).NotTo(HaveOccurred()) + Expect(databaseContextNames(contexts)).To(Equal([]string{defaultDatabaseContextName, "box2", "prod"})) + prod, err := lookupDatabaseContext("prod") + Expect(err).NotTo(HaveOccurred()) + Expect(prod).To(MatchFields(IgnoreExtras, Fields{ + "Label": Equal("Production"), "DSN": Equal("postgres://reader@prod/gavel"), + "ReadOnly": BeTrue(), "Default": BeFalse(), + })) + box2, err := lookupDatabaseContext("box2") + Expect(err).NotTo(HaveOccurred()) + Expect(box2.ReadOnly).To(BeFalse(), "readOnly:false must survive into the resolved context") + }) + + It("declares contexts from the environment", func() { + GinkgoT().Setenv(databaseContextsEnv, "ci=postgres://ci/gavel;box2=postgres://box2/gavel") + resetDatabaseContextCache() + + contexts, err := databaseContexts() + + Expect(err).NotTo(HaveOccurred()) + Expect(databaseContextNames(contexts)).To(Equal([]string{defaultDatabaseContextName, "box2", "ci"})) + }) + + It("prefers the flag over the environment over db.json for the same name", func() { + writeGavelDBConfig(`{"contexts":{"prod":{"dsn":"postgres://config/gavel"}}}`) + GinkgoT().Setenv(databaseContextsEnv, "prod=postgres://env/gavel") + databaseURLs = []string{"prod=postgres://flag/gavel"} + resetDatabaseContextCache() + + prod, err := lookupDatabaseContext("prod") + + Expect(err).NotTo(HaveOccurred()) + Expect(prod.DSN).To(Equal("postgres://flag/gavel")) + Expect(prod.Source).To(Equal("--" + databaseURLFlag)) + }) + + It("rejects a context named default", func() { + writeGavelDBConfig(`{"contexts":{"default":{"dsn":"postgres://other/gavel"}}}`) + + _, err := databaseContexts() + + Expect(err).To(MatchError(ContainSubstring("reserved context name"))) + }) + + It("rejects an environment entry whose name is not a valid context name", func() { + GinkgoT().Setenv(databaseContextsEnv, "Prod Box=postgres://prod/gavel") + resetDatabaseContextCache() + + _, err := databaseContexts() + + Expect(err).To(MatchError(ContainSubstring("must be name=dsn"))) + }) + + It("rejects a db.json context name that is not a valid context name", func() { + writeGavelDBConfig(`{"contexts":{"Prod Box":{"dsn":"postgres://prod/gavel"}}}`) + + _, err := databaseContexts() + + Expect(err).To(MatchError(ContainSubstring("invalid context name"))) + }) + + It("rejects an empty dsn", func() { + writeGavelDBConfig(`{"contexts":{"prod":{"dsn":" "}}}`) + + _, err := databaseContexts() + + Expect(err).To(MatchError(ContainSubstring(`context "prod" has an empty dsn`))) + }) + + It("reports the configured names when a context is unknown", func() { + GinkgoT().Setenv(databaseContextsEnv, "ci=postgres://ci/gavel") + resetDatabaseContextCache() + + _, err := lookupDatabaseContext("nope") + + Expect(err).To(MatchError(errUnknownDatabaseContext)) + Expect(err.Error()).To(ContainSubstring("configured: default, ci")) + }) + }) + + Describe("--db-url spec parsing", func() { + It("treats a bare URL as the default context's DSN", func() { + databaseURLs = []string{"postgres://flag/captain"} + resetDatabaseContextCache() + + override, err := defaultDatabaseURLOverride() + contexts, contextsErr := databaseContexts() + + Expect(err).NotTo(HaveOccurred()) + Expect(override).To(Equal("postgres://flag/captain")) + Expect(contextsErr).NotTo(HaveOccurred()) + Expect(databaseContextNames(contexts)).To(Equal([]string{defaultDatabaseContextName})) + }) + + It("treats a libpq keyword DSN as unnamed despite its = signs", func() { + databaseURLs = []string{"host=localhost dbname=gavel user=moshe"} + resetDatabaseContextCache() + + override, err := defaultDatabaseURLOverride() + + Expect(err).NotTo(HaveOccurred()) + Expect(override).To(Equal("host=localhost dbname=gavel user=moshe")) + }) + + It("rejects two unnamed DSNs", func() { + databaseURLs = []string{"postgres://one/captain", "postgres://two/captain"} + + _, err := defaultDatabaseURLOverride() + + Expect(err).To(MatchError(ContainSubstring("at most one unnamed DSN"))) + }) + + It("keeps a named DSN out of the default context", func() { + databaseURLs = []string{"prod=postgres://prod/gavel"} + resetDatabaseContextCache() + + override, err := defaultDatabaseURLOverride() + + Expect(err).NotTo(HaveOccurred()) + Expect(override).To(BeEmpty()) + }) + }) + + Describe("active context resolution", func() { + BeforeEach(func() { + GinkgoT().Setenv(databaseContextsEnv, "ctxflag=postgres://flag/gavel;ctxenv=postgres://env/gavel;ctxvalue=postgres://value/gavel") + resetDatabaseContextCache() + }) + + It("defaults when nothing selects a context", func(ctx SpecContext) { + Expect(activeDatabaseContextName(ctx)).To(Equal(defaultDatabaseContextName)) + }) + + It("uses the environment when no flag is set", func(ctx SpecContext) { + GinkgoT().Setenv(databaseContextEnv, "ctxenv") + + Expect(activeDatabaseContextName(ctx)).To(Equal("ctxenv")) + }) + + It("prefers the flag over the environment", func(ctx SpecContext) { + GinkgoT().Setenv(databaseContextEnv, "ctxenv") + databaseContextFlagValue = "ctxflag" + + Expect(activeDatabaseContextName(ctx)).To(Equal("ctxflag")) + }) + + It("prefers the context value over the flag", func(ctx SpecContext) { + GinkgoT().Setenv(databaseContextEnv, "ctxenv") + databaseContextFlagValue = "ctxflag" + + Expect(activeDatabaseContextName(ContextWithDatabaseContext(ctx, "ctxvalue"))).To(Equal("ctxvalue")) + }) + + It("rejects an unknown --context before the command runs", func(ctx SpecContext) { + databaseContextFlagValue = "missing" + + _, err := ResolveDatabaseContextName(ctx) + + Expect(err).To(MatchError(errUnknownDatabaseContext)) + }) + }) +}) diff --git a/pkg/cli/db_context_http.go b/pkg/cli/db_context_http.go new file mode 100644 index 00000000..f48cfe34 --- /dev/null +++ b/pkg/cli/db_context_http.go @@ -0,0 +1,160 @@ +package cli + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + + "github.com/flanksource/captain/pkg/database" +) + +const ( + // databaseContextCookie is how the web UI selects a context. A cookie + // rather than a header because EventSource streams and the embedded chat + // transport cannot set request headers. + databaseContextCookie = "captain_db_context" + // databaseContextHeader lets non-browser clients select a context. + databaseContextHeader = "X-Captain-DB-Context" +) + +// databaseContextError is the machine-readable body the web UI branches on: an +// unknown context clears its stale cookie, a read-only rejection explains why +// the control was refused. +type databaseContextError struct { + Error string `json:"error"` + Code string `json:"code"` + Contexts []string `json:"contexts,omitempty"` +} + +// DatabaseContextMiddleware binds one database context to each request, taken +// from the X-Captain-DB-Context header, else the captain_db_context cookie, +// else the default. An unknown context is rejected rather than silently +// defaulted, and writes are rejected on a read-only context. +func DatabaseContextMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + name := requestDatabaseContextName(r) + dbContext, err := lookupDatabaseContext(name) + if errors.Is(err, errUnknownDatabaseContext) { + names := []string{} + if contexts, listErr := databaseContexts(); listErr == nil { + names = databaseContextNames(contexts) + } + writeDatabaseContextError(w, http.StatusBadRequest, databaseContextError{ + Error: fmt.Sprintf("unknown database context %q", name), Code: "unknown_context", Contexts: names, + }) + return + } + if err != nil { + writeDatabaseContextError(w, http.StatusInternalServerError, databaseContextError{ + Error: err.Error(), Code: "database_context_config", + }) + return + } + if !dbContext.Default && !isReadOnlyMethod(r.Method) { + writeDatabaseContextError(w, http.StatusConflict, databaseContextError{ + Error: fmt.Sprintf("database context %q is read-only; switch to the %q context to write", name, defaultDatabaseContextName), + Code: "read_only_context", + }) + return + } + next.ServeHTTP(w, r.WithContext(ContextWithDatabaseContext(r.Context(), name))) + }) +} + +func requestDatabaseContextName(r *http.Request) string { + if name := strings.TrimSpace(r.Header.Get(databaseContextHeader)); name != "" { + return name + } + if cookie, err := r.Cookie(databaseContextCookie); err == nil { + if name := strings.TrimSpace(cookie.Value); name != "" { + return name + } + } + return defaultDatabaseContextName +} + +func isReadOnlyMethod(method string) bool { + switch method { + case http.MethodGet, http.MethodHead, http.MethodOptions: + return true + default: + return false + } +} + +func writeDatabaseContextError(w http.ResponseWriter, status int, body databaseContextError) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(body); err != nil { + log.Errorf("write database context error: %v", err) + } +} + +// serveRunStatus maps a Run* failure to a status code. A database context that +// is configured but unreachable is a server-side availability problem, not a +// malformed request. +func serveRunStatus(err error, fallback int) int { + if errors.Is(err, errUnknownDatabaseContext) || strings.Contains(err.Error(), "open captain database context") { + return http.StatusServiceUnavailable + } + return fallback +} + +// ContextsResult lists the database contexts this captain can read. +type ContextsResult struct { + Active string `json:"active" pretty:"label=Active"` + Default string `json:"default" pretty:"label=Default"` + Contexts []ContextRow `json:"contexts" pretty:"label=Contexts,table"` +} + +// ContextRow is one configured database context. +type ContextRow struct { + Name string `json:"name" pretty:"label=Name,table"` + Label string `json:"label" pretty:"label=Label,table"` + Source string `json:"source" pretty:"label=Source,table"` + DSN string `json:"dsn" pretty:"label=DSN,table"` + Default bool `json:"default" pretty:"label=Default,table"` + ReadOnly bool `json:"readOnly" pretty:"label=Read Only,table"` + Status string `json:"status,omitempty" pretty:"label=Status,table"` +} + +func handleContexts() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + result, err := RunContexts(r.Context(), ContextsOptions{}) + if err != nil { + http.Error(w, err.Error(), serveRunStatus(err, http.StatusInternalServerError)) + return + } + writeServeJSON(w, http.StatusOK, result) + } +} + +// describeDatabaseContexts renders the configured contexts. The default +// context's DSN is reported only once it has been opened, because resolving it +// can start captain's embedded postgres. +func describeDatabaseContexts(active string) (ContextsResult, error) { + contexts, err := databaseContexts() + if err != nil { + return ContextsResult{}, err + } + result := ContextsResult{Active: active, Default: defaultDatabaseContextName} + for _, dbContext := range contexts { + row := ContextRow{ + Name: dbContext.Name, Label: dbContext.Label, Source: dbContext.Source, + DSN: database.MaskDSN(dbContext.DSN), Default: dbContext.Default, ReadOnly: dbContext.ReadOnly, + } + if openedDSN, openedSource := contextDatabaseIdentity(dbContext.Name); openedSource != "" { + row.Source, row.DSN = openedSource, database.MaskDSN(openedDSN) + } + if row.Label == "" { + row.Label = row.Source + } + if row.Label == "" { + row.Label = row.Name + } + result.Contexts = append(result.Contexts, row) + } + return result, nil +} diff --git a/pkg/cli/db_context_http_ginkgo_test.go b/pkg/cli/db_context_http_ginkgo_test.go new file mode 100644 index 00000000..ada8bc4b --- /dev/null +++ b/pkg/cli/db_context_http_ginkgo_test.go @@ -0,0 +1,110 @@ +package cli + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("DatabaseContextMiddleware", Serial, func() { + const secondary = "prod" + + var ( + observed string + handler http.Handler + ) + + BeforeEach(func() { + databaseURLs = nil + databaseContextFlagValue = "" + GinkgoT().Setenv("HOME", GinkgoT().TempDir()) + GinkgoT().Setenv(databaseContextEnv, "") + GinkgoT().Setenv(databaseContextsEnv, secondary+"=postgres://reader@prod/gavel") + resetDatabaseContextCache() + + observed = "" + handler = DatabaseContextMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + observed = activeDatabaseContextName(r.Context()) + w.WriteHeader(http.StatusOK) + })) + }) + + AfterEach(func() { resetDatabaseContextCache() }) + + // serve runs one request through the middleware. + serve := func(request *http.Request) *httptest.ResponseRecorder { + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, request) + return recorder + } + + decodeError := func(recorder *httptest.ResponseRecorder) databaseContextError { + var body databaseContextError + Expect(json.Unmarshal(recorder.Body.Bytes(), &body)).To(Succeed()) + return body + } + + It("binds the default context when the request selects none", func() { + recorder := serve(httptest.NewRequest(http.MethodGet, "/api/captain/sessions/live", nil)) + + Expect(recorder.Code).To(Equal(http.StatusOK)) + Expect(observed).To(Equal(defaultDatabaseContextName)) + }) + + It("binds the context named by the cookie", func() { + request := httptest.NewRequest(http.MethodGet, "/api/captain/sessions/live", nil) + request.AddCookie(&http.Cookie{Name: databaseContextCookie, Value: secondary}) + + recorder := serve(request) + + Expect(recorder.Code).To(Equal(http.StatusOK)) + Expect(observed).To(Equal(secondary)) + }) + + It("prefers the header over the cookie", func() { + request := httptest.NewRequest(http.MethodGet, "/api/captain/sessions/live", nil) + request.AddCookie(&http.Cookie{Name: databaseContextCookie, Value: defaultDatabaseContextName}) + request.Header.Set(databaseContextHeader, secondary) + + recorder := serve(request) + + Expect(recorder.Code).To(Equal(http.StatusOK)) + Expect(observed).To(Equal(secondary)) + }) + + It("rejects an unknown context with the configured names", func() { + request := httptest.NewRequest(http.MethodGet, "/api/captain/sessions/live", nil) + request.Header.Set(databaseContextHeader, "missing") + + recorder := serve(request) + + Expect(recorder.Code).To(Equal(http.StatusBadRequest)) + Expect(decodeError(recorder)).To(Equal(databaseContextError{ + Error: `unknown database context "missing"`, + Code: "unknown_context", + Contexts: []string{defaultDatabaseContextName, secondary}, + })) + Expect(observed).To(BeEmpty(), "the request must not reach the handler") + }) + + It("rejects a write against a read-only context", func() { + request := httptest.NewRequest(http.MethodPost, "/api/chat/sessions", nil) + request.Header.Set(databaseContextHeader, secondary) + + recorder := serve(request) + + Expect(recorder.Code).To(Equal(http.StatusConflict)) + Expect(decodeError(recorder).Code).To(Equal("read_only_context")) + Expect(observed).To(BeEmpty(), "the request must not reach the handler") + }) + + It("allows a write against the default context", func() { + recorder := serve(httptest.NewRequest(http.MethodPost, "/api/chat/sessions", nil)) + + Expect(recorder.Code).To(Equal(http.StatusOK)) + Expect(observed).To(Equal(defaultDatabaseContextName)) + }) +}) diff --git a/pkg/cli/db_context_registry.go b/pkg/cli/db_context_registry.go new file mode 100644 index 00000000..bbec2b82 --- /dev/null +++ b/pkg/cli/db_context_registry.go @@ -0,0 +1,188 @@ +package cli + +import ( + "context" + "errors" + "fmt" + "sync" + + "github.com/flanksource/captain/pkg/database" + "gorm.io/gorm" +) + +// secondaryMaxOpenConns caps each non-default context's pool. Reads against +// another machine's database are occasional, and a process may hold several +// handles at once, so they get a fraction of the default pool size. +const secondaryMaxOpenConns = 5 + +type captainDatabaseMode uint8 + +const ( + captainDatabaseNoMigrations captainDatabaseMode = iota + captainDatabaseWithMigrations +) + +// databaseHandle memoizes one context's connection. Only successful opens are +// memoized: a transient failure must not poison a long-lived serve process. +type databaseHandle struct { + db *database.DB + dsn string + source string + migrated bool +} + +// databaseRegistry holds one handle per context name. The default context is +// what used to be the process-wide singleton; every other entry is read-only. +var databaseRegistry = struct { + mu sync.Mutex + handles map[string]*databaseHandle +}{handles: map[string]*databaseHandle{}} + +// openContextDB resolves and memoizes the handle for one context. Migrations +// are rejected for every context but the default, so no code path can migrate +// a database captain only reads. +func openContextDB(ctx context.Context, name string, mode captainDatabaseMode) (*database.DB, error) { + databaseRegistry.mu.Lock() + defer databaseRegistry.mu.Unlock() + + if handle, ok := databaseRegistry.handles[name]; ok { + if mode == captainDatabaseWithMigrations && !handle.migrated { + return nil, errors.New("captain serve cannot migrate after the process database was opened without migrations") + } + return handle.db, nil + } + if mode == captainDatabaseWithMigrations && name != defaultDatabaseContextName { + return nil, fmt.Errorf("captain only migrates the %q database context, not %q", defaultDatabaseContextName, name) + } + + dsn, source, err := contextDSN(name) + if err != nil { + return nil, err + } + options := []database.Option{database.WithDSN(dsn)} + if mode == captainDatabaseWithMigrations { + options = append(options, database.WithMigrations()) + } + if name != defaultDatabaseContextName { + options = append(options, database.WithMaxOpenConns(secondaryMaxOpenConns)) + } + log.Debugf("captain database context %q using %s", name, source) + db, err := database.Open(ctx, options...) + if err != nil { + return nil, fmt.Errorf("open captain database context %q (%s): %w", name, source, err) + } + databaseRegistry.handles[name] = &databaseHandle{ + db: db, + dsn: dsn, + source: source, + migrated: mode == captainDatabaseWithMigrations, + } + return db, nil +} + +// contextDSN resolves a context's connection string. The default keeps +// captain's established flag/env/db.json/embedded precedence; every other +// context is declared explicitly and opened read-only. +func contextDSN(name string) (dsn, source string, err error) { + if name == defaultDatabaseContextName { + return captainDSN() + } + dbContext, err := lookupDatabaseContext(name) + if err != nil { + return "", "", err + } + if !dbContext.ReadOnly { + return dbContext.DSN, dbContext.Source, nil + } + readOnly, err := readOnlyDSN(dbContext.DSN) + if err != nil { + return "", "", fmt.Errorf("database context %q: %w", name, err) + } + return readOnly, dbContext.Source, nil +} + +func contextDatabaseIdentity(name string) (dsn, source string) { + databaseRegistry.mu.Lock() + defer databaseRegistry.mu.Unlock() + handle, ok := databaseRegistry.handles[name] + if !ok { + return "", "" + } + return handle.dsn, handle.source +} + +// ConfigureNativeDatabase injects a host-owned GORM pool as the default +// database context 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. +// +// It deliberately does not resolve captain's context configuration: a host that +// never selects a secondary context must not fail to boot over a malformed +// db.json. Secondary contexts, if any are ever requested, open their own pools +// outside the host's. +func ConfigureNativeDatabase(gormDB *gorm.DB) error { + db, err := database.Use(gormDB) + if err != nil { + return err + } + + databaseRegistry.mu.Lock() + defer databaseRegistry.mu.Unlock() + if handle, ok := databaseRegistry.handles[defaultDatabaseContextName]; ok { + if handle.db != nil && handle.db.Gorm() == gormDB { + return nil + } + return fmt.Errorf("native Captain database is already configured with a different pool") + } + databaseRegistry.handles[defaultDatabaseContextName] = &databaseHandle{ + db: db, + source: "host-provided database", + migrated: true, + } + return nil +} + +// testDatabaseHandle describes a handle injected by a test. +type testDatabaseHandle struct { + Name string + DB *database.DB + // Source is the provenance string surfaced by contextDatabaseIdentity. + Source string + // Unmigrated installs the handle as if it had been opened without + // migrations, so tests can exercise the serve-after-plain-open guard. + Unmigrated bool +} + +// setCaptainDBForTest injects (or, with nil, resets) the default context's +// handle so tests run against their own embedded database instead of a +// configured DSN. Production code never calls this. +func setCaptainDBForTest(db *database.DB) { + setCaptainContextDBForTest(testDatabaseHandle{Name: defaultDatabaseContextName, DB: db}) +} + +// setCaptainContextDBForTest injects a handle under an arbitrary context name so +// tests can exercise context switching without a second postgres server. +func setCaptainContextDBForTest(handle testDatabaseHandle) { + databaseRegistry.mu.Lock() + defer databaseRegistry.mu.Unlock() + if handle.DB == nil { + delete(databaseRegistry.handles, handle.Name) + return + } + databaseRegistry.handles[handle.Name] = &databaseHandle{ + db: handle.DB, + source: handle.Source, + migrated: !handle.Unmigrated, + } +} + +// resetCaptainContextsForTest drops every memoized handle and the cached +// context configuration. +func resetCaptainContextsForTest() { + databaseRegistry.mu.Lock() + databaseRegistry.handles = map[string]*databaseHandle{} + databaseRegistry.mu.Unlock() + resetDatabaseContextCache() +} diff --git a/pkg/cli/db_context_registry_ginkgo_test.go b/pkg/cli/db_context_registry_ginkgo_test.go new file mode 100644 index 00000000..ca128936 --- /dev/null +++ b/pkg/cli/db_context_registry_ginkgo_test.go @@ -0,0 +1,160 @@ +package cli + +import ( + "context" + "os" + "path/filepath" + "time" + + "github.com/flanksource/captain/pkg/database" + "github.com/flanksource/captain/pkg/monitor" + "github.com/flanksource/commons-db/dbtest" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Database context registry", Serial, func() { + const secondaryContext = "secondary" + + var ( + defaultDB *database.DB + secondaryDB *database.DB + discovered int + ) + + // openLeasedDB leases an isolated migrated database from the shared test + // server, so two contexts can be exercised without a second postgres. + openLeasedDB := func(name string) *database.DB { + handle := dbtest.ForGinkgo(dbtest.Options{Name: name}) + db, err := database.Open(GinkgoT().Context(), database.WithDSN(handle.DSN()), database.WithMigrations()) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { Expect(db.Close()).To(Succeed()) }) + return db + } + + // seedLiveSession records one session and its running process, so the two + // databases are distinguishable by their contents alone. session live lists + // only sessions with an active process. + seedLiveSession := func(ctx context.Context, db *database.DB, project, providerSessionID string) { + session, err := db.CreateOrGetSession(ctx, database.CreateSessionInput{ + Source: "codex", ProviderSessionID: providerSessionID, CWD: project, + }) + Expect(err).NotTo(HaveOccurred()) + started := time.Now().UTC().Add(-time.Minute).Truncate(time.Second) + Expect(db.UpsertSessionProcess(ctx, database.SessionProcessInput{ + SessionID: session.ID, HostID: captainHostID(), BootID: "boot", PID: 24680, + ProcessStartedAt: started, SampledAt: started.Add(30 * time.Second), + Status: "sleeping", CWD: project, Source: "codex", + })).To(Succeed()) + } + + BeforeEach(func(ctx SpecContext) { + databaseURLs = nil + databaseContextFlagValue = "" + home := GinkgoT().TempDir() + GinkgoT().Setenv("HOME", home) + GinkgoT().Setenv(databaseContextEnv, "") + GinkgoT().Setenv(databaseContextsEnv, secondaryContext+"=postgres://unused/never-opened") + resetCaptainContextsForTest() + + Expect(os.MkdirAll(filepath.Join(home, "work", "project"), 0o755)).To(Succeed()) + GinkgoT().Chdir(filepath.Join(home, "work", "project")) + // Seed against the resolved working directory: on macOS the temp dir is + // reached through a symlink, and session scoping matches cwd exactly. + project, err := os.Getwd() + Expect(err).NotTo(HaveOccurred()) + + defaultDB = openLeasedDB("captain_ctx_default") + secondaryDB = openLeasedDB("captain_ctx_secondary") + seedLiveSession(ctx, defaultDB, project, "default-session") + seedLiveSession(ctx, secondaryDB, project, "secondary-session") + + // Inject both handles so the registry answers from them rather than + // dialing the placeholder DSN above. + setCaptainContextDBForTest(testDatabaseHandle{Name: defaultDatabaseContextName, DB: defaultDB}) + setCaptainContextDBForTest(testDatabaseHandle{Name: secondaryContext, DB: secondaryDB, Source: "test secondary"}) + + discovered = 0 + monitorDiscoverProcesses = func() ([]monitor.Process, error) { + discovered++ + return nil, nil + } + DeferCleanup(func() { + monitorDiscoverProcesses = nil + resetCaptainContextsForTest() + }) + }) + + It("reads the default context when nothing selects one", func(ctx SpecContext) { + result, err := RunSessionLive(ctx, SessionLiveOptions{Source: "all", All: true, Limit: 10}) + + Expect(err).NotTo(HaveOccurred()) + Expect(sessionProviderIDs(result)).To(ConsistOf("default-session")) + }) + + It("reads the context bound to the request context", func(ctx SpecContext) { + result, err := RunSessionLive(ContextWithDatabaseContext(ctx, secondaryContext), SessionLiveOptions{Source: "all", All: true, Limit: 10}) + + Expect(err).NotTo(HaveOccurred()) + Expect(sessionProviderIDs(result)).To(ConsistOf("secondary-session")) + }) + + It("reports the active context's database identity, not the default's", func(ctx SpecContext) { + result, err := RunSessionLive(ContextWithDatabaseContext(ctx, secondaryContext), SessionLiveOptions{Source: "all", All: true, Limit: 10}) + + Expect(err).NotTo(HaveOccurred()) + Expect(result.Database.Source).To(Equal("test secondary")) + }) + + It("keeps writes on the default context while a secondary is active", func(ctx SpecContext) { + db, err := captainDefaultDB(ContextWithDatabaseContext(ctx, secondaryContext)) + + Expect(err).NotTo(HaveOccurred()) + Expect(db).To(BeIdenticalTo(defaultDB)) + }) + + It("runs the monitor pass when freshening the default context", func(ctx SpecContext) { + _, err := freshenSessionDB(ctx) + + Expect(err).NotTo(HaveOccurred()) + Expect(discovered).To(BeNumerically(">", 0)) + }) + + It("never writes to a secondary context when freshening it", func(ctx SpecContext) { + db, err := freshenSessionDB(ContextWithDatabaseContext(ctx, secondaryContext)) + + Expect(err).NotTo(HaveOccurred()) + Expect(db).To(BeIdenticalTo(secondaryDB)) + Expect(discovered).To(Equal(0), "a read of another machine's database must not run a monitor pass") + }) + + It("refuses to migrate a secondary context", func(ctx SpecContext) { + setCaptainContextDBForTest(testDatabaseHandle{Name: secondaryContext, DB: nil}) + + _, err := openContextDB(ctx, secondaryContext, captainDatabaseWithMigrations) + + Expect(err).To(MatchError(ContainSubstring(`captain only migrates the "default" database context`))) + }) + + It("retries an open that failed rather than memoizing the error", func(ctx SpecContext) { + setCaptainContextDBForTest(testDatabaseHandle{Name: secondaryContext, DB: nil}) + + _, first := openContextDB(ctx, secondaryContext, captainDatabaseNoMigrations) + Expect(first).To(HaveOccurred()) + + setCaptainContextDBForTest(testDatabaseHandle{Name: secondaryContext, DB: secondaryDB}) + db, err := openContextDB(ctx, secondaryContext, captainDatabaseNoMigrations) + + Expect(err).NotTo(HaveOccurred()) + Expect(db).To(BeIdenticalTo(secondaryDB)) + }) +}) + +// sessionProviderIDs identifies which database answered a read. +func sessionProviderIDs(result SessionLiveResult) []string { + ids := make([]string, 0, len(result.Sessions)) + for _, session := range result.Sessions { + ids = append(ids, session.ID) + } + return ids +} diff --git a/pkg/cli/db_context_stores.go b/pkg/cli/db_context_stores.go new file mode 100644 index 00000000..705e55f1 --- /dev/null +++ b/pkg/cli/db_context_stores.go @@ -0,0 +1,50 @@ +package cli + +import ( + "context" + "sync" + + "github.com/flanksource/captain/pkg/aichat" +) + +// chatThreadStores memoizes one chat thread store per database context, so a +// request reading a secondary context lists that database's threads instead of +// the monitored one's. +var chatThreadStores struct { + mu sync.Mutex + byName map[string]aichat.ThreadStore +} + +// contextThreadStore resolves the chat thread store for the request's database +// context. Thread writes never reach it on a secondary context: the database +// context middleware rejects unsafe methods first. +func contextThreadStore(ctx context.Context) (aichat.ThreadStore, error) { + name := activeDatabaseContextName(ctx) + + chatThreadStores.mu.Lock() + store, ok := chatThreadStores.byName[name] + chatThreadStores.mu.Unlock() + if ok { + return store, nil + } + + db, err := openContextDB(ctx, name, captainDatabaseNoMigrations) + if err != nil { + return nil, err + } + store, err = aichat.NewDatabaseThreadStore(db) + if err != nil { + return nil, err + } + + chatThreadStores.mu.Lock() + defer chatThreadStores.mu.Unlock() + if existing, ok := chatThreadStores.byName[name]; ok { + return existing, nil + } + if chatThreadStores.byName == nil { + chatThreadStores.byName = map[string]aichat.ThreadStore{} + } + chatThreadStores.byName[name] = store + return store, nil +} diff --git a/pkg/cli/event_renderer.go b/pkg/cli/event_renderer.go index a4b2120d..626a7388 100644 --- a/pkg/cli/event_renderer.go +++ b/pkg/cli/event_renderer.go @@ -1,17 +1,160 @@ package cli import ( + "errors" + "fmt" + "io" "os" + "github.com/charmbracelet/x/ansi" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/session" + "golang.org/x/term" ) -// 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) +type EventRenderer struct { + output io.Writer + interactive bool + accumulator *promptEventAccumulator + pending *session.Message + rendered map[string]bool + err error + iteration int + hasIter bool +} + +func NewEventRenderer(output *os.File) *EventRenderer { + return newEventRenderer(output, output != nil && term.IsTerminal(int(output.Fd()))) +} + +func newEventRenderer(output io.Writer, interactive bool) *EventRenderer { + renderer := &EventRenderer{ + output: output, + interactive: interactive, + rendered: map[string]bool{}, + } + renderer.accumulator = newPromptEventAccumulator(renderer.consume, discardTaskSink{}, "", "") + if cwd, err := os.Getwd(); err == nil { + renderer.accumulator.cwd = cwd + } + return renderer +} + +func (r *EventRenderer) Handle(iteration int, event ai.Event) { + if r.hasIter && iteration != r.iteration { + r.flushPending() + r.accumulator.resetFrame() + clear(r.rendered) + } + r.iteration, r.hasIter = iteration, true + + if r.pendingBoundary(event.Kind) { + r.flushPending() + } + r.accumulator.handle(iteration, event) + if event.Kind == ai.EventError || event.Kind == ai.EventResult { + r.flushPending() + } +} + +func (r *EventRenderer) Flush() error { + r.flushPending() + return r.err +} + +func (r *EventRenderer) pendingBoundary(kind ai.EventKind) bool { + if r.pending == nil || len(r.pending.Parts) == 0 { + return false + } + pendingType := r.pending.Parts[0].Type + switch kind { + case ai.EventText: + return pendingType != session.PartText + case ai.EventThinking: + return pendingType != session.PartReasoning + default: + return true + } +} + +func (r *EventRenderer) consume(message session.Message) { + if len(message.Parts) == 0 { + return } + part := message.Parts[0] + switch part.Type { + case session.PartText, session.PartReasoning: + if r.pending != nil && r.pending.ID != message.ID { + r.flushPending() + } + copy := message + r.pending = © + if r.interactive { + r.redrawPending() + } + case session.PartTool: + if part.ToolName == "" { + r.err = errors.Join(r.err, fmt.Errorf("tool result for call %q has no matching tool use", part.ToolCallID)) + return + } + if r.rendered[message.ID] { + return + } + r.rendered[message.ID] = true + r.renderMessage(message) + } +} + +func (r *EventRenderer) redrawPending() { + if r.pending == nil { + return + } + text, ok := transcriptMessageANSI(*r.pending) + if !ok { + return + } + r.write("\r" + ansi.EraseEntireLine + text) +} + +func (r *EventRenderer) flushPending() { + if r.pending == nil { + return + } + if r.interactive { + r.write("\n") + } else { + r.renderMessage(*r.pending) + } + r.pending = nil +} + +func (r *EventRenderer) renderMessage(message session.Message) { + text, ok := transcriptMessageANSI(message) + if ok { + r.write(text + "\n") + } +} + +func transcriptMessageANSI(message session.Message) (string, bool) { + rows := (&session.Session{Messages: []session.Message{message}}).TranscriptRows() + if len(rows) == 0 { + return "", false + } + return rows[0].Pretty().ANSI(), true } + +func (r *EventRenderer) write(value string) { + if r.output == nil || r.err != nil { + return + } + _, err := io.WriteString(r.output, value) + r.err = errors.Join(r.err, err) +} + +type discardTaskSink struct{} + +func (discardTaskSink) SetDescription(string) {} +func (discardTaskSink) SetProgress(int, int) {} +func (discardTaskSink) Infof(string, ...interface{}) {} +func (discardTaskSink) Warnf(string, ...interface{}) {} +func (discardTaskSink) Errorf(string, ...interface{}) {} diff --git a/pkg/cli/event_renderer_ginkgo_test.go b/pkg/cli/event_renderer_ginkgo_test.go index 0a149384..d5063a32 100644 --- a/pkg/cli/event_renderer_ginkgo_test.go +++ b/pkg/cli/event_renderer_ginkgo_test.go @@ -1,50 +1,120 @@ package cli import ( - "io" - "os" + "bytes" "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()) + It("buffers captured deltas and emits canonical transcript rows", func() { + var output bytes.Buffer + renderer := newEventRenderer(&output, false) - 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{ + renderer.Handle(0, ai.Event{Kind: ai.EventText, Text: "a keyed "}) + renderer.Handle(0, ai.Event{Kind: ai.EventText, Text: "HMAC"}) + Expect(output.String()).To(BeEmpty()) + + renderer.Handle(0, ai.Event{ Kind: ai.EventToolUse, Tool: "Bash", - Input: map[string]any{"command": "pwd"}, + Input: map[string]any{"command": `/bin/zsh -lc 'pnpm test'`}, 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", - }, }) + renderer.Handle(0, ai.Event{ + Kind: ai.EventToolResult, + Tool: "Bash", + Text: "captured command output", + Success: true, + ToolCallID: "cmd-1", + }) + renderer.Handle(0, ai.Event{Kind: ai.EventText, Text: "done"}) + renderer.Handle(0, ai.Event{Kind: ai.EventResult, Success: true}) + Expect(renderer.Flush()).To(Succeed()) + + text := output.String() + Expect(text).To(And( + ContainSubstring("a keyed HMAC"), + ContainSubstring("zsh"), + ContainSubstring("pnpm test"), + ContainSubstring("done"), + Not(ContainSubstring("/bin/zsh -lc")), + Not(ContainSubstring("captured command output")), + Not(ContainSubstring("[tool-result]")), + )) + Expect(strings.Count(text, "pnpm test")).To(Equal(1)) + }) + + It("redraws an in-flight TTY message and finalizes it once", func() { + var output bytes.Buffer + renderer := newEventRenderer(&output, true) + + renderer.Handle(0, ai.Event{Kind: ai.EventText, Text: "hello "}) + renderer.Handle(0, ai.Event{Kind: ai.EventText, Text: "world"}) + Expect(renderer.Flush()).To(Succeed()) + + Expect(output.String()).To(And( + ContainSubstring("\r\x1b[2K"), + ContainSubstring("hello world"), + HaveSuffix("\n"), + )) + }) + + It("prints a hook notice on its own line, outside the model's prose", func() { + var output bytes.Buffer + renderer := newEventRenderer(&output, false) + + renderer.Handle(0, ai.Event{Kind: ai.EventText, Text: "applying the fix"}) + renderer.Handle(0, ai.Event{Kind: ai.EventSystem, Text: "[post-turn] committed abc1234: fix: the thing"}) + renderer.Handle(0, ai.Event{Kind: ai.EventResult, Success: true}) + Expect(renderer.Flush()).To(Succeed()) + + text := output.String() + Expect(text).To(ContainSubstring("[post-turn] committed abc1234")) + // On its own line: a commit landing inside the sentence the model was + // mid-way through writing is exactly what the buffering must prevent. + Expect(text).NotTo(ContainSubstring("applying the fix[post-turn]")) + Expect(strings.Count(text, "[post-turn] committed abc1234")).To(Equal(1)) + }) + + It("reports an unmatched result without dumping its payload", func() { + var output bytes.Buffer + renderer := newEventRenderer(&output, false) + + renderer.Handle(0, ai.Event{ + Kind: ai.EventToolResult, + ToolCallID: "missing-call", + Text: "sensitive payload", + Success: false, + }) + + err := renderer.Flush() + Expect(err).To(MatchError(ContainSubstring("missing-call"))) + Expect(output.String()).NotTo(ContainSubstring("sensitive payload")) + }) - Expect(writer.Close()).To(Succeed()) - output, err := io.ReadAll(reader) - Expect(err).NotTo(HaveOccurred()) - Expect(reader.Close()).To(Succeed()) + It("scopes tool-call identities to an iteration", func() { + var output bytes.Buffer + renderer := newEventRenderer(&output, false) + + for iteration, command := range []string{"first command", "second command"} { + renderer.Handle(iteration, ai.Event{ + Kind: ai.EventToolUse, + Tool: "Bash", + Input: map[string]any{"command": command}, + ToolCallID: "reused-call", + }) + } - 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]")) + Expect(renderer.Flush()).To(Succeed()) + Expect(output.String()).To(And( + ContainSubstring("first command"), + ContainSubstring("second command"), + )) }) }) diff --git a/pkg/cli/gavel_dsn.go b/pkg/cli/gavel_dsn.go index 1bd44d4e..556e18cf 100644 --- a/pkg/cli/gavel_dsn.go +++ b/pkg/cli/gavel_dsn.go @@ -36,6 +36,42 @@ func sessionDBDir() (string, error) { type gavelDBConfig struct { Mode string `json:"mode"` DSN string `json:"dsn,omitempty"` + // Contexts are additional read-only databases captain can be pointed at. + // They never participate in the default context's DSN resolution. + Contexts map[string]gavelContextConfig `json:"contexts,omitempty"` +} + +type gavelContextConfig struct { + DSN string `json:"dsn"` + Label string `json:"label,omitempty"` + // ReadOnly defaults to true. Set it false for DSNs behind a pooler that + // rejects the read-only session option; captain still never writes to a + // non-default context. + ReadOnly *bool `json:"readOnly,omitempty"` +} + +// configContextSpecs resolves the named contexts declared in db.json. Missing +// or legacy config files simply declare none. +func configContextSpecs() (map[string]databaseContextSpec, error) { + cfg, path, err := loadGavelDBConfig() + if err != nil { + return nil, err + } + specs := make(map[string]databaseContextSpec, len(cfg.Contexts)) + for name, entry := range cfg.Contexts { + source := path + "#" + name + if err := validateContextSpec(name, entry.DSN, path); err != nil { + return nil, err + } + specs[name] = databaseContextSpec{ + Name: name, + DSN: strings.TrimSpace(entry.DSN), + Label: entry.Label, + Source: source, + ReadOnly: entry.ReadOnly, + } + } + return specs, nil } // gavelConfiguredSessionDSN resolves a gavel-shared database from diff --git a/pkg/cli/har_test.go b/pkg/cli/har_test.go new file mode 100644 index 00000000..c941bd66 --- /dev/null +++ b/pkg/cli/har_test.go @@ -0,0 +1,133 @@ +package cli + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/flanksource/commons/har" + "github.com/flanksource/commons/properties" +) + +// withProperty sets a -P property for one test. Properties are process-global, +// so every test that sets one must restore it. +func withProperty(t *testing.T, key, value string) { + t.Helper() + properties.Set(key, value) + t.Cleanup(func() { properties.Set(key, "") }) +} + +// withCleanRegistry gives a test its own collectors. The package registry is +// process-wide because http.DefaultTransport is, and it retains collectors so a +// second Flush rewrites rather than truncates — which across tests would mean +// one test flushing another's archive into a deleted temp dir. +func withCleanRegistry(t *testing.T) { + t.Helper() + previous := harRegistry + harRegistry = har.NewRegistry(nil) + t.Cleanup(func() { harRegistry = previous }) +} + +// captureThrough issues one request against a JSON server through the HAR +// transport captain installs, and returns the entries written to path. +func captureThrough(t *testing.T, path string) []har.Entry { + t.Helper() + withCleanRegistry(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[]}`)) + })) + t.Cleanup(server.Close) + + transport, err := harRegistry.Transport(harFeature, http.DefaultTransport) + if err != nil { + t.Fatalf("Transport: %v", err) + } + req, err := http.NewRequest(http.MethodGet, server.URL+"/v1/models", nil) + if err != nil { + t.Fatal(err) + } + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + + FlushHAR() + + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("no HAR written to %s: %v", path, err) + } + var file har.File + if err := json.Unmarshal(data, &file); err != nil { + t.Fatalf("HAR is not valid JSON: %v", err) + } + return file.Log.Entries +} + +// TestHARDisabledWithoutProperty pins the default: captain must not pay for HAR +// capture, and the transport chain must be left exactly as it was. +func TestHARDisabledWithoutProperty(t *testing.T) { + base := http.DefaultTransport + transport, err := harRegistry.Transport(harFeature, base) + if err != nil { + t.Fatalf("Transport: %v", err) + } + if transport != base { + t.Error("transport must be unchanged when http.har is unset") + } +} + +// TestHARCapturesProviderTraffic is the end-to-end contract of -Phttp.har: a +// request through captain's default transport lands in a readable archive. +func TestHARCapturesProviderTraffic(t *testing.T) { + path := filepath.Join(t.TempDir(), "trace.har") + withProperty(t, "http.har", path) + + entries := captureThrough(t, path) + + if len(entries) != 1 { + t.Fatalf("expected 1 entry, got %d", len(entries)) + } + if !strings.HasSuffix(entries[0].Request.URL, "/v1/models") { + t.Errorf("unexpected captured URL: %s", entries[0].Request.URL) + } + if entries[0].Response.Content.Text != `{"data":[]}` { + t.Errorf("expected the response body, got %q", entries[0].Response.Content.Text) + } +} + +// TestHARFeaturePropertyOverridesGlobal covers -Phttp.captain.har, which lets a +// user capture captain's own traffic without touching a shared http.har. +func TestHARFeaturePropertyOverridesGlobal(t *testing.T) { + dir := t.TempDir() + featurePath := filepath.Join(dir, "captain.har") + withProperty(t, "http.har", filepath.Join(dir, "global.har")) + withProperty(t, "http."+harFeature+".har", featurePath) + + if entries := captureThrough(t, featurePath); len(entries) != 1 { + t.Fatalf("expected 1 entry in the feature archive, got %d", len(entries)) + } + if _, err := os.Stat(filepath.Join(dir, "global.har")); !os.IsNotExist(err) { + t.Error("the global archive must not be written when a feature override is set") + } +} + +// TestHARInvalidLevelFailsFast covers the reason EnableHTTPWireLogging returns +// an error: a typo'd level must stop the run, not capture the wrong thing. +func TestHARInvalidLevelFailsFast(t *testing.T) { + withProperty(t, "http.har", filepath.Join(t.TempDir(), "trace.har")) + withProperty(t, "http.har.level", "verbose") + + if _, err := harRegistry.Transport(harFeature, http.DefaultTransport); err == nil { + t.Fatal("expected an error for an unrecognised http.har.level") + } +} diff --git a/pkg/cli/history.go b/pkg/cli/history.go index ead5c1e5..1a91789b 100644 --- a/pkg/cli/history.go +++ b/pkg/cli/history.go @@ -484,6 +484,7 @@ func runHistoryAll(tl []tools.Tool, opts HistoryOptions, classifier *bash.Catego for _, t := range filtered { base := t.Base() + transcriptRow := session.NewTranscriptRow(t) result.Total++ approved := approvedStatus(t) @@ -500,8 +501,8 @@ func runHistoryAll(tl []tools.Tool, opts HistoryOptions, classifier *bash.Catego row := session.ScanResultRow{ Project: projectName, Tool: t.Name(), - Summary: firstLine(t.Pretty().String()), - Subject: t.Pretty(), + Summary: firstLine(transcriptRow.Pretty().String()), + Subject: transcriptRow.Pretty(), Detail: session.BuildRowDetail(t, session.RowOptions{Cost: opts.Cost, Raw: opts.Raw}), Paths: FormatPathsWithIcons(analysis.ReadPaths, analysis.WritePaths), ReadPaths: analysis.ReadPaths, @@ -543,6 +544,7 @@ func runHistorySingle(tl []tools.Tool, opts HistoryOptions, classifier *bash.Cat for _, t := range filtered { base := t.Base() + transcriptRow := session.NewTranscriptRow(t) if result.Project == "" && base.ProjectRoot != "" { result.Project = filepath.Base(base.ProjectRoot) } @@ -557,8 +559,8 @@ func runHistorySingle(tl []tools.Tool, opts HistoryOptions, classifier *bash.Cat analysis := AnalyzeToolUse(t) row := session.ScanResultRowSingle{ Tool: t.Name(), - Summary: firstLine(t.Pretty().String()), - Subject: t.Pretty(), + Summary: firstLine(transcriptRow.Pretty().String()), + Subject: transcriptRow.Pretty(), Detail: session.BuildRowDetail(t, session.RowOptions{Cost: opts.Cost, Raw: opts.Raw}), Paths: FormatPathsWithIcons(analysis.ReadPaths, analysis.WritePaths), ReadPaths: analysis.ReadPaths, diff --git a/pkg/cli/history_render.go b/pkg/cli/history_render.go index 424f7245..269eeb7b 100644 --- a/pkg/cli/history_render.go +++ b/pkg/cli/history_render.go @@ -7,6 +7,7 @@ import ( "strings" "github.com/flanksource/captain/pkg/claude/tools" + "github.com/flanksource/captain/pkg/session" "github.com/flanksource/clicky" "golang.org/x/term" ) @@ -43,8 +44,8 @@ func termWidth() int { // lineRenderer prints tool history rows to an io.Writer, emitting a synthetic // session-start banner whenever the (source, session, model, effort) key -// changes. Both `captain history` (batched) and `captain ai prompt` -// (streaming) drive the same renderer so output stays consistent. +// changes. Row content comes from session.TranscriptRow; this type only adds +// history's time, tool-name, usage, and session-boundary columns. type lineRenderer struct { w io.Writer width int @@ -72,7 +73,7 @@ func (r *lineRenderer) Render(t tools.Tool, compact bool) { r.prevKey = key r.hasPrev = true } - e := toLineEntry(t, compact, r.width, r.toolWidth) + e := toLineEntry(session.NewTranscriptRow(t), compact, r.width, r.toolWidth) printLeftTo(r.w, e, r.toolWidth) } @@ -161,14 +162,15 @@ func capitalize(s string) string { return strings.ToUpper(s[:1]) + s[1:] } -func toLineEntry(t tools.Tool, compact bool, width, toolWidth int) lineEntry { +func toLineEntry(row session.TranscriptRow, compact bool, width, toolWidth int) lineEntry { + t := row.Tool() base := t.Base() name := t.Name() e := lineEntry{ Tool: name, Time: base.PrettyTimestamp(), Denied: base.Denied && name != "Plan" && name != "User", - Command: t.Pretty().ANSI(), + Command: row.Pretty().ANSI(), } if base.IsSidechain { e.Command = sidechainBadge(base) + e.Command diff --git a/pkg/cli/history_render_test.go b/pkg/cli/history_render_test.go index a9c4b203..284190b5 100644 --- a/pkg/cli/history_render_test.go +++ b/pkg/cli/history_render_test.go @@ -1,11 +1,8 @@ package cli import ( - "bytes" - "strings" "testing" - "github.com/flanksource/captain/pkg/ai" "github.com/flanksource/captain/pkg/claude/tools" ) @@ -100,39 +97,6 @@ func TestLastSessionTools_TrimsToFinalSession(t *testing.T) { } } -func TestRenderResultEvent_RoutesThroughLineRenderer(t *testing.T) { - var buf bytes.Buffer - r := newLineRenderer(&buf, 8) - renderResultEvent(r, ai.Event{ - Kind: ai.EventResult, - Model: "claude-opus-4-7", - Success: true, - CostUSD: 0.0123, - Usage: &ai.Usage{InputTokens: 100, OutputTokens: 200}, - Input: map[string]any{"num_turns": float64(3), "duration_ms": float64(1500)}, - }) - - out := buf.String() - for _, want := range []string{"result", "$0.0123", "turns=3", "1.5s"} { - if !strings.Contains(out, want) { - t.Errorf("rendered output missing %q\nfull output:\n%s", want, out) - } - } -} - -func TestRenderResultEvent_FailureMarksError(t *testing.T) { - var buf bytes.Buffer - r := newLineRenderer(&buf, 8) - renderResultEvent(r, ai.Event{ - Kind: ai.EventResult, - Success: false, - Error: "timeout", - }) - if !strings.Contains(buf.String(), "ERROR") { - t.Errorf("failure result must include ERROR marker, got: %q", buf.String()) - } -} - func TestShortSessionID(t *testing.T) { tests := []struct{ in, want string }{ {"", ""}, diff --git a/pkg/cli/logging.go b/pkg/cli/logging.go index 4764ae3b..630a9235 100644 --- a/pkg/cli/logging.go +++ b/pkg/cli/logging.go @@ -2,7 +2,11 @@ package cli import ( "net/http" + "sync" + "github.com/flanksource/commons/har" + commonshttp "github.com/flanksource/commons/http" + "github.com/flanksource/commons/http/middlewares" "github.com/flanksource/commons/logger" ) @@ -10,30 +14,104 @@ import ( // global -v/--log-level and can be tuned independently with -Plog.level.cli=debug. var log = logger.GetLogger("cli") -// EnableHTTPWireLogging installs commons' HTTP request/response logger on the -// default transport when the "http" logger is verbose enough. Captain's API -// providers (genkit plugins) issue requests through http.DefaultClient, so -// wrapping the default transport captures their outbound calls with sensitive -// headers (Authorization, Cookie, ...) redacted by commons. +var httpLoggingOnce sync.Once + +// EnableHTTPWireLogging installs commons' HTTP trace middleware on the default +// transport. Captain's API providers (genkit plugins) and its own fetchers +// (pkg/ai/models_remote.go, pkg/ai/pricing/openrouter.go) all issue requests +// through http.DefaultClient, so wrapping the default transport is what +// captures their outbound calls. +// +// commons has the same hook built in (logger.onPropertyUpdate wraps +// http.DefaultTransport for log.level.http), but it never fires here: the +// listener is registered by logger.UseSlog(), which captain does not call, and +// -P values are bound straight into properties.commandlineProperties without +// going through Set/Update, so no listener is notified. Do not delete this in +// favour of the commons hook — and if captain ever adopts UseSlog(), the Once +// below is what keeps a second printer from stacking onto the first. +// +// The same transport carries HAR capture when -Phttp.har= is set; the +// archive is written by FlushHAR once the command finishes. // -// Enable it with -Plog.level.http=trace3 (headers + timing) or trace4 (+bodies). // Must run before the first provider request; cmd/captain wires it from // PersistentPreRun, after clicky applies the logging flags. -func EnableHTTPWireLogging() { - if wrapped := wrapHTTPLogging(http.DefaultTransport); wrapped != http.DefaultTransport { - http.DefaultTransport = wrapped - http.DefaultClient = &http.Client{Transport: wrapped} +func EnableHTTPWireLogging() error { + var err error + httpLoggingOnce.Do(func() { + var captured http.RoundTripper + // HAR sits innermost, closest to the transport, matching the ordering + // commons/http uses so the archive records the request as it went out. + if captured, err = harRegistry.Transport(harFeature, http.DefaultTransport); err != nil { + return + } + if wrapped := wrapHTTPLogging(captured); wrapped != http.DefaultTransport { + http.DefaultTransport = wrapped + http.DefaultClient = &http.Client{Transport: wrapped} + } + }) + return err +} + +// harFeature names captain's traffic for per-subsystem overrides, so +// -Phttp.captain.har= and the global -Phttp.har= both resolve. +const harFeature = "captain" + +// harRegistry resolves -Phttp.har* into collectors. It is process-wide because +// http.DefaultTransport is. Its own logger means capture announcements can be +// silenced or raised with -Plog.level.har. +var harRegistry = har.NewRegistry(logger.GetLogger("har")) + +// FlushHAR writes any archive requested with -Phttp.har=. cmd/captain +// calls it after Execute returns — including on the error path, which is the +// run most worth capturing. +func FlushHAR() { + if err := harRegistry.Flush(); err != nil { + log.Errorf("%v", err) } } -// wrapHTTPLogging returns base wrapped with commons' HTTP logger when the "http" -// logger is at trace3+, otherwise base unchanged. Split out from -// EnableHTTPWireLogging so the level gating is unit-testable without mutating +// wrapHTTPLogging returns base wrapped with the commons trace middleware, or +// base unchanged when the resolved level is too low to log anything. Split out +// from EnableHTTPWireLogging so the ladder is unit-testable without mutating // http.DefaultTransport. func wrapHTTPLogging(base http.RoundTripper) http.RoundTripper { - h := logger.GetLogger("http") - if !h.IsLevelEnabled(logger.Trace3) { + cfg, ok := httpTraceConfig() + if !ok { return base } - return logger.NewHttpLoggerWithLevels(h, base, logger.Trace3, logger.Trace4) + return middlewares.NewLogger(cfg)(base) +} + +// httpTraceConfig maps the effective HTTP log level onto commons' trace ladder, +// which is relative to a base level (Debug by default, overridable with +// -Phttp.log.base-level or HTTP_LOG_BASE_LEVEL): +// +// warn and below nothing is installed +// info failed requests only (status >= 400 or transport error) +// -v an access line per request +// -vv + request/response headers, query and form params +// -vvv + request bodies, TLS summary +// -vvvv + response bodies +func httpTraceConfig() (commonshttp.TraceConfig, bool) { + cfg, ok := commonshttp.TraceConfigForLogLevel(httpTraceLevel()) + if !ok { + return cfg, false + } + // commons redacts these already via CommonRedactedHeaders ("*-Key"); naming + // captain's own credential headers keeps that guarantee explicit, since the + // Anthropic and Gemini fetchers authenticate with them rather than Bearer. + cfg.RedactedHeaders = append(cfg.RedactedHeaders, "x-api-key", "x-goog-api-key") + return cfg, true +} + +// httpTraceLevel resolves the level the ladder is built from. Named loggers +// take their level from --log-level/-Plog.level.http, while the -v count is +// applied to the global logger, so the effective level is the more verbose of +// the two. +func httpTraceLevel() logger.LogLevel { + level := logger.GetLogger("http").GetLevel() + if root := logger.StandardLogger().GetLevel(); root > level { + level = root + } + return level } diff --git a/pkg/cli/logging_test.go b/pkg/cli/logging_test.go index fdd49435..b1305821 100644 --- a/pkg/cli/logging_test.go +++ b/pkg/cli/logging_test.go @@ -1,35 +1,185 @@ package cli import ( + "bytes" + "context" + "io" "net/http" + "net/http/httptest" + "strings" "testing" + commonsctx "github.com/flanksource/commons/context" "github.com/flanksource/commons/logger" ) -// TestWrapHTTPLogging verifies the level gate: the default transport is only -// wrapped once the "http" logger reaches trace3, which is how -// -Plog.level.http=trace3 enables redacted HTTP wire logging. -func TestWrapHTTPLogging(t *testing.T) { - h := logger.GetLogger("http") - orig := h.GetLevel() - t.Cleanup(func() { h.SetLogLevel(orig) }) +// withLogLevel pins both the "http" logger and the global logger for the +// duration of a test, because httpTraceLevel resolves the more verbose of the +// two. +func withLogLevel(t *testing.T, level logger.LogLevel) { + t.Helper() + h, root := logger.GetLogger("http"), logger.StandardLogger() + origHTTP, origRoot := h.GetLevel(), root.GetLevel() + t.Cleanup(func() { + h.SetLogLevel(origHTTP) + root.SetLogLevel(origRoot) + }) + h.SetLogLevel(level) + root.SetLogLevel(level) +} + +// TestHTTPTraceLadder pins the rungs captain advertises: nothing below warn, an +// errors-only access log at the default info level, then access lines, headers, +// request bodies, and response bodies as verbosity climbs. +func TestHTTPTraceLadder(t *testing.T) { + for _, tc := range []struct { + name string + level logger.LogLevel + wantInstalled bool + wantErrorsOnly bool + wantHeaders bool + wantRequestBody bool + wantResponse bool + }{ + {name: "warn installs nothing", level: logger.Warn}, + {name: "info logs failures only", level: logger.Info, wantInstalled: true, wantErrorsOnly: true}, + {name: "-v logs every request", level: logger.Debug, wantInstalled: true}, + {name: "-vv adds headers", level: logger.Trace, wantInstalled: true, wantHeaders: true}, + {name: "-vvv adds request bodies", level: logger.Trace1, wantInstalled: true, wantHeaders: true, wantRequestBody: true}, + {name: "-vvvv adds response bodies", level: logger.Trace2, wantInstalled: true, wantHeaders: true, wantRequestBody: true, wantResponse: true}, + } { + t.Run(tc.name, func(t *testing.T) { + withLogLevel(t, tc.level) + + cfg, ok := httpTraceConfig() + if ok != tc.wantInstalled { + t.Fatalf("installed = %v, want %v (level %v)", ok, tc.wantInstalled, tc.level) + } + base := http.DefaultTransport + if wrapped := wrapHTTPLogging(base); (wrapped != base) != tc.wantInstalled { + t.Fatalf("transport wrapped = %v, want %v", wrapped != base, tc.wantInstalled) + } + if !ok { + return + } + if !cfg.AccessLog { + t.Error("every installed rung must keep the access log so failures surface") + } + if cfg.AccessLogErrorsOnly != tc.wantErrorsOnly { + t.Errorf("AccessLogErrorsOnly = %v, want %v", cfg.AccessLogErrorsOnly, tc.wantErrorsOnly) + } + if cfg.Headers != tc.wantHeaders || cfg.ResponseHeaders != tc.wantHeaders { + t.Errorf("Headers/ResponseHeaders = %v/%v, want %v", cfg.Headers, cfg.ResponseHeaders, tc.wantHeaders) + } + if cfg.Body != tc.wantRequestBody { + t.Errorf("Body = %v, want %v", cfg.Body, tc.wantRequestBody) + } + if cfg.Response != tc.wantResponse { + t.Errorf("Response = %v, want %v", cfg.Response, tc.wantResponse) + } + }) + } +} + +// TestWrapHTTPLoggingRedactsProviderCredentials covers the headers captain +// actually authenticates with: Anthropic sends x-api-key and Gemini +// x-goog-api-key, neither of which httpretty's default sanitizers cover. +func TestWrapHTTPLoggingRedactsProviderCredentials(t *testing.T) { + const ( + anthropicKey = "sk-ant-api03-REDACTMEANTHROPIC" + geminiKey = "AIzaSyD-REDACTMEGEMINI" + bearerKey = "sk-proj-REDACTMEOPENAI" + ) - base := http.DefaultTransport + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[]}`)) + })) + t.Cleanup(server.Close) - h.SetLogLevel(logger.Info) - if got := wrapHTTPLogging(base); got != base { - t.Fatalf("at info level the transport must be returned unwrapped, got %T", got) + withLogLevel(t, logger.Trace) + + var out bytes.Buffer + ctx := commonsctx.NewContext(context.Background(), commonsctx.WithLogger(logger.NewWithWriter(&out))) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + if err != nil { + t.Fatalf("NewRequest: %v", err) } + req.Header.Set("x-api-key", anthropicKey) + req.Header.Set("x-goog-api-key", geminiKey) + req.Header.Set("Authorization", "Bearer "+bearerKey) - h.SetLogLevel(logger.Debug) - if got := wrapHTTPLogging(base); got != base { - t.Fatalf("at debug level (below trace3) the transport must stay unwrapped, got %T", got) + resp, err := wrapHTTPLogging(http.DefaultTransport).RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip: %v", err) } + _ = resp.Body.Close() + + logged := out.String() + if !strings.Contains(logged, "X-Api-Key") || !strings.Contains(logged, "X-Goog-Api-Key") { + t.Fatalf("expected both credential headers to be named in the trace, got:\n%s", logged) + } + for _, secret := range []string{anthropicKey, geminiKey, bearerKey} { + if strings.Contains(logged, secret) { + t.Errorf("credential %q leaked into the HTTP trace:\n%s", secret, logged) + } + } +} + +// TestErrorsOnlyAccessLogAtDefaultVerbosity covers the rung captain runs at +// with no flags: a failing provider call is reported (with its body, which +// carries the provider's error message) while successful calls stay silent. +func TestErrorsOnlyAccessLogAtDefaultVerbosity(t *testing.T) { + const providerError = `{"error":{"message":"invalid x-api-key"}}` + + for _, tc := range []struct { + name string + status int + wantLogged bool + }{ + {name: "success stays silent", status: http.StatusOK}, + {name: "failure is reported", status: http.StatusUnauthorized, wantLogged: true}, + } { + t.Run(tc.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(tc.status) + _, _ = w.Write([]byte(providerError)) + })) + t.Cleanup(server.Close) + + withLogLevel(t, logger.Info) + + var out bytes.Buffer + ctx := commonsctx.NewContext(context.Background(), commonsctx.WithLogger(logger.NewWithWriter(&out))) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) + if err != nil { + t.Fatalf("NewRequest: %v", err) + } + + resp, err := wrapHTTPLogging(http.DefaultTransport).RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip: %v", err) + } + body, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + + // The access log reads the error body to report it; downstream callers + // must still see it intact. + if string(body) != providerError { + t.Errorf("response body = %q, want it restored to %q", body, providerError) + } - h.SetLogLevel(logger.Trace3) - if got := wrapHTTPLogging(base); got == base { - t.Fatal("at trace3 the transport must be wrapped with the commons HTTP logger") + logged := out.String() + if !tc.wantLogged { + if strings.Contains(logged, server.URL) { + t.Errorf("a successful request must not be logged at default verbosity, got:\n%s", logged) + } + return + } + if !strings.Contains(logged, "401") || !strings.Contains(logged, "invalid x-api-key") { + t.Errorf("expected the failure and its body to be logged, got:\n%s", logged) + } + }) } } diff --git a/pkg/cli/plan.go b/pkg/cli/plan.go index a84a010c..64f7c1dd 100644 --- a/pkg/cli/plan.go +++ b/pkg/cli/plan.go @@ -2,7 +2,6 @@ package cli import ( "context" - "errors" "fmt" "os" "strings" @@ -57,29 +56,10 @@ func RunPlan(opts PlanOptions) (PlanResult, error) { id := strings.TrimSpace(opts.SessionID) if id != "" { - persisted, ok, err := resolveNativePlan(ctx, db, id, source) + plan, err := resolveIdentityPlan(ctx, db, id, source) if err != nil { return PlanResult{}, err } - if ok { - persisted.pathOnly = opts.PathOnly - return *persisted, nil - } - overview, err := db.GetSessionOverviewByIdentity(ctx, id) - if err != nil { - return PlanResult{}, err - } - 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 { - return PlanResult{}, err - } - if plan == nil { - return PlanResult{}, fmt.Errorf("session %q has no plan", id) - } plan.pathOnly = opts.PathOnly return *plan, nil } @@ -145,21 +125,70 @@ func resolveLatestTranscriptPlan( } } -// 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, "", "") +// planIdentityStore resolves an identity to every matching session overview and +// reads the plans recorded against a Captain session UUID. +type planIdentityStore interface { + sessionOverviewStore + ListPlans(context.Context, captaindb.PlanFilter) ([]captaindb.Plan, error) +} + +// resolveIdentityPlan resolves a Captain UUID or provider session ID to a plan. +// The identity lookup is plural because the same provider session ID may exist +// once per source (captain_sessions is unique on source+host+provider id): a +// gavel orchestration row carries no transcript while the provider row it +// parents carries the real one. Persisted plans win over transcript recovery, +// and transcript recovery uses the first match that actually has a transcript. +func resolveIdentityPlan(ctx context.Context, db planIdentityStore, identity, source string) (*PlanResult, error) { + overviews, err := resolveOverviewsByIdentity(ctx, db, identity) if err != nil { - if errors.Is(err, captaindb.ErrSessionNotFound) { - return nil, false, nil + return nil, err + } + if source != "all" { + filtered := make([]captaindb.SessionOverview, 0, len(overviews)) + for i := range overviews { + if overviews[i].Source == source { + filtered = append(filtered, overviews[i]) + } } - return nil, false, fmt.Errorf("resolve persisted Captain session %q: %w", identity, err) + overviews = filtered } + if len(overviews) == 0 { + return nil, fmt.Errorf("%w: %s", captaindb.ErrSessionNotFound, identity) + } + for i := range overviews { + persisted, ok, err := resolveNativePlan(ctx, db, overviews[i]) + if err != nil { + return nil, err + } + if ok { + return persisted, nil + } + } + for i := range overviews { + candidate := candidateFromOverview(overviews[i]) + if candidate.path == "" { + continue + } + plan, err := resolveSessionPlan(candidate) + if err != nil { + return nil, err + } + if plan != nil { + return plan, nil + } + return nil, fmt.Errorf("session %q has no plan", identity) + } + return nil, fmt.Errorf("session %q has no transcript recorded on this host", identity) +} + +// resolveNativePlan resolves persisted plan content for one Captain session +// 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 planIdentityStore, + session captaindb.SessionOverview, +) (*PlanResult, bool, error) { 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) diff --git a/pkg/cli/plan_identity_ginkgo_test.go b/pkg/cli/plan_identity_ginkgo_test.go new file mode 100644 index 00000000..4d1c356a --- /dev/null +++ b/pkg/cli/plan_identity_ginkgo_test.go @@ -0,0 +1,135 @@ +package cli + +import ( + "context" + "encoding/json" + "errors" + "os" + "path/filepath" + + "github.com/flanksource/captain/pkg/database" + "github.com/google/uuid" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// planIdentityStoreStub resolves one identity to several overviews, mirroring a +// provider session ID recorded once per source (captain_sessions is unique on +// source+host_id+provider_session_id, not on provider_session_id alone). +type planIdentityStoreStub struct { + overviews []database.SessionOverview + plans map[uuid.UUID][]database.Plan +} + +func (s *planIdentityStoreStub) ListSessionOverviewsByIdentity( + context.Context, string, +) ([]database.SessionOverview, error) { + return s.overviews, nil +} + +func (s *planIdentityStoreStub) ListThreadSessionOverviews( + context.Context, uuid.UUID, +) ([]database.SessionOverview, error) { + return nil, nil +} + +func (s *planIdentityStoreStub) ListPlans( + _ context.Context, filter database.PlanFilter, +) ([]database.Plan, error) { + if filter.SourceSessionID == nil { + return nil, nil + } + return s.plans[*filter.SourceSessionID], nil +} + +var _ = Describe("plan identity resolution", func() { + const providerSessionID = "7657484f-e2e6-4f71-85c7-c244577a4028" + + // writeClaudePlanTranscript emits a one-entry Claude transcript whose + // ExitPlanMode call carries the plan, and returns its path. + writeClaudePlanTranscript := func(planMarkdown string) (string, string) { + home := GinkgoT().TempDir() + planPath := filepath.Join(home, ".claude", "plans", "identity-plan.md") + historyPath := filepath.Join(home, ".claude", "projects", "identity.jsonl") + entry, err := json.Marshal(map[string]any{ + "type": "assistant", "sessionId": providerSessionID, "uuid": "assistant-1", + "timestamp": "2026-08-05T10:00:00Z", "cwd": home, "slug": "identity-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": planMarkdown}}, + }}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(os.MkdirAll(filepath.Dir(historyPath), 0o755)).To(Succeed()) + Expect(os.WriteFile(historyPath, append(entry, '\n'), 0o644)).To(Succeed()) + return historyPath, planPath + } + + It("selects the transcript-bearing source when one provider ID spans sources", func(ctx SpecContext) { + const planMarkdown = "# identity plan" + historyPath, planPath := writeClaudePlanTranscript(planMarkdown) + orchestrationID, transcriptID := uuid.New(), uuid.New() + providerID := providerSessionID + store := &planIdentityStoreStub{overviews: []database.SessionOverview{ + // The gavel orchestration row is listed first and has no transcript. + {ID: orchestrationID, ProviderSessionID: &providerID, Source: "gavel"}, + {ID: transcriptID, ProviderSessionID: &providerID, Source: "claude", Path: &historyPath}, + }} + + plan, err := resolveIdentityPlan(ctx, store, providerSessionID, "all") + + Expect(err).NotTo(HaveOccurred()) + Expect(plan.Content).To(Equal(planMarkdown)) + Expect(plan.Path).To(Equal(planPath)) + Expect(plan.Source).To(Equal("claude")) + }) + + It("prefers a persisted plan over transcript recovery", func(ctx SpecContext) { + historyPath, _ := writeClaudePlanTranscript("# transcript plan") + orchestrationID, transcriptID := uuid.New(), uuid.New() + providerID := providerSessionID + planID, revisionID := uuid.New(), uuid.New() + store := &planIdentityStoreStub{ + overviews: []database.SessionOverview{ + {ID: orchestrationID, ProviderSessionID: &providerID, Source: "gavel"}, + {ID: transcriptID, ProviderSessionID: &providerID, Source: "claude", Path: &historyPath}, + }, + plans: map[uuid.UUID][]database.Plan{orchestrationID: {{ + ID: planID, Slug: "persisted", + ApprovedRevision: &database.PlanRevision{ID: revisionID, Revision: 2, PlanMarkdown: "# persisted plan"}, + }}}, + } + + plan, err := resolveIdentityPlan(ctx, store, providerSessionID, "all") + + Expect(err).NotTo(HaveOccurred()) + Expect(plan.Content).To(Equal("# persisted plan")) + Expect(plan.SessionID).To(Equal(orchestrationID.String())) + }) + + It("reports a missing transcript when no matched source recorded one", func(ctx SpecContext) { + providerID := providerSessionID + store := &planIdentityStoreStub{overviews: []database.SessionOverview{ + {ID: uuid.New(), ProviderSessionID: &providerID, Source: "gavel"}, + {ID: uuid.New(), ProviderSessionID: &providerID, Source: "claude"}, + }} + + _, err := resolveIdentityPlan(ctx, store, providerSessionID, "all") + + Expect(err).To(MatchError(ContainSubstring("has no transcript recorded on this host"))) + Expect(errors.Is(err, database.ErrSessionConflict)).To(BeFalse()) + }) + + It("narrows matches to the requested source", func(ctx SpecContext) { + historyPath, _ := writeClaudePlanTranscript("# identity plan") + providerID := providerSessionID + store := &planIdentityStoreStub{overviews: []database.SessionOverview{ + {ID: uuid.New(), ProviderSessionID: &providerID, Source: "gavel"}, + {ID: uuid.New(), ProviderSessionID: &providerID, Source: "claude", Path: &historyPath}, + }} + + _, err := resolveIdentityPlan(ctx, store, providerSessionID, "codex") + + Expect(errors.Is(err, database.ErrSessionNotFound)).To(BeTrue()) + }) +}) diff --git a/pkg/cli/plan_native_integration_test.go b/pkg/cli/plan_native_integration_test.go index df164a45..1f590687 100644 --- a/pkg/cli/plan_native_integration_test.go +++ b/pkg/cli/plan_native_integration_test.go @@ -41,17 +41,15 @@ func TestResolveNativePlanUsesPersistedApprovedContentWithoutSourceFile(t *testi require.NoError(t, err) require.NoError(t, os.Remove(deletedPath)) - result, ok, err := resolveNativePlan(t.Context(), db, "provider-plan-session", "all") + result, err := resolveIdentityPlan(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") + byUUID, err := resolveIdentityPlan(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/prompt_batch_session.go b/pkg/cli/prompt_batch_session.go index f01bfc91..e9407dad 100644 --- a/pkg/cli/prompt_batch_session.go +++ b/pkg/cli/prompt_batch_session.go @@ -30,7 +30,7 @@ func createPromptBatchSessions(ctx context.Context, rendered PromptRenderResult, if err := validatePromptRuntimes(runtimes); err != nil { return promptBatchSession{}, err } - db, err := captainDB(ctx) + db, err := captainDefaultDB(ctx) if err != nil { return promptBatchSession{}, err } @@ -100,7 +100,7 @@ func promptBinding(batch promptBatchSession, index int) *promptSessionBinding { } func updatePromptSessionLifecycle(ctx context.Context, id uuid.UUID, lifecycle database.SessionLifecycleStatus, reason string) { - db, err := captainDB(ctx) + db, err := captainDefaultDB(ctx) if err != nil { log.Errorf("open database for session %s lifecycle: %v", id, err) return diff --git a/pkg/cli/prompt_entity.go b/pkg/cli/prompt_entity.go index 135629cb..154d32f4 100644 --- a/pkg/cli/prompt_entity.go +++ b/pkg/cli/prompt_entity.go @@ -78,11 +78,12 @@ func (p PromptSummary) Row() map[string]any { type PromptDetail struct { PromptSummary - Content string `json:"content"` - InputSchema map[string]any `json:"inputSchema,omitempty"` - InputDefault map[string]any `json:"inputDefault,omitempty"` - OutputSchema map[string]any `json:"outputSchema,omitempty"` - Metadata map[string]any `json:"metadata,omitempty"` + Content string `json:"content"` + InputSchema map[string]any `json:"inputSchema,omitempty"` + InputDefault map[string]any `json:"inputDefault,omitempty"` + OutputSchema map[string]any `json:"outputSchema,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` + Run PromptRenderRequest `json:"run"` } type PromptWriteRequest struct { @@ -287,6 +288,9 @@ func updatePrompt(ctx context.Context, id string, body map[string]any) (PromptDe if err != nil { return PromptDetail{}, err } + if !record.Source.Writable { + return PromptDetail{}, fmt.Errorf("prompt source %q is read-only; use create to save a copy", record.Source.Label) + } var req PromptWriteRequest if err := decodePromptBody(ctx, body, &req); err != nil { return PromptDetail{}, err @@ -294,12 +298,6 @@ func updatePrompt(ctx context.Context, id string, body map[string]any) (PromptDe if strings.TrimSpace(req.Content) == "" { return PromptDetail{}, fmt.Errorf("prompt content cannot be empty") } - if !record.Source.Writable { - if strings.TrimSpace(req.RelPath) == "" { - req.RelPath = localForkRelPath(record) - } - return writeNewLocalPrompt(ctx, req) - } full, err := safeLocalPromptPath(record.Source, record.Rel) if err != nil { return PromptDetail{}, err @@ -310,17 +308,6 @@ func updatePrompt(ctx context.Context, id string, body map[string]any) (PromptDe return promptDetail(record) } -// localForkRelPath derives the destination path for a read-only (embedded) -// prompt saved into a writable source, stripping the source walk root so an -// embedded "testdata/commit.prompt" lands as "commit.prompt". -func localForkRelPath(record promptRecord) string { - rel := record.Rel - if root := record.Source.WalkRoot; root != "" { - rel = strings.TrimPrefix(rel, root+"/") - } - return rel -} - func deletePrompt(ctx context.Context, id string) error { record, err := resolvePromptRecord(ctx, id) if err != nil { diff --git a/pkg/cli/prompt_entity_test.go b/pkg/cli/prompt_entity_test.go index 601e09f4..5f620d1f 100644 --- a/pkg/cli/prompt_entity_test.go +++ b/pkg/cli/prompt_entity_test.go @@ -256,7 +256,7 @@ func assertSchemaHasProps(t *testing.T, label string, schema map[string]any, key } } -func TestUpdateEmbeddedPromptForksToLocal(t *testing.T) { +func TestUpdateEmbeddedPromptRequiresSaveAs(t *testing.T) { isolateCaptainConfig(t) dir := t.TempDir() @@ -276,24 +276,39 @@ func TestUpdateEmbeddedPromptForksToLocal(t *testing.T) { } newContent := original.Content + "\n{{! local override }}\n" - forked, err := updatePrompt(ctx, embedded.ID, map[string]any{"content": newContent}) + if _, err := updatePrompt(ctx, embedded.ID, map[string]any{"content": newContent}); err == nil { + t.Fatal("updatePrompt(embedded) succeeded, want read-only error") + } else if !strings.Contains(err.Error(), "read-only") || !strings.Contains(err.Error(), "create") { + t.Fatalf("updatePrompt(embedded) error = %q, want read-only create guidance", err) + } + if entries, err := os.ReadDir(dir); err != nil { + t.Fatalf("read local prompt directory: %v", err) + } else if len(entries) != 0 { + t.Fatalf("updatePrompt(embedded) created %d local files, want none", len(entries)) + } + + savedAs, err := createPrompt(ctx, map[string]any{ + "name": "Commit Copy", + "relPath": "copies/commit.prompt", + "content": newContent, + }) if err != nil { - t.Fatalf("updatePrompt(embedded) err = %v", err) + t.Fatalf("createPrompt(save as) err = %v", err) } - if !forked.Writable || forked.SourceKind != "local" { - t.Fatalf("forked prompt = kind %q writable %v, want local writable", forked.SourceKind, forked.Writable) + if !savedAs.Writable || savedAs.SourceKind != "local" { + t.Fatalf("saved-as prompt = kind %q writable %v, want local writable", savedAs.SourceKind, savedAs.Writable) } - if forked.ID == embedded.ID { - t.Fatalf("forked prompt kept embedded id %q", forked.ID) + if savedAs.ID == embedded.ID { + t.Fatalf("saved-as prompt kept embedded id %q", savedAs.ID) } - if forked.RelPath != "commit.prompt" { - t.Fatalf("forked relPath = %q, want commit.prompt (testdata/ stripped)", forked.RelPath) + if savedAs.RelPath != "copies/commit.prompt" { + t.Fatalf("saved-as relPath = %q, want copies/commit.prompt", savedAs.RelPath) } - if !strings.Contains(forked.Content, "local override") { - t.Fatalf("forked content did not persist edit: %q", forked.Content) + if !strings.Contains(savedAs.Content, "local override") { + t.Fatalf("saved-as content did not persist edit: %q", savedAs.Content) } - if _, err := os.Stat(filepath.Join(dir, "commit.prompt")); err != nil { - t.Fatalf("forked prompt file missing: %v", err) + if _, err := os.Stat(filepath.Join(dir, "copies", "commit.prompt")); err != nil { + t.Fatalf("saved-as prompt file missing: %v", err) } stillEmbedded, err := getPrompt(ctx, embedded.ID) @@ -301,11 +316,7 @@ func TestUpdateEmbeddedPromptForksToLocal(t *testing.T) { t.Fatalf("getPrompt(embedded) after fork err = %v", err) } if strings.Contains(stillEmbedded.Content, "local override") { - t.Fatalf("embedded prompt was mutated by fork") - } - - if _, err := updatePrompt(ctx, embedded.ID, map[string]any{"content": newContent}); err == nil { - t.Fatalf("second fork of same prompt should fail with already-exists") + t.Fatal("embedded prompt was mutated by save as") } } diff --git a/pkg/cli/prompt_records.go b/pkg/cli/prompt_records.go index 664a530b..620b2887 100644 --- a/pkg/cli/prompt_records.go +++ b/pkg/cli/prompt_records.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io/fs" + "maps" "os" "path/filepath" "sort" @@ -13,6 +14,7 @@ import ( "time" promptlib "github.com/flanksource/captain/pkg/ai/prompt" + "github.com/flanksource/captain/pkg/api" dp "github.com/google/dotprompt/go/dotprompt" ) @@ -228,6 +230,10 @@ func promptDetailFromContent(record promptRecord, content string) (PromptDetail, if err != nil { return PromptDetail{}, err } + spec := &api.Spec{Model: api.Model{ + Name: summary.Model, + Backend: api.Backend(summary.Backend), + }} return PromptDetail{ PromptSummary: summary, Content: content, @@ -235,9 +241,34 @@ func promptDetailFromContent(record promptRecord, content string) (PromptDetail, InputDefault: inspection.InputDefault, OutputSchema: inspection.OutputSchema, Metadata: inspection.Metadata, + Run: PromptRenderRequest{ + Variables: maps.Clone(inspection.InputDefault), + Spec: spec, + Runtimes: promptRunModels(summary.Runtimes), + Chat: len(inspection.OutputSchema) == 0, + }, }, nil } +func promptRunModels(models []api.Model) []api.Model { + if len(models) == 0 { + return nil + } + out := make([]api.Model, len(models)) + for index, model := range models { + out[index] = api.Model{ + Name: model.Name, + ID: model.ID, + Backend: model.Backend, + Temperature: model.Temperature, + Effort: model.Effort, + NoCache: model.NoCache, + Fallbacks: promptRunModels(model.Fallbacks), + } + } + return out +} + func promptSummaryFromContent(record promptRecord, content string) (PromptSummary, error) { tmpl := promptlib.Load(content) req, cfg, err := tmpl.Render(map[string]any{}, nil) diff --git a/pkg/cli/prompt_run_events.go b/pkg/cli/prompt_run_events.go index 8a1b6967..938faa5f 100644 --- a/pkg/cli/prompt_run_events.go +++ b/pkg/cli/prompt_run_events.go @@ -6,6 +6,7 @@ import ( "sync" "github.com/flanksource/captain/pkg/ai" + "github.com/flanksource/captain/pkg/bash" "github.com/flanksource/captain/pkg/session" "github.com/segmentio/encoding/json" ) @@ -72,11 +73,16 @@ func (a *promptEventAccumulator) handle(_ int, ev ai.Event) { case ai.EventSystem: if ev.SessionID != "" { a.sessionID = ev.SessionID - } - a.task.SetDescription("starting") - if ev.SessionID != "" { + a.task.SetDescription("starting") a.task.Infof("session %s", ev.SessionID) } + // A lifecycle hook narrating what it did between turns (see + // HookContext.Notify). It stands on its own rather than joining the + // in-flight assistant turn, so the buffers are flushed first. + if ev.Text != "" { + a.flush() + a.emitNotice(ev.Text) + } case ai.EventThinking: a.appendThinking(ev.Text) a.task.SetDescription("thinking") @@ -161,8 +167,19 @@ func (a *promptEventAccumulator) flush() { a.thinkBuf.Reset() } +func (a *promptEventAccumulator) resetFrame() { + a.mu.Lock() + defer a.mu.Unlock() + a.flush() + clear(a.toolByID) +} + func (a *promptEventAccumulator) emitToolUse(ev ai.Event) { a.tools++ + input := ev.Input + if ev.Tool == "Bash" { + input = bash.TransformBashInput(input) + } msg := &session.Message{ ID: a.toolID(ev.ToolCallID), Role: "assistant", @@ -171,7 +188,7 @@ func (a *promptEventAccumulator) emitToolUse(ev ai.Event) { ToolName: ev.Tool, ToolCallID: ev.ToolCallID, State: session.ToolStateInputAvailable, - Input: mapToRaw(ev.Input), + Input: mapToRaw(input), }}, Provenance: a.provenance(), } @@ -219,6 +236,20 @@ func (a *promptEventAccumulator) emitError(ev ai.Event) { }) } +// emitNotice renders one lifecycle line as a discrete system message. Discrete +// rather than appended to a buffer like assistant text: each notice is already +// whole when it arrives, and giving it its own id keeps a later turn's text from +// overwriting it in a viewer that dedupes by id. +func (a *promptEventAccumulator) emitNotice(text string) { + a.task.Infof("%s", text) + a.emit(session.Message{ + ID: a.nextID("notice"), + Role: "system", + Parts: []session.Part{{Type: session.PartText, Text: text}}, + Provenance: a.provenance(), + }) +} + func (a *promptEventAccumulator) nextID(kind string) string { a.seq++ if a.idPrefix == "" { diff --git a/pkg/cli/prompt_run_events_test.go b/pkg/cli/prompt_run_events_test.go index 6c21c7b8..e3d9a576 100644 --- a/pkg/cli/prompt_run_events_test.go +++ b/pkg/cli/prompt_run_events_test.go @@ -137,6 +137,34 @@ func TestPromptRunAccumulator_EmitsErrorFrame(t *testing.T) { } } +// A hook narrating what it did between turns arrives as an EventSystem carrying +// text rather than a session id. It has to break out of the assistant's +// in-flight turn, not be appended to it, or a commit line lands inside the +// model's prose. +func TestPromptRunAccumulator_EmitsHookNoticeAsItsOwnSystemFrame(t *testing.T) { + msgs := collectEntries("m", "b", + ai.Event{Kind: ai.EventText, Text: "done, committing now"}, + ai.Event{Kind: ai.EventSystem, Text: "[post-turn] committed abc1234: fix: the thing"}, + ai.Event{Kind: ai.EventText, Text: "next turn"}, + ) + if len(msgs) != 3 { + t.Fatalf("want assistant/system/assistant frames, got %d: %+v", len(msgs), msgs) + } + notice := msgs[1] + if notice.Role != "system" { + t.Errorf("notice role = %q, want system", notice.Role) + } + if got := notice.Parts[0].Text; got != "[post-turn] committed abc1234: fix: the thing" { + t.Errorf("notice text = %q", got) + } + if msgs[2].ID == msgs[0].ID { + t.Error("the turn after a notice reused the earlier text id, so a viewer that dedupes by id would overwrite it") + } + if msgs[2].Parts[0].Text != "next turn" { + t.Errorf("text after a notice = %q, want a fresh buffer", msgs[2].Parts[0].Text) + } +} + func TestPromptRunAccumulator_CapturesSessionAndUsage(t *testing.T) { acc := newPromptEventAccumulator(func(session.Message) {}, fakeTaskSink{}, "m", "b") acc.handle(0, ai.Event{Kind: ai.EventSystem, SessionID: "sess-1"}) diff --git a/pkg/cli/prompt_run_persist.go b/pkg/cli/prompt_run_persist.go index 5e25ea49..3f24c742 100644 --- a/pkg/cli/prompt_run_persist.go +++ b/pkg/cli/prompt_run_persist.go @@ -36,7 +36,7 @@ func persistPromptRun(ctx context.Context, input promptRunRecordInput) { if source == "" { source = "claude" } - db, err := captainDB(ctx) + db, err := captainDefaultDB(ctx) if err != nil { log.Errorf("persist prompt run for session %s: %v", input.SessionID, err) return diff --git a/pkg/cli/prompt_runtimes_ginkgo_test.go b/pkg/cli/prompt_runtimes_ginkgo_test.go index 7c4bab12..cace1f90 100644 --- a/pkg/cli/prompt_runtimes_ginkgo_test.go +++ b/pkg/cli/prompt_runtimes_ginkgo_test.go @@ -54,6 +54,35 @@ Review the screenshot. )) }) + It("serves the canonical prompt run request with the detail", func() { + record, err := filePromptRecord(path) + Expect(err).NotTo(HaveOccurred()) + + detail, err := promptDetail(record) + + Expect(err).NotTo(HaveOccurred()) + Expect(detail.Run).To(Equal(PromptRenderRequest{ + Variables: map[string]any{}, + Spec: &api.Spec{Model: api.Model{ + Name: "gemini-3.5-flash", + Backend: api.BackendGemini, + }}, + Runtimes: []api.Model{ + { + Name: "gemini-3.5-flash", + Backend: api.BackendGemini, + Effort: api.EffortHigh, + }, + { + Name: "claude-sonnet-5", + Backend: api.BackendAnthropic, + Effort: api.EffortMedium, + }, + }, + Chat: true, + })) + }) + DescribeTable("resolves a discovered prompt by bare filename", func(id string) { ctx := ContextWithPromptDirs(context.Background(), []string{filepath.Dir(path)}) @@ -166,3 +195,34 @@ Review the screenshot. Expect(err).To(MatchError(ContainSubstring("field typo not found"))) }) }) + +var _ = Describe("prompt schema model catalog", func() { + It("keeps one exact runtime row per backend", func() { + models := flatModels([]AdapterStatus{ + { + Backend: string(api.BackendCodexCLI), + Type: "cli", + Authenticated: true, + Binary: "/usr/local/bin/codex", + Models: []string{"gpt-5.6-sol"}, + }, + { + Backend: string(api.BackendCodexCmux), + Type: "cli", + Authenticated: true, + Binary: "/usr/local/bin/codex", + Models: []string{"gpt-5.6-sol"}, + }, + }) + + Expect(models).To(HaveLen(2)) + Expect(models[0]["runtime"]).To(Equal(api.Model{ + Name: "gpt-5.6-sol", + Backend: api.BackendCodexCLI, + })) + Expect(models[1]["runtime"]).To(Equal(api.Model{ + Name: "gpt-5.6-sol", + Backend: api.BackendCodexCmux, + })) + }) +}) diff --git a/pkg/cli/prompt_schema_build.go b/pkg/cli/prompt_schema_build.go index fa8b13c5..c4e0b40c 100644 --- a/pkg/cli/prompt_schema_build.go +++ b/pkg/cli/prompt_schema_build.go @@ -327,16 +327,11 @@ func injectSpecConditionals(specMap map[string]any, adapters []AdapterStatus, ar return nil } -// 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. +// flatModels serves one display row per exact Captain runtime. A model exposed +// by multiple backends intentionally appears once per backend so selecting it +// produces a complete api.Model without client-side inference. func flatModels(adapters []AdapterStatus) []map[string]any { - type entry struct { - data map[string]any - backends []string - } - out := []entry{} - positions := map[string]int{} + out := []map[string]any{} for _, a := range adapters { provider := api.CatalogPrefixFor(api.Backend(a.Backend)) for _, model := range flatModelDetails(a) { @@ -344,35 +339,20 @@ func flatModels(adapters []AdapterStatus) []map[string]any { 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(), + out = append(out, map[string]any{ + "id": id, + "label": label, + "provider": provider, + "reasoning": modelSupportsReasoning(id), + "configured": a.Ready(), + "backends": []string{a.Backend}, + "runtime": api.Model{ + Name: id, + Backend: api.Backend(a.Backend), }, }) if len(model.supportedEfforts) > 0 { @@ -380,18 +360,14 @@ func flatModels(adapters []AdapterStatus) []map[string]any { for _, effort := range model.supportedEfforts { values = append(values, string(effort)) } - out[len(out)-1].data["supportedEfforts"] = values + out[len(out)-1]["supportedEfforts"] = values } if model.defaultEffort != api.EffortNone { - out[len(out)-1].data["defaultEffort"] = string(model.defaultEffort) + out[len(out)-1]["defaultEffort"] = string(model.defaultEffort) } } } - flat := make([]map[string]any, 0, len(out)) - for _, item := range out { - flat = append(flat, item.data) - } - return flat + return out } type flatModelDetail struct { diff --git a/pkg/cli/prompt_schema_test.go b/pkg/cli/prompt_schema_test.go index 9016852b..5edc6bfa 100644 --- a/pkg/cli/prompt_schema_test.go +++ b/pkg/cli/prompt_schema_test.go @@ -94,7 +94,12 @@ func TestPromptSchemaDocumentBackendsAndConditionals(t *testing.T) { 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)) + if got := codexModel["runtime"]; !reflect.DeepEqual(got, api.Model{ + Name: "gpt-5.6-sol", + Backend: api.BackendCodexCLI, + }) { + t.Errorf("flat model runtime = %#v, want exact codex-cli runtime", got) + } anthropic := byName[string(api.BackendAnthropic)] if _, hasModels := anthropic["models"]; hasModels { @@ -327,19 +332,6 @@ func schemaModelForBackend(t *testing.T, models []map[string]any, id, backend st 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() diff --git a/pkg/cli/result_costs.go b/pkg/cli/result_costs.go new file mode 100644 index 00000000..3a00c4c6 --- /dev/null +++ b/pkg/cli/result_costs.go @@ -0,0 +1,116 @@ +package cli + +import ( + "context" + + "github.com/flanksource/captain/pkg/api" + "github.com/flanksource/captain/pkg/claude" + "github.com/flanksource/captain/pkg/database" +) + +// resultCostLookup resolves the result-derived cost captain recorded for the +// sessions it ran itself. +// +// A stored claude transcript carries no result record — no invocation summary, +// no running total — so replaying one can only rebuild the numbers from each +// response's usage and price them from the registry. That reconstruction is an +// estimate: it cannot see pricing the provider applied but never published +// (1-hour cache writes, for one), and on a real session it came out ~9% under +// the billed figure. +// +// Captain does hold the provider's own answer for any session it executed: +// pkg/aichat writes each EventResult's Usage and CostUSD to captain_model_calls. +// This looks that up so the replay surfaces report the billed figure rather than +// their own recomputation. +type resultCostLookup struct { + db *database.DB +} + +// newResultCostLookup opens the captain database, returning a lookup that +// reports nothing when no database is configured. Cost reporting over local +// transcripts must keep working without one, so an unavailable database +// degrades to the reconstruction rather than failing the command. +func newResultCostLookup(ctx context.Context) *resultCostLookup { + db, err := captainDB(ctx) + if err != nil || db == nil { + return &resultCostLookup{} + } + return &resultCostLookup{db: db} +} + +// resultCost is a session's recorded usage plus both readings of its cost: the +// resolved total, and the portion of it the provider actually reported. +type resultCost struct { + Usage api.Usage + Model string + TotalUSD float64 + ProviderUSD float64 +} + +// find returns the result captain recorded for a session identity (a provider +// session id or captain id). +// +// It reports a hit only when the stored rows carry a provider-reported cost. +// A row without one is not a result — it is another reconstruction, written by +// transcript ingest from the same per-message usage the caller already has, and +// possibly staler. Preferring it would trade a fresh estimate for an old one. +func (l *resultCostLookup) find(ctx context.Context, identity string) (resultCost, bool) { + if l == nil || l.db == nil || identity == "" { + return resultCost{}, false + } + overview, err := l.db.GetSessionOverviewByIdentity(ctx, identity) + // An ambiguous identity (SessionConflictError) or a plain miss both mean + // there is no single authoritative row to prefer. + if err != nil || overview == nil { + return resultCost{}, false + } + rows, err := l.db.ListThreadCosts(ctx, overview.ID) + if err != nil || len(rows) == 0 { + return resultCost{}, false + } + out := resultCost{Model: rows[0].Model} + for i := range rows { + out.Usage.InputTokens += int(rows[i].InputTokens) + out.Usage.OutputTokens += int(rows[i].OutputTokens) + out.Usage.ReasoningTokens += int(rows[i].ReasoningTokens) + out.Usage.CacheReadTokens += int(rows[i].CacheReadTokens) + out.Usage.CacheWriteTokens += int(rows[i].CacheWriteTokens) + // total_cost already resolves provider-reported against list-price per + // underlying call (see 67_view_session_costs.sql); provider_cost_usd is + // how much of it the provider itself reported. + out.TotalUSD += rows[i].TotalCost + out.ProviderUSD += rows[i].ProviderCostUSD + } + return out, out.ProviderUSD > 0 +} + +// applyResultCosts replaces each session's reconstructed usage and cost with the +// figures captain recorded from the provider's results, where it has them. +// Sessions captain never ran keep their reconstruction, which stays marked as an +// estimate because no provider cost is set on it. +func applyResultCosts(ctx context.Context, sessions []claude.SessionCost) { + if len(sessions) == 0 { + return + } + lookup := newResultCostLookup(ctx) + if lookup.db == nil { + return + } + for i := range sessions { + cost, ok := lookup.find(ctx, sessions[i].SessionID) + if !ok { + continue + } + sessions[i].Tokens = claude.TokenSummary{ + InputTokens: cost.Usage.InputTokens, + OutputTokens: cost.Usage.OutputTokens, + CacheWriteTokens: cost.Usage.CacheWriteTokens, + CacheReadTokens: cost.Usage.CacheReadTokens, + TotalCost: cost.TotalUSD, + ProviderCostUSD: cost.ProviderUSD, + } + if sessions[i].Model == "" { + sessions[i].Model = cost.Model + } + } +} diff --git a/pkg/cli/serve.go b/pkg/cli/serve.go index 3664dd23..249aa4c0 100644 --- a/pkg/cli/serve.go +++ b/pkg/cli/serve.go @@ -38,22 +38,20 @@ import ( var captainWebappFS embed.FS type ServeOptions struct { - Host string - Port int - Dev bool - UIPort int - Open bool - ThreadsFile string - PromptDirs []string - MCPServers []aichat.MCPServer + Host string + Port int + Dev bool + UIPort int + Open bool + PromptDirs []string + MCPServers []aichat.MCPServer } func NewServeCommand(version string) *cobra.Command { opts := ServeOptions{ - Host: "localhost", - Port: 9020, - UIPort: 0, - ThreadsFile: ".captain/chat-threads.json", + Host: "localhost", + Port: 9020, + UIPort: 0, } cmd := &cobra.Command{ @@ -69,21 +67,18 @@ 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 { - runOpts := opts - runOpts.Port = effectiveServePort(runOpts.Dev, cmd.Flags().Changed("port"), runOpts.Port) - if err := runOpts.validate(); err != nil { + if err := opts.validate(); err != nil { return err } - return RunServe(cmd.Context(), cmd.Root(), runOpts, version, cmd.OutOrStdout(), cmd.ErrOrStderr()) + return RunServe(cmd.Context(), cmd.Root(), opts, 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 (random when --dev is set)") + cmd.Flags().IntVarP(&opts.Port, "port", "p", opts.Port, "Port to bind the API server to") 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 (random by default)") cmd.Flags().BoolVar(&opts.Open, "open", false, "Open the web UI in the default browser") - cmd.Flags().StringVar(&opts.ThreadsFile, "threads-file", opts.ThreadsFile, "Path to persisted chat thread JSON") cmd.Flags().StringArrayVar(&opts.PromptDirs, "prompt-dir", nil, "Additional local directory containing .prompt files (repeatable)") return cmd @@ -104,7 +99,18 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve if err != nil { return err } - threadStore := newFileThreadStore(opts.ThreadsFile) + db, err := captainServeDB(ctx) + if err != nil { + return err + } + threadStore, err := aichat.NewDatabaseThreadStore(db) + if err != nil { + return err + } + authority, err := aichat.NewDatabaseExecutionAuthority(db) + if err != nil { + return err + } attachmentStore, err := newAttachmentStore(cwd) if err != nil { return err @@ -141,7 +147,7 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve addCaptainProviderTokenPaths(openAPISpec) addCaptainProviderDefaultsPaths(openAPISpec) addCaptainDisabledPaths(openAPISpec) - chat, mcpTools, err := newCaptainChatService(ctx, rootCmd, opts, cwd, threadStore, attachmentStore) + chat, mcpTools, err := newCaptainChatService(ctx, rootCmd, opts, cwd, authority, attachmentStore) if err != nil { return err } @@ -164,6 +170,7 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve 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/contexts", handleContexts()) mux.HandleFunc("GET /api/captain/projects", handleProjects()) mux.HandleFunc("GET /api/captain/sessions/live", handleSessionsLive()) mux.HandleFunc("GET /api/captain/sessions/throughput", handleSessionsThroughput()) @@ -205,7 +212,8 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve } httpSrv := &http.Server{ Addr: addr, - Handler: rpchttp.TimingMiddleware(PromptDirsMiddleware(mux, opts.PromptDirs)), + Handler: rpchttp.TimingMiddleware( + DatabaseContextMiddleware(PromptDirsMiddleware(mux, opts.PromptDirs))), ReadTimeout: 30 * time.Second, // /api/chat streams SSE; a fixed write timeout truncates long turns. IdleTimeout: 60 * time.Second, @@ -214,10 +222,6 @@ 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 @@ -237,8 +241,15 @@ func RunServe(ctx context.Context, rootCmd *cobra.Command, opts ServeOptions, ve if err != nil { return err } - dsn, source := captainDatabaseIdentity() + dsn, source := contextDatabaseIdentity(defaultDatabaseContextName) log.Infof("Database Info: source=%q dsn=%q live_sessions=%d", source, database.MaskDSN(dsn), liveSessions) + if contexts, err := databaseContexts(); err != nil { + // Reads can select a context per request, so a malformed context + // configuration is a startup error rather than a surprise mid-session. + return err + } else if len(contexts) > 1 { + log.Infof("Read-only database contexts: %s", strings.Join(databaseContextNames(contexts[1:]), ", ")) + } go prunePromptRuns(ctx, promptRuns) defer promptChats.stopAll() diff --git a/pkg/cli/serve_chat.go b/pkg/cli/serve_chat.go index ca953951..8d6d1c51 100644 --- a/pkg/cli/serve_chat.go +++ b/pkg/cli/serve_chat.go @@ -23,7 +23,7 @@ func newCaptainChatService( rootCmd *cobra.Command, opts ServeOptions, cwd string, - threadStore aichat.ThreadStore, + authority aichat.ExecutionAuthority, attachmentStore *attachments.Store, ) (*aichat.Service, *aichat.MCPToolProvider, error) { chatTools, err := clickyaichat.NewCobraToolProvider(clickyaichat.CobraToolProviderOptions{ @@ -47,7 +47,11 @@ func newCaptainChatService( }, }, nil }), - Tools: chatTools, MCP: mcpTools, Threads: threadStore, + // Thread reads follow the request's database context; writes never reach + // a secondary because the context middleware rejects unsafe methods. + Tools: chatTools, MCP: mcpTools, + Threads: aichat.ThreadStoreProviderFunc(contextThreadStore), + Authority: authority, Attachments: chatAttachmentResolver{store: attachmentStore}, }) return chat, mcpTools, nil diff --git a/pkg/cli/serve_disabled_test.go b/pkg/cli/serve_disabled_test.go index 4d9f794d..57a78d51 100644 --- a/pkg/cli/serve_disabled_test.go +++ b/pkg/cli/serve_disabled_test.go @@ -7,6 +7,7 @@ import ( "path/filepath" "reflect" "strings" + "sync" "testing" "github.com/flanksource/captain/pkg/api" @@ -67,6 +68,45 @@ func TestDisabledPutPreservesUnrelatedConfiguration(t *testing.T) { } } +func TestDisabledPutKeepsPersistedAndRuntimeSelectionsConsistentUnderConcurrency(t *testing.T) { + setupDisabledTest(t) + bodies := []string{ + `{"modes":["cmux"],"providers":[],"backends":[],"models":[],"efforts":[]}`, + `{"modes":[],"providers":["deepseek"],"backends":[],"models":[],"efforts":[]}`, + } + + for range 50 { + start := make(chan struct{}) + responses := make(chan *httptest.ResponseRecorder, len(bodies)) + var wg sync.WaitGroup + for _, body := range bodies { + wg.Add(1) + go func() { + defer wg.Done() + <-start + responses <- serveDisabledRequest(t, body) + }() + } + close(start) + wg.Wait() + close(responses) + for response := range responses { + if response.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", response.Code, response.Body.String()) + } + } + + config, _, err := captainconfig.Load() + if err != nil { + t.Fatalf("Load(): %v", err) + } + if !reflect.DeepEqual(api.Disabled(), config.AI.Disabled.Set()) { + t.Fatalf("runtime disabled set does not match persisted config: runtime=%+v persisted=%+v", + api.Disabled(), config.AI.Disabled.Set()) + } + } +} + func TestDisabledPutRejectsInvalidSets(t *testing.T) { for name, body := range map[string]string{ "unknown mode": `{"modes":["telepathy"],"providers":[],"backends":[],"models":[],"efforts":[]}`, diff --git a/pkg/cli/serve_port.go b/pkg/cli/serve_port.go index f1888062..d1524e76 100644 --- a/pkg/cli/serve_port.go +++ b/pkg/cli/serve_port.go @@ -7,13 +7,6 @@ import ( "strings" ) -func effectiveServePort(dev, portFlagSet bool, configuredPort int) int { - if dev && !portFlagSet { - return 0 - } - return configuredPort -} - func viteDevServerArgs(configuredPort int, open bool) ([]string, error) { port := configuredPort if port == 0 { @@ -41,15 +34,12 @@ 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) { + if o.Port < 1 || o.Port > 65535 { return fmt.Errorf("invalid --port %d", o.Port) } if o.Dev && (o.UIPort < 0 || 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) } diff --git a/pkg/cli/serve_port_ginkgo_test.go b/pkg/cli/serve_port_ginkgo_test.go index 74cc0d02..8495b4a9 100644 --- a/pkg/cli/serve_port_ginkgo_test.go +++ b/pkg/cli/serve_port_ginkgo_test.go @@ -9,19 +9,19 @@ import ( ) 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: 0, ThreadsFile: "threads.json"} + It("keeps the configured Captain port in development", func() { + options := ServeOptions{Host: "localhost", Port: 9020, Dev: true, UIPort: 0} Expect(options.validate()).To(Succeed()) + flag := NewServeCommand("test").Flags().Lookup("port") + Expect(flag).NotTo(BeNil()) + Expect(flag.DefValue).To(Equal("9020")) + }) + + It("rejects an ephemeral API port in every mode", func() { + options := ServeOptions{Host: "localhost", Port: 0, Dev: true, UIPort: 0} + Expect(options.validate()).To(MatchError("invalid --port 0")) + options.Dev = false Expect(options.validate()).To(MatchError("invalid --port 0")) }) @@ -50,7 +50,7 @@ var _ = Describe("serve ports", func() { Expect(args).To(Equal([]string{"exec", "vite", "--port", "62183", "--strictPort", "--host", "localhost"})) }) - It("keeps the ephemeral API port reserved for the server", func() { + It("reserves an automatically selected port while choosing it", func() { listener, addr, port, err := listenCaptainServer("127.0.0.1", 0) Expect(err).NotTo(HaveOccurred()) DeferCleanup(listener.Close) diff --git a/pkg/cli/serve_sessions.go b/pkg/cli/serve_sessions.go index 9de1c4ab..083be25d 100644 --- a/pkg/cli/serve_sessions.go +++ b/pkg/cli/serve_sessions.go @@ -67,7 +67,7 @@ func handleSessionsLiveWithRunner(run func(context.Context, SessionLiveOptions) } result, err := run(r.Context(), opts) if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) + http.Error(w, err.Error(), serveRunStatus(err, http.StatusBadRequest)) return } writeServeJSON(w, http.StatusOK, result) @@ -128,7 +128,7 @@ func handleSessionsThroughput() http.HandlerFunc { } result, err := RunSessionThroughput(r.Context(), opts) if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) + http.Error(w, err.Error(), serveRunStatus(err, http.StatusBadRequest)) return } writeServeJSON(w, http.StatusOK, result) @@ -139,7 +139,7 @@ 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) + http.Error(w, err.Error(), serveRunStatus(err, http.StatusBadRequest)) return } writeServeJSON(w, http.StatusOK, result) @@ -158,7 +158,7 @@ func handleSessionGet() http.HandlerFunc { 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) + http.Error(w, err.Error(), serveRunStatus(err, http.StatusNotFound)) return } writeServeJSON(w, http.StatusOK, s) diff --git a/pkg/cli/serve_test.go b/pkg/cli/serve_test.go index 28f275e5..6fa5c86d 100644 --- a/pkg/cli/serve_test.go +++ b/pkg/cli/serve_test.go @@ -9,12 +9,13 @@ import ( "strings" "testing" + "github.com/flanksource/captain/pkg/aichat" "github.com/flanksource/captain/pkg/claude" "github.com/flanksource/captain/pkg/database" ) func TestHandleThreadFromAgentCreatesThread(t *testing.T) { - store := newFileThreadStore(filepath.Join(t.TempDir(), "threads.json")) + store := aichat.NewMemoryThreadStore() body := `{"title":"Fix flaky test","providerSessionId":"sess-123","model":"codex-gpt-5-codex"}` req := httptest.NewRequest(http.MethodPost, "/api/captain/chat/threads/from-agent", strings.NewReader(body)) rec := httptest.NewRecorder() @@ -48,7 +49,7 @@ func TestHandleThreadFromAgentCreatesThread(t *testing.T) { } func TestHandleThreadFromAgentRequiresProviderSession(t *testing.T) { - store := newFileThreadStore(filepath.Join(t.TempDir(), "threads.json")) + store := aichat.NewMemoryThreadStore() req := httptest.NewRequest(http.MethodPost, "/api/captain/chat/threads/from-agent", strings.NewReader(`{"title":"missing"}`)) rec := httptest.NewRecorder() @@ -109,8 +110,9 @@ func TestHandleSessionGetReturnsAllMatches(t *testing.T) { 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" { + item := got.Sessions[0] + detail := item.Detail + if detail == nil || detail.ID != item.CaptainID || detail.ProviderSessionID != "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" { diff --git a/pkg/cli/session_get.go b/pkg/cli/session_get.go index df084e36..cef1eb2a 100644 --- a/pkg/cli/session_get.go +++ b/pkg/cli/session_get.go @@ -7,6 +7,7 @@ import ( "strings" "time" + "github.com/flanksource/captain/pkg/aichat" "github.com/flanksource/captain/pkg/api" "github.com/flanksource/captain/pkg/database" "github.com/flanksource/captain/pkg/session" @@ -31,10 +32,10 @@ type SessionGetItem struct { 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"` - transcriptFiltered bool + ActiveRunID string `json:"activeRunId,omitempty"` + Chat *ChatCapabilities `json:"chat,omitempty"` + ChatState *ChatStateFrame `json:"chatState,omitempty"` + notice transcriptNotice } type sessionGetStore interface { @@ -92,8 +93,7 @@ func buildSessionGetItem(ctx context.Context, db sessionGetStore, overview datab item := SessionGetItem{ CaptainID: overview.ID.String(), ProviderSessionID: stringOr(overview.ProviderSessionID, ""), Host: overview.HostID, Aggregate: stringOr(overview.AgentType, "") == "batch", - Summary: recordFromOverview(overview), - transcriptFiltered: len(opts.Tools) > 0 || len(opts.Categories) > 0, + Summary: recordFromOverview(overview), } if overview.ParentSessionID != nil { item.ParentSessionID = overview.ParentSessionID.String() @@ -127,15 +127,40 @@ func buildSessionGetItem(ctx context.Context, db sessionGetStore, overview datab 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) - if err := filterSessionTranscript(detail, opts); err != nil { + notice, err := applyTranscriptWindow(detail, opts, item.Summary.Messages) + if err != nil { return SessionGetItem{}, fmt.Errorf("filter Captain session %s transcript: %w", overview.ID, err) } - pageSessionTranscript(detail, opts) + item.notice = notice item.Detail = detail return item, nil } func loadSessionDetail(ctx context.Context, db sessionGetStore, overview database.SessionOverview) (*session.Session, error) { + if overview.MessageCount > 0 { + captainDB, ok := db.(*database.DB) + if !ok { + return nil, fmt.Errorf("captain session %s has database messages but its store cannot load the canonical aggregate", overview.ID) + } + store, err := aichat.NewDatabaseThreadStore(captainDB) + if err != nil { + return nil, err + } + detail, err := store.GetSession(ctx, overview.ID.String()) + if err != nil { + return nil, fmt.Errorf("load canonical Captain session %s: %w", overview.ID, err) + } + runs, err := db.ListPromptRuns(ctx, database.PromptRunFilter{SessionID: &overview.ID}) + if err != nil { + return nil, fmt.Errorf("list prompt runs for Captain session %s: %w", overview.ID, err) + } + if len(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 + } path := stringOr(overview.HistoryFile, stringOr(overview.Path, "")) var detail *session.Session if path != "" { @@ -146,6 +171,9 @@ func loadSessionDetail(ctx context.Context, db sessionGetStore, overview databas return nil, fmt.Errorf("parse Captain session %s: %w", overview.ID, err) } detail = parsed + detail.ID = overview.ID.String() + detail.ProviderSessionID = stringOr(overview.ProviderSessionID, "") + detail.Revision = overview.StateVersion } stopPromptRuns := rpchttp.Track(ctx, "prompt_runs") runs, err := db.ListPromptRuns(ctx, database.PromptRunFilter{SessionID: &overview.ID}) @@ -153,36 +181,62 @@ func loadSessionDetail(ctx context.Context, db sessionGetStore, overview databas 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 len(runs) > 0 { + if detail == nil { + if detail, err = sessionFromPromptRun(overview, runs[0]); err != nil { + return nil, err + } + } else 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) + } } if detail == nil { - return sessionFromPromptRun(overview, runs[0]) + return nil, nil } - 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) + // The transcript and prompt-run branches build the aggregate from the source + // that produced the session, so neither carries the stored projection the + // database branch gets from GetSession. + if err := applyOverviewProjection(ctx, db, overview, detail); err != nil { + return nil, fmt.Errorf("project Captain session %s overview: %w", overview.ID, err) } return detail, nil } +// applyOverviewProjection adapts the store interface; the projection itself is +// shared with the database branch. +func applyOverviewProjection( + ctx context.Context, + db sessionGetStore, + overview database.SessionOverview, + detail *session.Session, +) error { + store, ok := db.(aichat.OverviewProjectionStore) + if !ok { + return nil + } + return aichat.ApplyOverviewProjection(ctx, store, overview, detail) +} + 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), + ID: overview.ID.String(), + ProviderSessionID: stringOr(overview.ProviderSessionID, ""), + Revision: overview.StateVersion, + 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)) @@ -322,30 +376,73 @@ func (r SessionGetResult) Pretty() clickyapi.Text { return clickyapi.Text{}.Add(list) } +// Tree roots the forest at every session whose parent is absent from the +// result set, not only at sessions with no parent at all. Resolving a provider +// session ID that exists under several sources (the schema allows one row per +// source) returns a mid-thread slice whose parents all live outside the slice, +// and anchoring roots at "" alone rendered that slice as an empty forest. func (r SessionGetResult) Tree() clickyapi.TreeNode { byParent := map[string][]SessionGetItem{} + present := make(map[string]struct{}, len(r.Sessions)) for i := range r.Sessions { byParent[r.Sessions[i].ParentSessionID] = append(byParent[r.Sessions[i].ParentSessionID], r.Sessions[i]) + present[r.Sessions[i].CaptainID] = struct{}{} + } + children := make([]clickyapi.TreeNode, 0, len(r.Sessions)) + rendered := map[string]struct{}{} + addRoot := func(item SessionGetItem) { + node := sessionGetTreeNode{ + item: item, byParent: byParent, seen: map[string]struct{}{item.CaptainID: {}}, + } + markRendered(node, rendered) + children = append(children, node) + } + for i := range r.Sessions { + if _, parented := present[r.Sessions[i].ParentSessionID]; !parented { + addRoot(r.Sessions[i]) + } } - children := make([]clickyapi.TreeNode, 0, len(byParent[""])) - for _, item := range byParent[""] { - children = append(children, sessionGetTreeNode{item: item, byParent: byParent}) + // Every session must appear once. Only a parent cycle can leave one + // unreachable from the roots above, so promote whatever is left over rather + // than dropping it from the render. + for i := range r.Sessions { + if _, shown := rendered[r.Sessions[i].CaptainID]; !shown { + addRoot(r.Sessions[i]) + } } return &clickyapi.ConcreteBranchNode{Children: children} } +func markRendered(node sessionGetTreeNode, rendered map[string]struct{}) { + rendered[node.item.CaptainID] = struct{}{} + for _, child := range node.GetChildren() { + markRendered(child.(sessionGetTreeNode), rendered) + } +} + type sessionGetTreeNode struct { item SessionGetItem byParent map[string][]SessionGetItem + // seen carries the ancestors already rendered on this branch so a cyclic + // parent reference cannot recurse forever now that roots are derived. + seen map[string]struct{} } 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)) + children := make([]clickyapi.TreeNode, 0, len(items)) for i := range items { - children[i] = sessionGetTreeNode{item: items[i], byParent: n.byParent} + if _, cycle := n.seen[items[i].CaptainID]; cycle { + continue + } + seen := make(map[string]struct{}, len(n.seen)+1) + for id := range n.seen { + seen[id] = struct{}{} + } + seen[items[i].CaptainID] = struct{}{} + children = append(children, sessionGetTreeNode{item: items[i], byParent: n.byParent, seen: seen}) } return children } @@ -383,25 +480,14 @@ func (i SessionGetItem) Pretty() clickyapi.Text { 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. +// hiddenRowsNotice reports messages dropped between the session's full count +// and the rendered transcript so a bounded view never reads as the whole +// session. See transcriptNotice for how the causes are attributed. 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{} - } - if i.transcriptFiltered { - return clickyapi.Text{}.NewLine().Append( - fmt.Sprintf(" … showing %d of %d messages after transcript filters and windowing", - len(i.Detail.Messages), i.Summary.Messages), "text-amber-600") - } - 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") + return i.notice.text(len(i.Detail.Messages)) } type sessionGetListItem struct { diff --git a/pkg/cli/session_get_compact_test.go b/pkg/cli/session_get_compact_test.go index a1ad7f9f..d64fdab1 100644 --- a/pkg/cli/session_get_compact_test.go +++ b/pkg/cli/session_get_compact_test.go @@ -31,6 +31,29 @@ func transcriptFixture(messages, events int) *session.Session { return s } +// reasoningTranscriptFixture interleaves plain assistant text with reasoning +// parts so `--category '!reasoning'` has something to exclude. +func reasoningTranscriptFixture(kept, reasoning int) *session.Session { + s := &session.Session{ID: "d3521f3b-38a2-43b7-b80f-77450a9cb30c", Source: "codex"} + for i := 0; i < max(kept, reasoning); i++ { + if i < reasoning { + s.Messages = append(s.Messages, session.Message{ + ID: fmt.Sprintf("r%d", i), + Role: "assistant", + Parts: []session.Part{{Type: session.PartReasoning, Text: fmt.Sprintf("thinking %d", i)}}, + }) + } + if i < kept { + s.Messages = append(s.Messages, session.Message{ + ID: fmt.Sprintf("m%d", i), + Role: "assistant", + Parts: []session.Part{{Type: session.PartText, Text: fmt.Sprintf("message %d", i)}}, + }) + } + } + return s +} + var _ = Describe("session get transcript paging", func() { It("bounds events alongside messages under --tail", func() { detail := transcriptFixture(50, 40) @@ -188,42 +211,119 @@ var _ = Describe("session get header", func() { }) It("reports how many transcript rows the window hid", func() { + opts := SessionGetOptions{Tail: 5} detail := transcriptFixture(200, 0) - pageSessionTranscript(detail, SessionGetOptions{Tail: 5}) + notice, err := applyTranscriptWindow(detail, opts, 200) + Expect(err).NotTo(HaveOccurred()) item := SessionGetItem{ CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", Summary: SessionRecord{Messages: 200}, Detail: detail, + notice: notice, } rendered := item.Pretty().String() - Expect(rendered).To(ContainSubstring("195 of 200 messages hidden")) + Expect(rendered).To(ContainSubstring("195 of 200 messages hidden by --tail 5")) Expect(rendered).To(ContainSubstring("--limit 0")) }) It("stays silent when the whole transcript is shown", func() { + opts := SessionGetOptions{Limit: 1000} detail := transcriptFixture(4, 0) + notice, err := applyTranscriptWindow(detail, opts, 4) + Expect(err).NotTo(HaveOccurred()) item := SessionGetItem{ CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", Summary: SessionRecord{Messages: 4}, Detail: detail, + notice: notice, } Expect(item.Pretty().String()).NotTo(ContainSubstring("hidden")) }) - It("does not claim --limit 0 restores rows excluded by filters", func() { + // Regression: `sessions get -c '!reasoning' -l 1000` used to report + // "after transcript filters and windowing", which reads as if --limit 1000 + // had truncated the transcript when the category filter alone removed the + // rows. The two causes are now attributed separately. + It("attributes filter exclusions to the filter, not to the window", func() { + opts := SessionGetOptions{Categories: []string{"!reasoning"}, Limit: 1000} + detail := reasoningTranscriptFixture(148, 117) + notice, err := applyTranscriptWindow(detail, opts, 265) + Expect(err).NotTo(HaveOccurred()) item := SessionGetItem{ - CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", - Summary: SessionRecord{Messages: 4}, - Detail: transcriptFixture(1, 0), - transcriptFiltered: true, + CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", + Summary: SessionRecord{Messages: 265}, + Detail: detail, + notice: notice, } rendered := item.Pretty().String() - Expect(rendered).To(ContainSubstring("showing 1 of 4 messages after transcript filters and windowing")) - Expect(rendered).NotTo(ContainSubstring("--limit 0")) + Expect(detail.Messages).To(HaveLen(148)) + Expect(rendered).To(ContainSubstring("117 of 265 messages excluded by --category '!reasoning'")) + Expect(rendered).To(ContainSubstring("none hidden by --limit 1000")) + Expect(rendered).NotTo(ContainSubstring("--limit 0"), + "--limit 0 cannot restore rows the category filter excluded") }) + + It("reports filter exclusions and window truncation as separate causes", func() { + opts := SessionGetOptions{Categories: []string{"!reasoning"}, Limit: 100} + detail := reasoningTranscriptFixture(148, 117) + notice, err := applyTranscriptWindow(detail, opts, 265) + Expect(err).NotTo(HaveOccurred()) + item := SessionGetItem{ + CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", + Summary: SessionRecord{Messages: 265}, + Detail: detail, + notice: notice, + } + + rendered := item.Pretty().String() + + Expect(detail.Messages).To(HaveLen(100)) + Expect(rendered).To(ContainSubstring("117 of 265 messages excluded by --category '!reasoning'")) + Expect(rendered).To(ContainSubstring("48 more hidden by --limit 100")) + Expect(rendered).To(ContainSubstring("--limit 0")) + }) + + It("reports rows the overview counted but the recorded transcript lacks", func() { + opts := SessionGetOptions{Limit: 1000} + detail := transcriptFixture(4, 0) + notice, err := applyTranscriptWindow(detail, opts, 6) + Expect(err).NotTo(HaveOccurred()) + item := SessionGetItem{ + CaptainID: "d81f885d-3d60-47c0-8122-a8124f2fbdd1", + Summary: SessionRecord{Messages: 6}, + Detail: detail, + notice: notice, + } + + Expect(item.Pretty().String()).To(ContainSubstring("2 not present in the recorded transcript")) + }) +}) + +var _ = Describe("session get transcript flag descriptions", func() { + DescribeTable("renders the flags that bounded the transcript", + func(opts SessionGetOptions, filters, window string) { + Expect(transcriptFilterFlags(opts)).To(Equal(filters)) + Expect(transcriptWindowFlags(opts)).To(Equal(window)) + }, + Entry("negated category and an explicit limit", + SessionGetOptions{Categories: []string{"!reasoning"}, Limit: 1000}, + "--category '!reasoning'", "--limit 1000"), + Entry("bare category needs no quoting", + SessionGetOptions{Categories: []string{"test"}, Limit: 200}, + "--category test", "--limit 200"), + Entry("tool globs are quoted and tail wins over offset/limit", + SessionGetOptions{Tools: []string{"B*"}, Offset: 10, Limit: 5, Tail: 3}, + "--tool 'B*'", "--tail 3"), + Entry("offset joins the limit", + SessionGetOptions{Offset: 10, Limit: 5}, + "", "--offset 10 --limit 5"), + Entry("cleared limit is reported as unbounded", + SessionGetOptions{Limit: 0}, + "", "--limit 0"), + ) }) diff --git a/pkg/cli/session_get_multi_test.go b/pkg/cli/session_get_multi_test.go index 96813a70..4723bbfe 100644 --- a/pkg/cli/session_get_multi_test.go +++ b/pkg/cli/session_get_multi_test.go @@ -386,7 +386,7 @@ var _ = Describe("session get multi-result output", func() { "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}} + "detail": {"id":"ad4c854e-cde6-4b99-99f3-667bf74112e3","revision":0,"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", diff --git a/pkg/cli/session_get_notice.go b/pkg/cli/session_get_notice.go new file mode 100644 index 00000000..8a4890d2 --- /dev/null +++ b/pkg/cli/session_get_notice.go @@ -0,0 +1,122 @@ +package cli + +import ( + "fmt" + "strings" + + "github.com/flanksource/captain/pkg/session" + clickyapi "github.com/flanksource/clicky/api" +) + +// transcriptNotice records why a rendered transcript is shorter than the +// session's message count. Filter exclusions and window truncation are counted +// separately: --limit 0 restores windowed rows but can never restore rows a +// --tool/--category filter removed, so blaming one cause for the other's rows +// sends the reader chasing a bigger --limit that changes nothing. +type transcriptNotice struct { + total int + filterExcluded int + windowHidden int + filterFlags string + windowFlags string +} + +// applyTranscriptWindow filters then windows the transcript, attributing the +// dropped messages to whichever step removed them. total is the session's full +// message count before either step runs. +func applyTranscriptWindow(detail *session.Session, opts SessionGetOptions, total int) (transcriptNotice, error) { + notice := transcriptNotice{ + total: total, + filterFlags: transcriptFilterFlags(opts), + windowFlags: transcriptWindowFlags(opts), + } + before := len(detail.Messages) + if err := filterSessionTranscript(detail, opts); err != nil { + return notice, err + } + notice.filterExcluded = before - len(detail.Messages) + + before = len(detail.Messages) + pageSessionTranscript(detail, opts) + notice.windowHidden = before - len(detail.Messages) + return notice, nil +} + +func (n transcriptNotice) text(shown int) clickyapi.Text { + clauses := n.clauses(shown) + if len(clauses) == 0 { + return clickyapi.Text{} + } + text := clickyapi.Text{}.NewLine().Append(" … "+strings.Join(clauses, "; "), "text-amber-600") + // Only the window is recoverable, and the advice earns its own line so a + // multi-cause notice does not run past the terminal width and clip it. + if n.windowHidden > 0 { + text = text.NewLine().Append(" use --limit 0 for the full transcript", "text-amber-600") + } + return text +} + +func (n transcriptNotice) clauses(shown int) []string { + var clauses []string + if n.filterExcluded > 0 { + clauses = append(clauses, fmt.Sprintf("%d of %d messages excluded by %s", + n.filterExcluded, n.total, n.filterFlags)) + } + switch { + case n.windowHidden > 0 && len(clauses) > 0: + clauses = append(clauses, fmt.Sprintf("%d more hidden by %s", n.windowHidden, n.windowFlags)) + case n.windowHidden > 0: + clauses = append(clauses, fmt.Sprintf("%d of %d messages hidden by %s", + n.windowHidden, n.total, n.windowFlags)) + case len(clauses) > 0: + clauses = append(clauses, "none hidden by "+n.windowFlags) + } + // The overview's message count and the recorded transcript can disagree when + // a session was ingested but never fully parsed; report the gap rather than + // letting it read as a complete transcript. + if missing := n.total - shown - n.filterExcluded - n.windowHidden; missing > 0 { + clauses = append(clauses, fmt.Sprintf("%d not present in the recorded transcript", missing)) + } + return clauses +} + +func transcriptFilterFlags(opts SessionGetOptions) string { + parts := make([]string, 0, len(opts.Tools)+len(opts.Categories)) + for _, tool := range opts.Tools { + parts = append(parts, "--tool "+quoteFlagValue(tool)) + } + for _, category := range opts.Categories { + parts = append(parts, "--category "+quoteFlagValue(category)) + } + return strings.Join(parts, " ") +} + +func transcriptWindowFlags(opts SessionGetOptions) string { + if opts.Tail > 0 { + return fmt.Sprintf("--tail %d", opts.Tail) + } + if opts.Offset > 0 { + return fmt.Sprintf("--offset %d --limit %d", opts.Offset, opts.Limit) + } + return fmt.Sprintf("--limit %d", opts.Limit) +} + +// quoteFlagValue echoes a filter value the way the user had to type it, so +// shell-significant patterns such as !reasoning or B* stay copy-pasteable. +func quoteFlagValue(value string) string { + if value != "" && strings.IndexFunc(value, flagValueNeedsQuote) < 0 { + return value + } + return "'" + strings.ReplaceAll(value, "'", `'\''`) + "'" +} + +func flagValueNeedsQuote(r rune) bool { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9': + return false + case r == '-', r == '_', r == '.', r == '/', r == ':': + return false + default: + return true + } +} diff --git a/pkg/cli/session_get_tree_ginkgo_test.go b/pkg/cli/session_get_tree_ginkgo_test.go new file mode 100644 index 00000000..791f5a9e --- /dev/null +++ b/pkg/cli/session_get_tree_ginkgo_test.go @@ -0,0 +1,73 @@ +package cli + +import ( + clickyapi "github.com/flanksource/clicky/api" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// collectTreeIDs walks the rendered forest and returns every Captain ID in +// visit order, so a dropped subtree shows up as a missing ID rather than as +// silently empty output. +func collectTreeIDs(nodes []clickyapi.TreeNode) []string { + ids := []string{} + for _, node := range nodes { + typed, ok := node.(sessionGetTreeNode) + if !ok { + continue + } + ids = append(ids, typed.item.CaptainID) + ids = append(ids, collectTreeIDs(typed.GetChildren())...) + } + return ids +} + +func treeChildren(result SessionGetResult) []clickyapi.TreeNode { + branch, ok := result.Tree().(*clickyapi.ConcreteBranchNode) + Expect(ok).To(BeTrue()) + return branch.Children +} + +var _ = Describe("session get tree", func() { + const ( + rootID = "11de1d82-7622-494f-8219-3ca32bd13dff" + orchestration = "2832bf9b-4aca-5bcd-8127-94eb8dc505d1" + transcript = "a255a8d1-a1c9-4da2-a420-51bbd2c8a561" + ) + + It("renders a mid-thread slice whose parents are all outside the result set", func() { + // Resolving a provider session ID shared across sources returns the + // orchestration row and the transcript row it parents; the thread root + // is not in the slice, so no item has an empty parent. + result := SessionGetResult{Sessions: []SessionGetItem{ + {CaptainID: transcript, ParentSessionID: orchestration, RootSessionID: rootID}, + {CaptainID: orchestration, ParentSessionID: rootID, RootSessionID: rootID}, + }, Total: 2} + + Expect(collectTreeIDs(treeChildren(result))).To(Equal([]string{orchestration, transcript})) + }) + + It("still nests children under a root present in the result set", func() { + result := SessionGetResult{Sessions: []SessionGetItem{ + {CaptainID: rootID}, + {CaptainID: orchestration, ParentSessionID: rootID}, + {CaptainID: transcript, ParentSessionID: orchestration}, + }, Total: 3} + + children := treeChildren(result) + + Expect(children).To(HaveLen(1)) + Expect(collectTreeIDs(children)).To(Equal([]string{rootID, orchestration, transcript})) + }) + + It("renders every session when parents form a cycle", func() { + result := SessionGetResult{Sessions: []SessionGetItem{ + {CaptainID: orchestration, ParentSessionID: transcript}, + {CaptainID: transcript, ParentSessionID: orchestration}, + }, Total: 2} + + // Both parents are present, so the cycle yields no natural root. The + // render must still terminate and must not swallow either session. + Expect(collectTreeIDs(treeChildren(result))).To(ConsistOf(orchestration, transcript)) + }) +}) diff --git a/pkg/cli/session_live.go b/pkg/cli/session_live.go index 9c53487a..96f0945f 100644 --- a/pkg/cli/session_live.go +++ b/pkg/cli/session_live.go @@ -67,6 +67,7 @@ func RunSessionLive(ctx context.Context, opts SessionLiveOptions) (SessionLiveRe return buildSessionLiveResult(sessionLiveResultOptions{ Page: page, Source: source, Scope: scope, Project: projectResultValue(scope, projectRoot), ReadAt: time.Now().UTC(), DatabaseCoverage: coverage, + DatabaseContext: activeDatabaseContextName(ctx), }), nil } @@ -77,6 +78,7 @@ type sessionLiveResultOptions struct { Project string ReadAt time.Time DatabaseCoverage string + DatabaseContext string } func buildSessionLiveResult(options sessionLiveResultOptions) SessionLiveResult { @@ -84,7 +86,7 @@ func buildSessionLiveResult(options sessionLiveResultOptions) SessionLiveResult if coverage == "" { coverage = "page" } - databaseStatus := sessionDatabaseStatus(options.Page.Records, options.ReadAt) + databaseStatus := sessionDatabaseStatus(options.DatabaseContext, options.Page.Records, options.ReadAt) databaseStatus.Coverage = coverage return SessionLiveResult{ Sessions: options.Page.Records, Total: options.Page.Total, Source: options.Source, @@ -125,8 +127,13 @@ func enrichLiveSessionSurfaces(records []SessionRecord) { } } -func sessionDatabaseStatus(records []SessionRecord, readAt time.Time) SessionDatabaseStatusWire { - dsn, source := captainDatabaseIdentity() +// sessionDatabaseStatus reports the database the records were actually read +// from, which is the active context rather than the monitored default. +func sessionDatabaseStatus(contextName string, records []SessionRecord, readAt time.Time) SessionDatabaseStatusWire { + if contextName == "" { + contextName = defaultDatabaseContextName + } + dsn, source := contextDatabaseIdentity(contextName) status := SessionDatabaseStatusWire{Source: source, DSN: database.MaskDSN(dsn), ReadAt: readAt} for _, record := range records { if record.Live == nil { diff --git a/pkg/cli/session_record_db.go b/pkg/cli/session_record_db.go index ecb22f59..390e1ac0 100644 --- a/pkg/cli/session_record_db.go +++ b/pkg/cli/session_record_db.go @@ -122,24 +122,6 @@ func candidateFromOverview(overview database.SessionOverview) sessionCandidate { } } -// 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 "" @@ -183,7 +165,7 @@ func overviewFromSummary(summary database.SessionListSummary) database.SessionOv // recordFromOverview projects one overview row to the SessionRecord wire shape. func recordFromOverview(overview database.SessionOverview) SessionRecord { - metadata := overviewMetadata(overview) + metadata := session.DecodeMetadata(overview.Metadata) path := stringOr(overview.HistoryFile, stringOr(overview.Path, "")) id := stringOr(overview.ProviderSessionID, overview.ID.String()) record := SessionRecord{ diff --git a/pkg/cli/sessions_test.go b/pkg/cli/sessions_test.go index 2f6e8a90..61792b41 100644 --- a/pkg/cli/sessions_test.go +++ b/pkg/cli/sessions_test.go @@ -98,8 +98,9 @@ func TestRunSessionListAndGetClaude(t *testing.T) { if detail.Total != 1 || len(detail.Sessions) != 1 || detail.Sessions[0].Detail == nil { t.Fatalf("session detail result = %+v", detail) } - parsed := detail.Sessions[0].Detail - if parsed.Source != "claude" || parsed.ID != "sess-claude" { + item := detail.Sessions[0] + parsed := item.Detail + if parsed.Source != "claude" || parsed.ID != item.CaptainID || parsed.ProviderSessionID != "sess-claude" { t.Fatalf("session detail meta = %+v", parsed) } entries := parsed.ToReplayEntries() diff --git a/pkg/cli/stdin_claude_command_test.go b/pkg/cli/stdin_claude_command_test.go index ca56667d..c741e712 100644 --- a/pkg/cli/stdin_claude_command_test.go +++ b/pkg/cli/stdin_claude_command_test.go @@ -5,6 +5,7 @@ import ( "github.com/flanksource/captain/pkg/session" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/segmentio/encoding/json" ) // claudeGoalTranscript is the exact three-record shape Claude writes for a @@ -14,6 +15,8 @@ const claudeGoalTranscript = `{"type":"attachment","uuid":"uuid-goal","sessionId {"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"}}` +const claudeWrappedShellTranscript = `{"type":"assistant","sessionId":"s1","uuid":"uuid-bash","timestamp":"2026-07-14T12:16:09.300Z","cwd":"/repo","message":{"role":"assistant","content":[{"type":"tool_use","id":"tool-1","name":"Bash","input":{"command":"/bin/zsh -lc 'pnpm test'"}}]}}` + func historyToolNames(uses []claude.ToolUse) []string { names := make([]string, 0, len(uses)) for _, u := range uses { @@ -70,4 +73,22 @@ var _ = Describe("Claude /goal transcript history from a reader", func() { Expect(cmdSummary).To(ContainSubstring("/goal")) Expect(outSummary).To(ContainSubstring("Goal set: ship the docker build")) }) + + It("serializes transformed shell input unless raw history was requested", func() { + out, err := runHistoryFromReader([]byte(claudeWrappedShellTranscript), HistoryOptions{Limit: 1}) + Expect(err).NotTo(HaveOccurred()) + encoded, err := json.Marshal(out) + Expect(err).NotTo(HaveOccurred()) + Expect(string(encoded)).To(And( + ContainSubstring("zsh"), + ContainSubstring("pnpm test"), + Not(ContainSubstring("/bin/zsh -lc")), + )) + + rawOut, err := runHistoryFromReader([]byte(claudeWrappedShellTranscript), HistoryOptions{Limit: 1, Raw: true}) + Expect(err).NotTo(HaveOccurred()) + rawEncoded, err := json.Marshal(rawOut) + Expect(err).NotTo(HaveOccurred()) + Expect(string(rawEncoded)).To(ContainSubstring(`/bin/zsh -lc`)) + }) }) diff --git a/pkg/cli/webapp/dist/index.html b/pkg/cli/webapp/dist/index.html index a48eaf9a..902b3a23 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 79c4bd16..98e89f31 100644 --- a/pkg/cli/webapp/package.json +++ b/pkg/cli/webapp/package.json @@ -31,8 +31,8 @@ "@testing-library/react": "^16.1.0", "@tailwindcss/vite": "^4.2.2", "@types/node": "^24.0.0", - "@types/react": "18.3.28", - "@types/react-dom": "18.3.7", + "@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", diff --git a/pkg/cli/webapp/pnpm-lock.yaml b/pkg/cli/webapp/pnpm-lock.yaml index 9fb73fe6..84dfa030 100644 --- a/pkg/cli/webapp/pnpm-lock.yaml +++ b/pkg/cli/webapp/pnpm-lock.yaml @@ -5,8 +5,8 @@ settings: excludeLinksFromLockfile: false overrides: - '@types/react': 18.3.28 - '@types/react-dom': 18.3.7 + '@types/react': 19.2.17 + '@types/react-dom': 19.2.3 importers: @@ -66,16 +66,16 @@ importers: version: 6.9.1 '@testing-library/react': specifier: ^16.1.0 - version: 16.3.2(@testing-library/dom@10.4.1)(@types/react-dom@18.3.7(@types/react@18.3.28))(@types/react@18.3.28)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) + 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/node': specifier: ^24.0.0 version: 24.13.3 '@types/react': - specifier: 18.3.28 - version: 18.3.28 + specifier: 19.2.17 + version: 19.2.17 '@types/react-dom': - specifier: 18.3.7 - version: 18.3.7(@types/react@18.3.28) + specifier: 19.2.3 + version: 19.2.3(@types/react@19.2.17) '@vitejs/plugin-react': specifier: ^5.0.4 version: 5.2.0(vite@7.3.6(@types/node@24.13.3)(jiti@2.7.0)(lightningcss@1.32.0)(yaml@2.9.0)) @@ -717,8 +717,8 @@ packages: engines: {node: '>=18'} peerDependencies: '@testing-library/dom': ^10.0.0 - '@types/react': 18.3.28 - '@types/react-dom': 18.3.7 + '@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: @@ -865,16 +865,13 @@ packages: '@types/node@24.13.3': resolution: {integrity: sha512-Dh8vAsV36ig5wa9OX4pXvMc9D3Veibfw2wix0CUwYODLD8nkj9UsLjASr49nPg+2eKzxhBV+v7L8pXvT4e639Q==} - '@types/prop-types@15.7.15': - resolution: {integrity: sha512-F6bEyamV9jKGAFBEmlQnesRPGOQqS2+Uwi0Em15xenOxHaf2hv6L8YCVn3rPdPJOiJfPiCnLIRyvwVaqMY3MIw==} - - '@types/react-dom@18.3.7': - resolution: {integrity: sha512-MEe3UeoENYVFXzoXEWsvcpg6ZvlrFNlOQ7EOsvhI3CfAXwzPfO8Qwuxd40nepsYKqyyVQnTdEfv68q91yLcKrQ==} + '@types/react-dom@19.2.3': + resolution: {integrity: sha512-jp2L/eY6fn+KgVVQAOqYItbF0VY/YApe5Mz2F0aykSO8gx31bYCZyvSeYxCHKvzHG5eZjc+zyaS5BrBWya2+kQ==} peerDependencies: - '@types/react': 18.3.28 + '@types/react': 19.2.17 - '@types/react@18.3.28': - resolution: {integrity: sha512-z9VXpC7MWrhfWipitjNdgCauoMLRdIILQsAEV+ZesIzBq/oUlxk0m3ApZuMFCXdnS4U7KrI+l3WRUEGQ8K1QKw==} + '@types/react@19.2.17': + resolution: {integrity: sha512-MXfmqaVPEVgkBT/aY0aGCkRWWtByiYQXo3xdQ8r5RzuFrPiRn8Gar2tQdXSUQ2GKV3bkXckek89V8wQBY2Q/Aw==} '@types/trusted-types@2.0.7': resolution: {integrity: sha512-ScaPdn1dQczgbl0QFTeTOmVHFULt394XJgOQNoyVhZ6r2vLnMLJfBPd53SB52T/3G36VI1/g2MZaX0cwDuXsfw==} @@ -2692,15 +2689,15 @@ snapshots: picocolors: 1.1.1 redent: 3.0.0 - '@testing-library/react@16.3.2(@testing-library/dom@10.4.1)(@types/react-dom@18.3.7(@types/react@18.3.28))(@types/react@18.3.28)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)': + '@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': 18.3.28 - '@types/react-dom': 18.3.7(@types/react@18.3.28) + '@types/react': 19.2.17 + '@types/react-dom': 19.2.3(@types/react@19.2.17) '@types/aria-query@5.0.4': {} @@ -2875,15 +2872,12 @@ snapshots: dependencies: undici-types: 7.18.2 - '@types/prop-types@15.7.15': {} - - '@types/react-dom@18.3.7(@types/react@18.3.28)': + '@types/react-dom@19.2.3(@types/react@19.2.17)': dependencies: - '@types/react': 18.3.28 + '@types/react': 19.2.17 - '@types/react@18.3.28': + '@types/react@19.2.17': dependencies: - '@types/prop-types': 15.7.15 csstype: 3.2.3 '@types/trusted-types@2.0.7': diff --git a/pkg/cli/webapp/pnpm-workspace.yaml b/pkg/cli/webapp/pnpm-workspace.yaml index f93d9c40..188e72f4 100644 --- a/pkg/cli/webapp/pnpm-workspace.yaml +++ b/pkg/cli/webapp/pnpm-workspace.yaml @@ -7,5 +7,5 @@ trustPolicy: no-downgrade allowBuilds: esbuild: true overrides: - '@types/react': 18.3.28 - '@types/react-dom': 18.3.7 + '@types/react': 19.2.17 + '@types/react-dom': 19.2.3 diff --git a/pkg/cli/webapp/src/AgentLauncher.tsx b/pkg/cli/webapp/src/AgentLauncher.tsx index 4570547e..d6eff4d5 100644 --- a/pkg/cli/webapp/src/AgentLauncher.tsx +++ b/pkg/cli/webapp/src/AgentLauncher.tsx @@ -8,6 +8,7 @@ import { } from "@flanksource/clicky-ui/rpc"; import { useChatWindowManager } from "@flanksource/clicky-ui/ai"; import { apiClient } from "./api"; +import { isReadOnlyDbContext } from "./dbContext"; import { agentModelFor, extractSessionId, @@ -86,6 +87,21 @@ export function AgentLauncher({ onNavigate }: AgentLauncherProps) { ); } + // Launching an agent creates a chat thread, which a read-only database + // context rejects. The form is withheld rather than left to fail: its submit + // control lives inside clicky-ui's OperationCommandPage. + if (isReadOnlyDbContext()) { + return ( +
+
+ Launching agents is disabled while a read-only database context is selected. + Switch back to the monitored database from the project picker to launch one. +
+ +
+ ); + } + if (!operation) { return (
diff --git a/pkg/cli/webapp/src/ChatLayer.tsx b/pkg/cli/webapp/src/ChatLayer.tsx index 89c397f8..d7a52383 100644 --- a/pkg/cli/webapp/src/ChatLayer.tsx +++ b/pkg/cli/webapp/src/ChatLayer.tsx @@ -3,20 +3,29 @@ import { ChatFab, ChatWindowLayer } from "@flanksource/clicky-ui/ai"; import { clickyOperationsToTools } from "@flanksource/clicky-ui/chat"; import { useOperations } from "@flanksource/clicky-ui/rpc"; import { apiClient } from "./api"; +import { isReadOnlyDbContext } from "./dbContext"; import { isChatToolOperation } from "./session"; export function ChatLayer() { const { operations } = useOperations(apiClient); + // Every chat action creates or appends to a thread, and a read-only database + // context rejects those writes. The composer is withheld entirely rather than + // left to fail: the chat transport is built inside clicky-ui, so there is no + // per-control disable to reach from here. + const readOnly = isReadOnlyDbContext(); const tools = useMemo( () => clickyOperationsToTools(operations.filter(isChatToolOperation)), [operations], ); + if (readOnly) return null; + return ( <> { - it("offers no cmux mode once the served catalog marks it disabled", () => { - render( - , - ); - - const modes = within( - screen.getByRole("radiogroup", { name: "Runtime mode" }), - ).getAllByRole("radio"); - expect(modes.map((mode) => mode.textContent)).toEqual([ - "API", - "Agent", - "CLI", - ]); - }); - - // The registry has one Anthropic provider with four modes where the offline - // default split it into "Claude" (agent/cli/cmux) and "Anthropic" (api). - it("renders one family per provider rather than one per backend group", () => { - render( - , - ); - - const families = within( - screen.getByRole("radiogroup", { name: "Provider family" }), - ).getAllByRole("radio"); - expect(families.map((family) => family.textContent)).toEqual([ - "Claude", - "Codex", - ]); - }); - - it("offers only the effort tiers the server served", () => { - render( - , - ); - - fireEvent.focus(screen.getByRole("combobox", { name: "Reasoning effort" })); - - expect(screen.getAllByRole("option").map((option) => option.textContent)).toEqual([ - "None", - "Low", - "Extra high", - ]); - }); - - it("hides the effort control when the server served no tiers", () => { - render( - , - ); - - expect(screen.queryByRole("combobox", { name: "Reasoning effort" })).toBeNull(); - }); - - it("adds a runtime with the primary backend and an intentionally blank model", () => { - expect( - addRuntimeRow([{ backend: "codex-cmux", model: "gpt-5.6-sol" }]), - ).toEqual([ - { backend: "codex-cmux", model: "gpt-5.6-sol" }, - { backend: "codex-cmux" }, - ]); - }); - - it("rejects incomplete and duplicate comparison rows", () => { - expect( - validateRuntimeRows([ - { backend: "codex-cmux", model: "gpt-5.6-sol" }, - { backend: "codex-cmux" }, - ]), - ).toEqual("Runtime 2 needs a model"); - expect( - validateRuntimeRows([ - { backend: "codex-cmux", model: "gpt-5.6-sol", effort: "high" }, - { backend: "codex-cmux", model: "gpt-5.6-sol", effort: "high" }, - ]), - ).toEqual("Runtime 2 duplicates codex-cmux:gpt-5.6-sol:high"); - }); -}); diff --git a/pkg/cli/webapp/src/PromptRuntimeRows.tsx b/pkg/cli/webapp/src/PromptRuntimeRows.tsx deleted file mode 100644 index a580f86e..00000000 --- a/pkg/cli/webapp/src/PromptRuntimeRows.tsx +++ /dev/null @@ -1,142 +0,0 @@ -import { Button } from "@flanksource/clicky-ui/components"; -import { Icon, UiAdd, UiTrash } from "@flanksource/clicky-ui/data"; -import { - RuntimeModePicker, - SPEC_RUNTIME_FAMILIES, - effortOptionsForModel, - familyById, - modelsForFamily, - reconcileModelCapabilities, - selectionForBackend, - type AISpecRuntimeValue, - type SpecRuntimeFamily, -} from "@flanksource/clicky-ui/ai"; -import { - EffortSelector, - ModelSelector, - type ChatModel, -} from "@flanksource/clicky-ui/chat"; -import { - addRuntimeRow, - validateRuntimeRows, -} from "./promptRuntimeRowsHelpers"; -import { backendForRow } from "./promptWorkbenchHelpers"; - -export function PromptRuntimeRows({ - rows, - models, - families = SPEC_RUNTIME_FAMILIES, - efforts: effortUniverse = [], - onChange, -}: { - rows: AISpecRuntimeValue[]; - models: ChatModel[]; - /** - * The runtime catalog the server projected from its model registry, already - * stripped of the backends the user disabled. Falling back to the offline - * default would re-offer them, so the schema query is the only source. - */ - families?: SpecRuntimeFamily[]; - /** - * Fallback tiers for a model whose catalog entry carries no supportedEfforts. - * The prompt schema serves this and has already dropped disabled tiers, so an - * empty list means the server said nothing — not that every tier is off. - */ - efforts?: string[]; - onChange: (rows: AISpecRuntimeValue[]) => void; -}) { - const error = validateRuntimeRows(rows); - const update = (index: number, value: AISpecRuntimeValue) => - onChange(rows.map((row, rowIndex) => (rowIndex === index ? value : row))); - - return ( -
- {rows.map((row, index) => { - const backend = backendForRow(row, models); - const selection = selectionForBackend(families, backend); - const family = familyById(families, selection.family); - const availableModels = modelsForFamily(models, family, backend); - const selectedModel = models.find((model) => model.id === row.model); - const efforts = effortOptionsForModel(selectedModel, effortUniverse); - return ( -
-
- - Runtime {index + 1} - - {index > 0 && ( - - )} -
- update(index, value)} - models={models} - families={families} - /> -
- - {efforts.length > 0 && ( - - )} -
-
- ); - })} -
- - {rows.length > 1 && error && ( - {error} - )} -
-
- ); -} diff --git a/pkg/cli/webapp/src/PromptWorkbench.test.ts b/pkg/cli/webapp/src/PromptWorkbench.test.ts index 0d15c4c7..de154f16 100644 --- a/pkg/cli/webapp/src/PromptWorkbench.test.ts +++ b/pkg/cli/webapp/src/PromptWorkbench.test.ts @@ -1,8 +1,8 @@ import { describe, expect, it } from "vitest"; +import type { ChatModel } from "@flanksource/clicky-ui/chat"; import { + mergePromptModelCatalogs, promptOptions, - runtimeModelsPayload, - runtimeRowsFromPrompt, } from "./promptWorkbenchHelpers"; function prompt( @@ -69,11 +69,14 @@ describe("promptOptions", () => { it("does not duplicate a selected prompt already present in the list", () => { const selected = prompt("alpha", "embedded"); - const options = promptOptions([selected, prompt("beta", "local")], selected); - - expect(options.filter((option) => option.value === selected.id)).toHaveLength( - 1, + const options = promptOptions( + [selected, prompt("beta", "local")], + selected, ); + + expect( + options.filter((option) => option.value === selected.id), + ).toHaveLength(1); }); it("titles each option with its description, falling back to the path", () => { @@ -89,62 +92,57 @@ describe("promptOptions", () => { }); }); -describe("runtimeRowsFromPrompt", () => { - it("uses prompt-declared runtimes as the initial comparison rows", () => { - expect( - runtimeRowsFromPrompt({ - ...prompt("compare", "local"), - runtimes: [ - { - model: "gemini-3.5-flash", - backend: "gemini", - effort: "high", - }, - { - model: "claude-sonnet-5", - backend: "anthropic", - effort: "medium", - }, - ], - }), - ).toEqual([ - { - model: "gemini-3.5-flash", - backend: "gemini", - effort: "high", - }, - { - model: "claude-sonnet-5", - backend: "anthropic", - effort: "medium", - }, +describe("mergePromptModelCatalogs", () => { + const model = ( + id: string, + backend: string, + configured: boolean, + state: "available" | "disabled" = "available", + ): ChatModel => ({ + id, + provider: backend.startsWith("claude") ? "claude-agent" : "anthropic", + label: `${backend} ${id}`, + runtime: { model: id, backend }, + reasoning: true, + configured, + availability: { state }, + }); + + it("keeps prompt selections and adds only distinct unavailable status rows", () => { + const promptModels = [ + model("claude-opus-5", "claude-agent", true), + model("claude-opus-5", "claude-cli", true), + ]; + const result = mergePromptModelCatalogs(promptModels, [ + model("claude-opus-5", "claude-agent", true), + model("claude-sonnet-5", "claude-agent", false, "disabled"), + model("claude-opus-5", "anthropic", false, "disabled"), + ]); + + expect(result).toEqual([ + ...promptModels, + model("claude-sonnet-5", "claude-agent", false, "disabled"), + model("claude-opus-5", "anthropic", false, "disabled"), ]); }); - it("preserves model options when declared rows become a run override", () => { + it("treats an exact prompt model as authoritative over stale availability", () => { + const promptModel = model("claude-opus-5", "claude-agent", true); + expect( - runtimeModelsPayload( - [ - { - model: "gemini-3.5-flash", - backend: "gemini", - effort: "high", - temperature: 0, - noCache: true, - fallbacks: [{ model: "gemini-3-flash" }], - }, - ], - [], + mergePromptModelCatalogs( + [promptModel], + [model("claude-opus-5", "claude-agent", false, "disabled")], ), - ).toEqual([ - { - model: "gemini-3.5-flash", - backend: "gemini", - effort: "high", - temperature: 0, - noCache: true, - fallbacks: [{ model: "gemini-3-flash" }], - }, - ]); + ).toEqual([promptModel]); + }); + + it("adds a configured-false model even when no availability detail was served", () => { + const unavailable = { + ...model("claude-sonnet-5", "claude-agent", false), + availability: undefined, + }; + + expect(mergePromptModelCatalogs([], [unavailable])).toEqual([unavailable]); }); }); diff --git a/pkg/cli/webapp/src/PromptWorkbench.tsx b/pkg/cli/webapp/src/PromptWorkbench.tsx index b4d9c5c1..a0c232f0 100644 --- a/pkg/cli/webapp/src/PromptWorkbench.tsx +++ b/pkg/cli/webapp/src/PromptWorkbench.tsx @@ -1,10 +1,9 @@ -import { useMemo, useReducer, useState, type ReactNode } from "react"; +import { useMemo, useReducer, useState } from "react"; import { useQuery } from "@tanstack/react-query"; import { AppShell, Button, Combobox, - Modal, SegmentedControl, Tabs, type AppShellProps, @@ -22,18 +21,14 @@ import { UiListTree, UiPlay, UiRefresh, - UiSave, UiTerminal, UiTrash, } from "@flanksource/clicky-ui/data"; -import "@flanksource/clicky-ui/mdx-editor.css"; -import { MdxEditorField } from "@flanksource/clicky-ui/mdx-editor"; import { PromptRunEditor, - buildAISpecRuntimePayload, familiesFromRuntimeCatalog, + type AIPromptRunValue, type AISpecRuntimePermissionCatalog, - type AISpecRuntimeValue, type RuntimeCatalogFamily, type ToolMeta, } from "@flanksource/clicky-ui/ai"; @@ -43,11 +38,10 @@ import { type ResolvedOperation, } from "@flanksource/clicky-ui/rpc"; import { apiClient } from "./api"; +import { isReadOnlyDbContext } from "./dbContext"; import { PromptRunStream } from "./PromptRunStream"; import { PromptBatchInspector } from "./PromptBatchInspector"; import { PromptSchemaEditor } from "./PromptSchemaEditor"; -import { PromptRuntimeRows } from "./PromptRuntimeRows"; -import { validateRuntimeRows } from "./promptRuntimeRowsHelpers"; import { RunningPromptsBadge, RunningPromptsRunsTab } from "./RunningPrompts"; import { isPromptBatchHandle, @@ -60,30 +54,28 @@ import { requiredOperation, resolvePromptOps, unwrapResponse, + type PromptDetail, type PromptSourceFilter, type PromptSummary, } from "./promptData"; import type { PromptSchemaKind } from "./promptSchemaSource"; import { - normalizeRuntimeModel, + mergePromptModelCatalogs, promptOptions, - runtimeModelsPayload, - runtimeRowsFromPrompt, } from "./promptWorkbenchHelpers"; +import { + PromptSourceMarkdownEditor, + PromptWriteAction, + PromptWriteModal, + type PromptWriteInput, + type PromptWriteMode, +} from "./PromptWriteModal"; type Navigate = (to: string, opts?: { replace?: boolean }) => void; type SourceFilter = PromptSourceFilter; type DetailTab = "source" | "runner" | "schema" | "runs"; -type PromptDetail = PromptSummary & { - content: string; - inputSchema?: Record; - inputDefault?: Record; - outputSchema?: Record; - metadata?: Record; -}; - type PromptPreviewResult = { id: string; name: string; @@ -113,7 +105,11 @@ const SOURCE_OPTIONS = [ { id: "local", label: "Local" }, ] satisfies Array<{ id: SourceFilter; label: string }>; -const EMPTY_RUNTIME: AISpecRuntimeValue = { budget: { timeout: "2h" } }; +const EMPTY_RUN_REQUEST: AIPromptRunValue = { + variables: {}, + spec: { budget: { timeout: "2h" } }, + chat: true, +}; const EMPTY_PROMPTS: PromptSummary[] = []; const EMPTY_MODELS: ChatModel[] = []; const SCRATCH_PROMPT_ID = "__scratch__"; @@ -128,6 +124,7 @@ const SCRATCH_PROMPT: PromptDetail = { relPath: "scratch.prompt", writable: false, content: "", + run: EMPTY_RUN_REQUEST, }; const AGENT_TOOLS = [ @@ -220,11 +217,9 @@ const AGENT_TOOLS = [ type PromptDetailState = { detailId?: string; draft: string; - variables: Record; + runRequest: AIPromptRunValue; variablesValid: boolean; schemaValidity: Record; - runtime: AISpecRuntimeValue; - additionalRuntimes: AISpecRuntimeValue[]; previewResult?: PromptPreviewResult; activeRunID?: string; activeBatch?: PromptBatchHandle; @@ -234,7 +229,7 @@ type PromptDetailState = { type PromptDetailStateAction = | { type: "draft"; detail?: PromptDetail; value: string } - | { type: "variables"; detail?: PromptDetail; value: Record } + | { type: "run-request"; detail?: PromptDetail; value: AIPromptRunValue } | { type: "variables-validity"; detail?: PromptDetail; value: boolean } | { type: "schema-validity"; @@ -242,12 +237,6 @@ type PromptDetailStateAction = kind: PromptSchemaKind; value: boolean; } - | { type: "runtime"; detail?: PromptDetail; value: AISpecRuntimeValue } - | { - type: "runtime-rows"; - detail?: PromptDetail; - value: AISpecRuntimeValue[]; - } | { type: "preview-result"; detail?: PromptDetail; @@ -271,8 +260,8 @@ function promptDetailReducer( switch (action.type) { case "draft": return { ...current, draft: action.value }; - case "variables": - return { ...current, variables: action.value }; + case "run-request": + return { ...current, runRequest: action.value }; case "variables-validity": return { ...current, variablesValid: action.value }; case "schema-validity": @@ -283,14 +272,6 @@ function promptDetailReducer( [action.kind]: action.value, }, }; - case "runtime": - return { ...current, runtime: action.value }; - case "runtime-rows": - return { - ...current, - runtime: action.value[0] ?? {}, - additionalRuntimes: action.value.slice(1), - }; case "preview-result": return { ...current, previewResult: action.value }; case "active-run": @@ -307,15 +288,16 @@ function promptDetailReducer( } function initialPromptDetailState(detail?: PromptDetail): PromptDetailState { - const runtimeRows = detail ? runtimeRowsFromPrompt(detail) : []; + const runRequest = detail?.run ?? EMPTY_RUN_REQUEST; return { detailId: detail?.id, draft: detail?.content ?? "", - variables: detail?.inputDefault ?? {}, + runRequest: { + ...runRequest, + spec: { ...EMPTY_RUN_REQUEST.spec, ...runRequest.spec }, + }, variablesValid: true, schemaValidity: { input: true, output: true }, - runtime: { ...EMPTY_RUNTIME, ...runtimeRows[0] }, - additionalRuntimes: runtimeRows.slice(1), previewResult: undefined, activeRunID: undefined, activeBatch: undefined, @@ -359,7 +341,7 @@ function usePromptWorkbenchView({ undefined, () => initialPromptDetailState(), ); - const [createOpen, setCreateOpen] = useState(false); + const [writeMode, setWriteMode] = useState(); const listQuery = useQuery({ queryKey: [ @@ -380,13 +362,34 @@ function usePromptWorkbenchView({ queryKey: ["prompt-schema"], queryFn: fetchPromptSchema, }); + const modelCatalogQuery = useQuery({ + queryKey: ["chat-model-catalog"], + queryFn: () => fetchCatalog("/api/chat/models", "Model catalog"), + }); + const runtimeCatalogQuery = useQuery({ + queryKey: ["chat-runtime-catalog"], + queryFn: () => + fetchCatalog( + "/api/chat/runtimes", + "Runtime catalog", + ), + }); const permissionCatalogQuery = useQuery({ queryKey: ["permission-catalog"], queryFn: () => fetchPermissionCatalog(), }); const prompts = listQuery.data ?? EMPTY_PROMPTS; - const models = promptSchemaQuery.data?.models ?? EMPTY_MODELS; + const models = useMemo( + () => + mergePromptModelCatalogs( + promptSchemaQuery.data?.models ?? EMPTY_MODELS, + modelCatalogQuery.data ?? EMPTY_MODELS, + ), + [modelCatalogQuery.data, promptSchemaQuery.data?.models], + ); + const runtimeCatalog = + runtimeCatalogQuery.data ?? promptSchemaQuery.data?.runtimes; const activePromptId = selectedId; const selectedSummary = useMemo( @@ -415,10 +418,6 @@ function usePromptWorkbenchView({ const detail = activePromptId ? detailQuery.data : SCRATCH_PROMPT; const selected = detail ?? selectedSummary; const selectedDetailState = promptDetailStateFor(detailState, detail); - const runtimeRows = [ - selectedDetailState.runtime, - ...selectedDetailState.additionalRuntimes, - ]; const scratch = isScratchPrompt(detail); const writableSources = useMemo( () => uniqueWritableSources(prompts), @@ -427,10 +426,19 @@ function usePromptWorkbenchView({ const canSave = Boolean( detail && !scratch && + detail.writable && promptOps.update && selectedDetailState.schemaValidity.input && selectedDetailState.schemaValidity.output && - (detail.writable ? selectedDetailState.draft !== detail.content : true), + selectedDetailState.draft !== detail.content, + ); + const canSaveAs = Boolean( + detail && + !scratch && + !detail.writable && + promptOps.create && + selectedDetailState.schemaValidity.input && + selectedDetailState.schemaValidity.output, ); const hasSelection = Boolean(detail || activePromptId); const operationsReady = Boolean( @@ -443,7 +451,7 @@ function usePromptWorkbenchView({ } async function saveDraft() { - if (!detail || scratch || !promptOps.update) return; + if (!detail?.writable || scratch || !promptOps.update) return; dispatchDetailState({ type: "action-error", detail, value: undefined }); dispatchDetailState({ type: "action-loading", detail, value: "save" }); try { @@ -453,14 +461,8 @@ function usePromptWorkbenchView({ { content: selectedDetailState.draft }, ); await listQuery.refetch(); - if (saved.id === detail.id) { - dispatchDetailState({ type: "saved", detail, content: saved.content }); - await detailQuery.refetch(); - } else { - // Saving a read-only (embedded) prompt forks it to a local copy; - // switch to the new writable prompt. - onNavigate(`/prompts/${encodeURIComponent(saved.id)}`); - } + dispatchDetailState({ type: "saved", detail, content: saved.content }); + await detailQuery.refetch(); } catch (error) { dispatchDetailState({ type: "action-error", @@ -472,6 +474,18 @@ function usePromptWorkbenchView({ } } + async function createPromptCopy(input: PromptWriteInput) { + const create = requiredOperation(promptOps.create, "create"); + const created = await submitPromptOperation( + create, + {}, + { ...input }, + ); + setWriteMode(undefined); + await listQuery.refetch(); + onNavigate(`/prompts/${encodeURIComponent(created.id)}`); + } + async function previewPrompt() { if (!detail || !promptOps.preview) return; dispatchDetailState({ type: "action-error", detail, value: undefined }); @@ -480,10 +494,7 @@ function usePromptWorkbenchView({ const preview = await submitPromptOperation( promptOps.preview, promptActionParams(detail), - { - variables: selectedDetailState.variables, - ...runtimePayload(selectedDetailState.runtime, models), - }, + selectedDetailState.runRequest, ); dispatchDetailState({ type: "preview-result", detail, value: preview }); dispatchDetailState({ type: "active-run", detail, value: undefined }); @@ -507,14 +518,7 @@ function usePromptWorkbenchView({ const handle = await submitPromptOperation( promptOps.run, promptActionParams(detail), - { - variables: selectedDetailState.variables, - ...runtimePayload(selectedDetailState.runtime, models), - ...(runtimeRows.length > 1 - ? { runtimes: runtimeModelsPayload(runtimeRows, models) } - : {}), - chat: promptChatEligible(detail, selectedDetailState.runtime), - }, + selectedDetailState.runRequest, ); dispatchDetailState({ type: "preview-result", detail, value: undefined }); if (isPromptBatchHandle(handle)) { @@ -560,169 +564,175 @@ function usePromptWorkbenchView({ } return ( - Captain
} - navSections={navSections} - collapsedStorageKey={CAPTAIN_SIDEBAR_COLLAPSE_KEY} - actions={actions} - search={search} - bodySidebar={ - onNavigate(`/prompts/${encodeURIComponent(id)}`)} - onRefresh={() => void refreshAll()} - onCreate={() => setCreateOpen(true)} - /> - } - bodyHeader={ - - } - bodyActions={ -
- { - dispatchDetailState({ - type: "active-batch", - detail, - value: undefined, - }); - dispatchDetailState({ type: "active-run", detail, value: id }); - setTab("runner"); - }} + <> + Captain
} + navSections={navSections} + collapsedStorageKey={CAPTAIN_SIDEBAR_COLLAPSE_KEY} + actions={actions} + search={search} + bodySidebar={ + onNavigate(`/prompts/${encodeURIComponent(id)}`)} + onRefresh={() => void refreshAll()} + onCreate={() => setWriteMode("create")} /> - {detail?.writable && !scratch && promptOps.delete && ( - - )} - {detail && !scratch && promptOps.update && ( + } + bodyHeader={ + + } + bodyActions={ +
+ { + dispatchDetailState({ + type: "active-batch", + detail, + value: undefined, + }); + dispatchDetailState({ type: "active-run", detail, value: id }); + setTab("runner"); + }} + /> + {detail?.writable && !scratch && promptOps.delete && ( + + )} + {detail && + !scratch && + ((detail.writable && promptOps.update) || + (!detail.writable && promptOps.create)) && ( + { + if (detail.writable) { + void saveDraft(); + } else { + setWriteMode("save-as"); + } + }} + /> + )} - )} - -
- } - bodySplit={30} - contentClassName="p-0 overflow-hidden" - > - setTab(next as DetailTab)} - draft={selectedDetailState.draft} - onDraftChange={(value) => - dispatchDetailState({ type: "draft", detail, value }) - } - onSchemaValidityChange={(kind, value) => - dispatchDetailState({ - type: "schema-validity", - detail, - kind, - value, - }) - } - variables={selectedDetailState.variables} - variablesValid={selectedDetailState.variablesValid} - onVariablesChange={(value) => - dispatchDetailState({ type: "variables", detail, value }) - } - onVariablesValidityChange={(value) => - dispatchDetailState({ type: "variables-validity", detail, value }) - } - runtime={selectedDetailState.runtime} - onRuntimeChange={(value) => - dispatchDetailState({ type: "runtime", detail, value }) - } - runtimeRows={runtimeRows} - onRuntimeRowsChange={(value) => - dispatchDetailState({ type: "runtime-rows", detail, value }) - } - models={models} - promptSchema={promptSchemaQuery.data} - tools={AGENT_TOOLS} - permissionCatalog={permissionCatalogQuery.data} - previewResult={selectedDetailState.previewResult} - activeRunID={selectedDetailState.activeRunID} - activeBatch={selectedDetailState.activeBatch} - onEditBatch={() => - dispatchDetailState({ - type: "active-batch", - detail, - value: undefined, - }) + } - onSelectRun={(id) => { - dispatchDetailState({ - type: "active-batch", - detail, - value: undefined, - }); - dispatchDetailState({ type: "active-run", detail, value: id }); - if (id) setTab("runner"); - }} - onPreview={() => void previewPrompt()} - onRun={() => void runPrompt()} - previewLoading={selectedDetailState.actionLoading === "preview"} - runLoading={selectedDetailState.actionLoading === "run"} - previewEnabled={Boolean(promptOps.preview && detail)} - runEnabled={Boolean(promptOps.run && detail)} - /> - setCreateOpen(false)} + bodySplit={30} + contentClassName="p-0 overflow-hidden" + > + setTab(next as DetailTab)} + draft={selectedDetailState.draft} + onDraftChange={(value) => + dispatchDetailState({ type: "draft", detail, value }) + } + onSchemaValidityChange={(kind, value) => + dispatchDetailState({ + type: "schema-validity", + detail, + kind, + value, + }) + } + variablesValid={selectedDetailState.variablesValid} + onVariablesValidityChange={(value) => + dispatchDetailState({ type: "variables-validity", detail, value }) + } + runRequest={selectedDetailState.runRequest} + onRunRequestChange={(value) => + dispatchDetailState({ type: "run-request", detail, value }) + } + models={models} + runtimeCatalog={runtimeCatalog} + promptSchema={promptSchemaQuery.data} + tools={AGENT_TOOLS} + permissionCatalog={permissionCatalogQuery.data} + previewResult={selectedDetailState.previewResult} + activeRunID={selectedDetailState.activeRunID} + activeBatch={selectedDetailState.activeBatch} + onEditBatch={() => + dispatchDetailState({ + type: "active-batch", + detail, + value: undefined, + }) + } + onSelectRun={(id) => { + dispatchDetailState({ + type: "active-batch", + detail, + value: undefined, + }); + dispatchDetailState({ type: "active-run", detail, value: id }); + if (id) setTab("runner"); + }} + onPreview={() => void previewPrompt()} + onRun={() => void runPrompt()} + previewLoading={selectedDetailState.actionLoading === "preview"} + runLoading={selectedDetailState.actionLoading === "run"} + previewEnabled={Boolean(promptOps.preview && detail)} + // Running a prompt writes its session; a read-only database context + // would have the POST rejected, so the control is disabled instead. + runEnabled={Boolean(promptOps.run && detail) && !isReadOnlyDbContext()} + /> + + setWriteMode(undefined)} sources={writableSources} - createOp={promptOps.create} - seedContent={scratch ? undefined : detail?.content} - onCreated={(prompt) => { - setCreateOpen(false); - void listQuery.refetch(); - onNavigate(`/prompts/${encodeURIComponent(prompt.id)}`); - }} + onSubmit={createPromptCopy} + {...(writeMode === "save-as" && detail + ? { + initialName: detail.name, + initialContent: selectedDetailState.draft, + } + : !scratch && detail + ? { initialContent: detail.content } + : {})} /> - - ); -} - -function promptChatEligible(detail: PromptDetail, runtime: AISpecRuntimeValue) { - return ( - !detail.outputSchema && - !runtime.workflow?.verify && - !runtime.workflow?.commits?.length + ); } @@ -913,15 +923,12 @@ function PromptDetailPane({ draft, onDraftChange, onSchemaValidityChange, - variables, variablesValid, - onVariablesChange, onVariablesValidityChange, - runtime, - onRuntimeChange, - runtimeRows, - onRuntimeRowsChange, + runRequest, + onRunRequestChange, models, + runtimeCatalog, promptSchema, tools, permissionCatalog, @@ -946,15 +953,12 @@ function PromptDetailPane({ draft: string; onDraftChange: (value: string) => void; onSchemaValidityChange: (kind: PromptSchemaKind, valid: boolean) => void; - variables: Record; variablesValid: boolean; - onVariablesChange: (value: Record) => void; onVariablesValidityChange: (valid: boolean) => void; - runtime: AISpecRuntimeValue; - onRuntimeChange: (value: AISpecRuntimeValue) => void; - runtimeRows: AISpecRuntimeValue[]; - onRuntimeRowsChange: (value: AISpecRuntimeValue[]) => void; + runRequest: AIPromptRunValue; + onRunRequestChange: (value: AIPromptRunValue) => void; models: ChatModel[]; + runtimeCatalog?: RuntimeCatalogFamily[]; promptSchema?: PromptSchemaDoc; tools: ToolMeta[]; permissionCatalog?: AISpecRuntimePermissionCatalog; @@ -1010,16 +1014,13 @@ function PromptDetailPane({ ? undefined : normalizeObjectSchema(detail.inputSchema); const backendCliArgs = promptSchema?.backends?.find( - (backend) => backend.backend === runtime.backend, + (backend) => backend.backend === runRequest.spec?.backend, )?.args; - // The picker's families come from the same document as its models, so a - // backend the user disabled is absent from both. - const runtimeFamilies = familiesFromRuntimeCatalog(promptSchema?.runtimes); + const runtimeFamilies = familiesFromRuntimeCatalog(runtimeCatalog); const promptReady = !scratch || - Boolean(runtime.prompt?.user?.trim()) || - Boolean(runtime.prompt?.attachments?.length); - const runtimeRowsError = validateRuntimeRows(runtimeRows); + Boolean(runRequest.spec?.prompt?.user?.trim()) || + Boolean(runRequest.spec?.prompt?.attachments?.length); return (
@@ -1058,23 +1059,12 @@ function PromptDetailPane({
- } + value={runRequest} + onChange={onRunRequestChange} families={runtimeFamilies} - models={promptSelectableModels(models)} + models={models} tools={tools} secretSelector={CAPTAIN_SECRET_SELECTOR} - variables={variables} - onVariablesChange={onVariablesChange} onVariablesValidityChange={onVariablesValidityChange} enableAttachments {...(permissionCatalog ? { permissionCatalog } : {})} @@ -1109,10 +1099,7 @@ function PromptDetailPane({ size="sm" loading={runLoading} disabled={ - !runEnabled || - !promptReady || - Boolean(runtimeRowsError) || - (!schema && !variablesValid) + !runEnabled || !promptReady || (!schema && !variablesValid) } onClick={onRun} > @@ -1146,8 +1133,8 @@ function SourceEditor({
{!detail.writable && (
- This is an embedded prompt. Saving your edits creates a local, - editable copy. + This is an embedded prompt. Use Save as… to create a local, editable + copy.
)} void; - readOnly?: boolean; - minHeight: string | number; -}) { - return ( -
-
- {label} -
-
- -
-
- ); -} - -function Field({ label, children }: { label: string; children: ReactNode }) { - return ( - - ); -} - function RunnerOutput({ previewResult, activeRunID, @@ -1250,151 +1180,6 @@ function RunnerOutput({ ); } -function CreatePromptModal({ - open, - ...props -}: { - open: boolean; - onClose: () => void; - sources: Array<{ id: string; label: string }>; - createOp?: ResolvedOperation; - seedContent?: string; - onCreated: (prompt: PromptDetail) => void; -}) { - if (!open) return null; - return ; -} - -function CreatePromptModalForm({ - onClose, - sources, - createOp, - seedContent, - onCreated, -}: { - onClose: () => void; - sources: Array<{ id: string; label: string }>; - createOp?: ResolvedOperation; - seedContent?: string; - onCreated: (prompt: PromptDetail) => void; -}) { - const [name, setName] = useState(""); - const [relPath, setRelPath] = useState(""); - const [target, setTarget] = useState(() => sources[0]?.id ?? ""); - const [content, setContent] = useState( - () => seedContent || defaultPromptContent(""), - ); - const [loading, setLoading] = useState(false); - const [error, setError] = useState(); - - async function submit() { - if (!createOp) return; - setLoading(true); - setError(undefined); - try { - const created = await submitPromptOperation( - createOp, - {}, - { - target, - name, - relPath, - content, - }, - ); - onCreated(created); - } catch (err) { - setError(errorMessage(err)); - } finally { - setLoading(false); - } - } - - return ( - - - -
- } - > -
- {error &&
{error}
} -
- - { - const next = event.target.value; - setName(next); - if (!relPath) setContent(defaultPromptContent(next)); - }} - className="h-control-h w-full rounded-md border border-border bg-background px-density-3 text-sm outline-none focus:ring-2 focus:ring-ring" - /> - - - setRelPath(event.target.value)} - className="h-control-h w-full rounded-md border border-border bg-background px-density-3 text-sm outline-none focus:ring-2 focus:ring-ring" - placeholder="name.prompt" - /> - - - - -
- -
- - ); -} - -function createPromptModalKey({ - seedContent, - sources, -}: { - seedContent?: string; - sources: Array<{ id: string; label: string }>; -}) { - return `${sources[0]?.id ?? ""}:${seedContent ?? ""}`; -} - async function fetchPermissionCatalog() { const response = await fetch("/api/captain/ai/permissions/catalog", { headers: { Accept: "application/json" }, @@ -1448,6 +1233,19 @@ async function fetchPromptSchema(): Promise { return (await response.json()) as PromptSchemaDoc; } +async function fetchCatalog(endpoint: string, label: string): Promise { + const response = await fetch(endpoint, { + headers: { Accept: "application/json" }, + }); + if (!response.ok) { + const message = await response.text(); + throw new Error(message || `${label} failed with ${response.status}`); + } + const data: unknown = await response.json(); + if (!Array.isArray(data)) throw new Error(`${label} must be an array`); + return data as T[]; +} + const CAPTAIN_SECRET_SELECTOR = { loadResources: fetchSecretResources, loadKeyPreview: fetchSecretKeyPreview, @@ -1579,67 +1377,6 @@ function normalizeObjectSchema( } as JsonSchemaObject; } -// runtime is the single source of truth: the inline PromptRunEditor and its -// "Edit spec" modal both edit this one AISpecRuntimeValue, so the payload is -// just the compacted spec (plus catalog model/backend normalization). -function runtimePayload(runtime: AISpecRuntimeValue, models: ChatModel[]) { - return normalizeSpecRuntimePayload( - buildAISpecRuntimePayload(runtime), - models, - runtime.backend, - ); -} - -function normalizeSpecRuntimePayload( - payload: Record, - models: ChatModel[], - backend?: string, -) { - const spec = payload.spec; - if (!spec || typeof spec !== "object" || Array.isArray(spec)) return payload; - const specRecord = { ...(spec as Record) }; - if (typeof specRecord.model === "string") { - const selected = normalizeRuntimeModel( - specRecord.model, - models, - typeof specRecord.backend === "string" ? specRecord.backend : backend, - ); - if (selected.model && selected.model !== specRecord.model) { - if (typeof specRecord.id !== "string" || !specRecord.id.trim()) { - specRecord.id = specRecord.model; - } - specRecord.model = selected.model; - } - if ( - selected.backend && - (typeof specRecord.backend !== "string" || !specRecord.backend.trim()) - ) { - specRecord.backend = selected.backend; - } - } - return { ...payload, spec: specRecord }; -} - -function promptSelectableModels(models: ChatModel[]) { - return models.map((model) => - model.configured === false ? { ...model, configured: true } : model, - ); -} - -function defaultPromptContent(name: string) { - const promptName = name.trim() || "new prompt"; - return `--- -name: ${JSON.stringify(promptName)} -description: "" -input: - schema: - input: string ---- -{{role "user"}} -{{input}} -`; -} - function errorMessage(error: unknown) { if (error instanceof Error) return error.message; if (typeof error === "string") return error; diff --git a/pkg/cli/webapp/src/PromptWriteModal.test.tsx b/pkg/cli/webapp/src/PromptWriteModal.test.tsx new file mode 100644 index 00000000..29a691b9 --- /dev/null +++ b/pkg/cli/webapp/src/PromptWriteModal.test.tsx @@ -0,0 +1,105 @@ +import { cleanup, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { PromptWriteAction, PromptWriteModal } from "./PromptWriteModal"; + +vi.mock("@flanksource/clicky-ui/mdx-editor", () => ({ + MdxEditorField: ({ + value, + onChange, + }: { + value: string; + onChange?: (value: string) => void; + }) => ( +