diff --git a/cmd/captain/main.go b/cmd/captain/main.go index da48e814..39966e0f 100644 --- a/cmd/captain/main.go +++ b/cmd/captain/main.go @@ -125,7 +125,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, or --backend to inspect a single adapter." configureCmd := clicky.AddNamedCommandWithContext("configure", rootCmd, cli.ConfigureOptions{}, cli.RunConfigure) configureCmd.Use = "configure [provider]" diff --git a/go.mod b/go.mod index 01817de1..7e0deee1 100644 --- a/go.mod +++ b/go.mod @@ -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/migrations/02_merge_duplicate_sessions.sql b/migrations/02_merge_duplicate_sessions.sql index 577c7366..505dacd1 100644 --- a/migrations/02_merge_duplicate_sessions.sql +++ b/migrations/02_merge_duplicate_sessions.sql @@ -17,9 +17,9 @@ -- duplicate groups no group had more than one member with messages, so this -- picks a winner rather than merging two transcripts; ties fall back to the -- oldest row and then to the id, so the choice is deterministic on re-run. --- Prompt runs and child sessions pointing at a ghost are re-pointed, the ghosts --- are deleted, and the winner then absorbs a ghost's provider label -- that label --- exists nowhere else once the ghost is gone. +-- Prompt runs, child sessions, and processes pointing at a ghost are re-pointed, +-- the ghosts are deleted, and the winner then absorbs a ghost's provider label -- +-- that label exists nowhere else once the ghost is gone. -- -- Every reference re-pointed below is one the ghost's ON DELETE CASCADE would -- otherwise destroy. The ghost's own transcript rows -- messages, turns, events, @@ -101,6 +101,13 @@ BEGIN FROM captain_duplicate_session_map m WHERE w.id = m.winner AND w.root_session_id = m.loser; + -- A running process is the durable identity of the execution that created the + -- ghost. Its session foreign key also cascades, so preserve it on the row that + -- actually owns the transcript before removing the duplicate. + UPDATE captain_session_processes p SET session_id = m.winner + FROM captain_duplicate_session_map m + WHERE p.session_id = m.loser; + DELETE FROM captain_sessions s USING captain_duplicate_session_map m WHERE s.id = m.loser; 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..1928bad8 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 @@ -314,432 +313,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..9abf6006 --- /dev/null +++ b/migrations/31_execution_events.pg.hcl @@ -0,0 +1,281 @@ +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" + } +} 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..d83545fa 100644 --- a/migrations/60_view_session_overview.sql +++ b/migrations/60_view_session_overview.sql @@ -100,7 +100,8 @@ 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 FROM public.captain_sessions s LEFT JOIN LATERAL ( SELECT p.* @@ -205,7 +206,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/merge_duplicate_sessions_integration_test.go b/migrations/merge_duplicate_sessions_integration_test.go index aa87f31c..f5e8e6de 100644 --- a/migrations/merge_duplicate_sessions_integration_test.go +++ b/migrations/merge_duplicate_sessions_integration_test.go @@ -9,12 +9,11 @@ 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) { + It("re-points dependent rows at the survivor instead of cascading them away", func(ctx SpecContext) { handle := dbtest.ForGinkgo(dbtest.Options{Name: "captain_merge_duplicates"}) dsn, db := handle.DSN(), handle.SQL() @@ -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_test.go b/migrations/migrations_test.go index cee76efd..7d03e3f6 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", @@ -82,8 +84,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 +153,8 @@ 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_") for _, name := range expectedFiles { assertContainsNone(t, name, `table "captain_outbox"`, diff --git a/pkg/ai/adapters.go b/pkg/ai/adapters.go index b14251e4..daeb3866 100644 --- a/pkg/ai/adapters.go +++ b/pkg/ai/adapters.go @@ -17,9 +17,10 @@ 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"` } // AdapterStatus is the resolved auth/availability of a single agent adapter 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_disabled_ginkgo_test.go b/pkg/ai/catalog_disabled_ginkgo_test.go index 71ad3fec..f8cdfe27 100644 --- a/pkg/ai/catalog_disabled_ginkgo_test.go +++ b/pkg/ai/catalog_disabled_ginkgo_test.go @@ -125,4 +125,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..e619d9e3 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,13 @@ 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"` // 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 @@ -67,10 +69,18 @@ func catalogInfoFrom(models []Model, configuredProviders []string) []ModelInfo { } else { configured = slices.Contains(configuredProviders, BackendToProvider(m.Backend)) } + 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, 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/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/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/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/options.go b/pkg/ai/provider/genkit/options.go index 26564af4..365d985e 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(captureGenkitModelRequest)), + } toolOptions, err := p.toolOptions(req.ToolPreferences, emit) if err != nil { return nil, err @@ -132,6 +136,18 @@ func generateOptions(p *Provider, req ai.Request, stream gkai.ModelStreamCallbac return opts, nil } +func captureGenkitModelRequest(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 + }, + }, 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..3857fbf6 100644 --- a/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go +++ b/pkg/ai/provider/genkit/tool_approval_ginkgo_test.go @@ -47,7 +47,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 +63,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 +176,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/tools.go b/pkg/ai/provider/genkit/tools.go index 4eb5f84d..258ba87c 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 { diff --git a/pkg/ai/provider/init.go b/pkg/ai/provider/init.go index b5b4298b..b6590848 100644 --- a/pkg/ai/provider/init.go +++ b/pkg/ai/provider/init.go @@ -24,7 +24,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_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..e9b9dc2a --- /dev/null +++ b/pkg/aichat/aimock_lifecycle_integration_test.go @@ -0,0 +1,500 @@ +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) 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: 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..c806a0dc --- /dev/null +++ b/pkg/aichat/approval_execution.go @@ -0,0 +1,90 @@ +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()) + } + } + thread, err := s.options.Threads.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) + for event := range s.persistedEvents(streamContext, request, execution.TurnID(), events) { + 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..45da9009 --- /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) + 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..970ad060 --- /dev/null +++ b/pkg/aichat/database_threads.go @@ -0,0 +1,376 @@ +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) { + record, err := s.db.CreateOrGetSession(ctx, database.CreateSessionInput{ + ID: uuid.New(), Source: "aichat", Provider: "captain", HostID: "local", + Title: strings.TrimSpace(title), Metadata: map[string]any{"aichat": true}, + }) + 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 + } + 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 || overview.Source != "aichat" { + return nil, fmt.Errorf("captain chat session %s has source %q", parsed, overview.Source) + } + 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 + } + aggregate := sessionFromOverview(overview) + 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) AddUsage(ctx context.Context, id string, _ TurnUsage) (*Thread, error) { + return s.Get(ctx, id) +} + +func sessionFromOverview(overview database.SessionOverview) *session.Session { + return &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), + }, + 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), + ProviderCostUSD: overview.CostUSD, + }, + } +} + +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), ProviderCostUSD: rows[i].CostUSD}, + } + } + 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..68a9e110 --- /dev/null +++ b/pkg/aichat/database_threads_integration_test.go @@ -0,0 +1,205 @@ +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/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: 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")) + }) +}) diff --git a/pkg/aichat/events.go b/pkg/aichat/events.go index 588b43f1..f7396877 100644 --- a/pkg/aichat/events.go +++ b/pkg/aichat/events.go @@ -17,6 +17,7 @@ type toolState struct { type eventStream struct { writer *SSEWriter + messageID string blockType string blockID string nextBlock int @@ -31,6 +32,7 @@ type eventStream struct { // EventStreamOptions carries request state needed to finish the resumed UI message. type EventStreamOptions struct { ToolApproval *api.ToolApprovalResume + MessageID string } // WriteEventStream translates a Captain event channel into one complete UI Message Stream. @@ -42,7 +44,7 @@ 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{}} if err := stream.start(); err != nil { return err } @@ -95,7 +97,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 +130,31 @@ 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.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.tools = map[string]toolState{} + s.terminal = true + return nil +} + func (s *eventStream) delta(kind, delta string) error { if delta == "" { return nil @@ -189,10 +211,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 +267,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 { diff --git a/pkg/aichat/execution.go b/pkg/aichat/execution.go new file mode 100644 index 00000000..e035abaa --- /dev/null +++ b/pkg/aichat/execution.go @@ -0,0 +1,170 @@ +package aichat + +import ( + "context" + + "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 { + sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error(), Model: event.Model}) + 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..d6e42f87 --- /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: 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: 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: 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: 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: 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: 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: 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: 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..82add5bf --- /dev/null +++ b/pkg/aichat/execution_database.go @@ -0,0 +1,494 @@ +package aichat + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "time" + + "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 + 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{} +} + +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.db.FinishChatModelCall(ctx, database.FinishChatModelCallInput{ + ID: e.modelCallID, Status: database.ModelCallStatusSucceeded, + StopReason: "tool_approval", Event: 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.db.FinishChatModelCall(ctx, database.FinishChatModelCallInput{ + ID: e.modelCallID, Status: database.ModelCallStatusCancelled, StopReason: "interrupt", + Event: 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.db.FinishChatModelCall(ctx, database.FinishChatModelCallInput{ + ID: e.modelCallID, Status: callStatus, StopReason: stopReason, Event: 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..86fc7f62 --- /dev/null +++ b/pkg/aichat/execution_database_authority.go @@ -0,0 +1,300 @@ +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) + } + turn, created, err := a.db.CreateChatTurn(ctx, database.CreateChatTurnInput{ + SessionID: session.ID, ProviderTurnID: request.RequestID, + }) + if err != nil { + return nil, err + } + if !created { + return nil, fmt.Errorf("chat turn %q already exists in state %s", request.RequestID, turn.Status) + } + renderedSpec, err := renderedSpecMap(request.Spec) + if err != nil { + return nil, err + } + run, err := a.db.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 err != nil { + return nil, err + } + if run.State != database.PromptRunStatePending { + return nil, fmt.Errorf("chat request %q already has prompt run %s in state %s", request.RequestID, run.ID, run.State) + } + modelCallID, err := a.db.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 err != nil { + return nil, err + } + execution := &databaseExecution{ + db: a.db, ctx: ctx, session: session, turn: turn, run: run, modelCallID: modelCallID, + events: make(chan api.Event, 16), definitions: append([]api.ToolDefinition(nil), request.Definitions...), + approvalIDs: map[string]uuid.UUID{}, providerToolUseReady: make(chan struct{}, 1), + } + if err := execution.markRunning(ctx); err != nil { + return nil, err + } + 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, + 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..6cb57c95 --- /dev/null +++ b/pkg/aichat/execution_database_integration_test.go @@ -0,0 +1,271 @@ +package aichat_test + +import ( + "context" + "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("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..f30f1227 --- /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) + 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..68857aa8 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 { @@ -175,7 +218,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/persistence.go b/pkg/aichat/persistence.go index 676741f9..4e81e4cd 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,7 @@ func newAssistantMessageBuilder(options assistantMessageBuilderOptions) (*assist return builder, nil } -func (s *Service) persistedEvents(ctx context.Context, request ChatRequest, source <-chan api.Event) <-chan api.Event { +func (s *Service) persistedEvents(ctx context.Context, request ChatRequest, turnID string, source <-chan api.Event) <-chan api.Event { if request.ThreadID == "" { return source } @@ -88,8 +87,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()}) @@ -107,7 +108,7 @@ func (s *Service) persistedEvents(ctx context.Context, request ChatRequest, sour return } } - if !persisted && (event.Kind == api.EventResult || event.Kind == api.EventError) { + if !persisted && (event.Kind == api.EventResult || event.Kind == api.EventError || event.Kind == api.EventInterrupted) { if err := s.persistCompletedTurn(ctx, request.ThreadID, builder, event); err != nil { sendEvent(ctx, out, api.Event{Kind: api.EventError, Error: err.Error(), Model: event.Model}) return @@ -203,11 +204,26 @@ func (b *assistantMessageBuilder) apply(event api.Event) error { return err } b.message.Parts = append(b.message.Parts, UIPart{Type: "data-error", Data: payload}) + case api.EventInterrupted: + return b.interrupted() case api.EventSystem: } return nil } +func (b *assistantMessageBuilder) interrupted() error { + 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) appendText(partType, text string) { if text == "" { return @@ -256,12 +272,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 +321,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 @@ -351,12 +369,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 +402,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..f4147f52 100644 --- a/pkg/aichat/provider_config.go +++ b/pkg/aichat/provider_config.go @@ -44,23 +44,23 @@ func (s *Service) annotateConfiguredModels(ctx context.Context, models ModelCata 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/service.go b/pkg/aichat/service.go index fd1ecab0..387b8b2d 100644 --- a/pkg/aichat/service.go +++ b/pkg/aichat/service.go @@ -6,6 +6,8 @@ import ( "fmt" "net/http" "strings" + "sync" + "time" aitools "github.com/flanksource/captain/pkg/ai/tools" "github.com/flanksource/captain/pkg/api" @@ -55,12 +57,15 @@ type ServiceOptions struct { MCP ToolProvider Attachments AttachmentResolver Threads ThreadStore + 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,7 +73,7 @@ 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 { @@ -112,7 +117,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 +131,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,7 +158,8 @@ 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 } @@ -155,8 +167,46 @@ func (s *Service) handleChat(w http.ResponseWriter, request *http.Request) { 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 +216,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 +234,90 @@ 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{ + if err := WriteEventStream(writer, s.persistedEvents(streamContext, chat, turnID, events), EventStreamOptions{ ToolApproval: chat.ToolApproval, + MessageID: assistantMessageID(chat, turnID), }); 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,25 +325,36 @@ 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") + return nil, fmt.Errorf("thread persistence is not configured") } thread, err := s.options.Threads.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] diff --git a/pkg/aichat/service_ginkgo_test.go b/pkg/aichat/service_ginkgo_test.go index 299ad6f5..72d4a505 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" @@ -53,9 +52,12 @@ func (f *fakeResolver) Provider(_ context.Context, config api.Config) (api.Strea } 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,27 +77,22 @@ 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() { @@ -219,6 +216,7 @@ 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, + Runtime: api.Model{Name: "test-model", Backend: api.BackendOpenAI}, }}} service := aichat.NewService(aichat.ServiceOptions{ Resolver: resolver, @@ -240,7 +238,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,"contextWindow":0,"inputMediaTypes":null}]`)) tools := httptest.NewRecorder() service.Handler().ServeHTTP(tools, httptest.NewRequest(http.MethodGet, "/api/chat/tools", nil)) @@ -309,48 +307,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: 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()}) 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)) }) @@ -368,8 +424,8 @@ var _ = Describe("Captain aichat service", func() { service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}, Threads: store}) response := httptest.NewRecorder() service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ - ID: "chat-1", ThreadID: thread.ID, Model: "openai/test-model", - Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "inspect"}}}}, + 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 +433,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")) @@ -421,8 +478,8 @@ var _ = Describe("Captain aichat service", func() { service := aichat.NewService(aichat.ServiceOptions{Resolver: &fakeResolver{provider: provider}, Threads: store}) response := httptest.NewRecorder() service.Handler().ServeHTTP(response, requestJSON(http.MethodPost, "/api/chat", aichat.ChatRequest{ - ThreadID: thread.ID, Model: "openai/test-model", - Messages: []aichat.UIMessage{{Role: "user", Parts: []aichat.UIPart{{Type: "text", Text: "inspect"}}}}, + 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/sse.go b/pkg/aichat/sse.go index 5f9c29f8..b9b41461 100644 --- a/pkg/aichat/sse.go +++ b/pkg/aichat/sse.go @@ -44,6 +44,7 @@ type MessageMetadata struct { Cost float64 `json:"cost,omitempty"` ContextTokens int `json:"contextTokens,omitempty"` Success *bool `json:"success,omitempty"` + Interrupted bool `json:"interrupted,omitempty"` } // 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..adb67c3b 100644 --- a/pkg/aichat/stream_ginkgo_test.go +++ b/pkg/aichat/stream_ginkgo_test.go @@ -81,6 +81,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 +113,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 +136,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"}`})) @@ -162,10 +168,26 @@ var _ = Describe("AI SDK v6 event stream", func() { 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.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", + "data-result", "finish-step", "finish", + })) + Expect(parts[5]["data"]).To(Equal(map[string]any{"success": false, "interrupted": true})) + Expect(parts[7]["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 +201,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 +244,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 +266,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 +279,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/threads.go b/pkg/aichat/threads.go index 6830618b..c0f4be0b 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 { @@ -48,9 +51,12 @@ type ThreadStore interface { 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 } @@ -61,9 +67,8 @@ func NewMemoryThreadStore() ThreadStore { 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 return cloneThread(thread), nil } @@ -133,6 +138,20 @@ 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 @@ -180,5 +199,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..95e4649b 100644 --- a/pkg/aichat/threads_http.go +++ b/pkg/aichat/threads_http.go @@ -8,10 +8,12 @@ import ( ) 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("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 { @@ -69,6 +71,17 @@ func (s *Service) handleGetThread(w http.ResponseWriter, request *http.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) 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..31881bbd 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")), ) }) @@ -75,6 +92,7 @@ var _ = Describe("AI SDK v6 wire types", func() { ID: "openai/gpt", Provider: "openai", Label: "GPT", Reasoning: true, Temperature: true, Configured: true, 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,"contextWindow":128000,"inputMediaTypes":["image/*"]}]`)) toolJSON, err := json.Marshal(tools) Expect(err).NotTo(HaveOccurred()) Expect(toolJSON).To(MatchJSON(`{"tools":[{"name":"invoice_get","source":"custom","group":"billing","preferenceKey":"billing","defaultPermission":"ask","strict":true,"method":"GET","path":"/invoices/{id}","operationName":"invoice get","inputSchema":{"type":"object"}}]}`)) diff --git a/pkg/aiflags/flags.go b/pkg/aiflags/flags.go index 8c036e83..b7ad44d5 100644 --- a/pkg/aiflags/flags.go +++ b/pkg/aiflags/flags.go @@ -128,10 +128,28 @@ func (f ModelFlags) Resolve() (registry.Model, error) { // ResolveWith is the pure core: no ambient I/O, no globals. Saved defaults arrive // as a parameter so tests and spec-overlaying callers can drive it directly. func (f ModelFlags) ResolveWith(saved captainconfig.AIDefaults) (registry.Model, error) { + return f.ResolveWithMode(saved, "") +} + +// ResolveWithMode resolves flags and saved defaults while requesting one +// runtime mechanism for the primary model and all fallbacks. +func (f ModelFlags) ResolveWithMode(saved captainconfig.AIDefaults, mode registry.RuntimeMode) (registry.Model, error) { + if mode != "" && strings.TrimSpace(f.Mode) != "" { + explicit, ok := registry.ParseRuntimeMode(f.Mode) + if !ok { + return registry.Model{}, fmt.Errorf("invalid --mode %q (valid: %s)", f.Mode, registry.RuntimeModeList()) + } + if explicit != mode { + return registry.Model{}, fmt.Errorf("mode %q contradicts requested mode %q", explicit, mode) + } + } m, err := f.ToModel() if err != nil { return registry.Model{}, err } + if m, err = m.WithMode(mode); err != nil { + return registry.Model{}, err + } if !f.NoCache && saved.NoCache { m.NoCache = true } 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/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/runtime_config.go b/pkg/api/runtime_config.go index 16d6e3c0..2d5f86e4 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. @@ -58,9 +117,14 @@ type Config struct { CacheTTL time.Duration NoCache bool MaxConcurrent int - SessionID string - ProjectName string - SchemaRepair SchemaRepairConfig + // SessionID is the provider-native session/thread used for resume. + SessionID string + // CaptainSessionID is the authoritative Captain chat/run identity used for + // caller-tool capabilities and approvals. It must not be inferred from the + // provider-native SessionID. + CaptainSessionID string + ProjectName string + SchemaRepair SchemaRepairConfig // CanUseTool, when set, brokers tool permissions over the stream-json control // protocol: the streaming provider asks this callback before a tool that needs @@ -71,8 +135,14 @@ 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:"-"` } 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 aa5f7475..d457bd90 100644 --- a/pkg/api/spec_merge_differential_test.go +++ b/pkg/api/spec_merge_differential_test.go @@ -196,7 +196,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/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/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/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 4de8f9a6..02eb9cf0 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" ) @@ -500,7 +499,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, @@ -511,14 +510,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 } @@ -534,6 +533,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 } @@ -592,148 +594,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/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/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..f2e24bfd 100644 --- a/pkg/cli/event_renderer_ginkgo_test.go +++ b/pkg/cli/event_renderer_ginkgo_test.go @@ -1,50 +1,103 @@ 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("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, }) - Expect(writer.Close()).To(Succeed()) - output, err := io.ReadAll(reader) - Expect(err).NotTo(HaveOccurred()) - Expect(reader.Close()).To(Succeed()) + err := renderer.Flush() + Expect(err).To(MatchError(ContainSubstring("missing-call"))) + Expect(output.String()).NotTo(ContainSubstring("sensitive payload")) + }) + + 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/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/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..60d85d62 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" ) @@ -161,8 +162,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 +183,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(), } 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 2f01175d..cc8d8846 100644 --- a/pkg/cli/prompt_schema_build.go +++ b/pkg/cli/prompt_schema_build.go @@ -256,16 +256,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) { @@ -273,35 +268,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 { @@ -309,18 +289,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 0f355b6e..7d3399ae 100644 --- a/pkg/cli/prompt_schema_test.go +++ b/pkg/cli/prompt_schema_test.go @@ -93,7 +93,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 { @@ -300,19 +305,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/serve.go b/pkg/cli/serve.go index 3664dd23..2c66b8db 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{ @@ -83,7 +81,6 @@ proxies /api back to this Go process.`, 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 +101,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 +149,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, threadStore, authority, attachmentStore) if err != nil { return err } @@ -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 diff --git a/pkg/cli/serve_chat.go b/pkg/cli/serve_chat.go index ca953951..f6e309da 100644 --- a/pkg/cli/serve_chat.go +++ b/pkg/cli/serve_chat.go @@ -24,6 +24,7 @@ func newCaptainChatService( 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 +48,7 @@ func newCaptainChatService( }, }, nil }), - Tools: chatTools, MCP: mcpTools, Threads: threadStore, + Tools: chatTools, MCP: mcpTools, Threads: threadStore, 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..295f8c48 100644 --- a/pkg/cli/serve_port.go +++ b/pkg/cli/serve_port.go @@ -47,9 +47,6 @@ func (o ServeOptions) validate() error { 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..d03c9069 100644 --- a/pkg/cli/serve_port_ginkgo_test.go +++ b/pkg/cli/serve_port_ginkgo_test.go @@ -19,7 +19,7 @@ var _ = Describe("serve ports", func() { ) It("accepts an ephemeral port only in development", func() { - options := ServeOptions{Host: "localhost", Port: 0, Dev: true, UIPort: 0, ThreadsFile: "threads.json"} + options := ServeOptions{Host: "localhost", Port: 0, Dev: true, UIPort: 0} Expect(options.validate()).To(Succeed()) options.Dev = false diff --git a/pkg/cli/serve_test.go b/pkg/cli/serve_test.go index 28f275e5..2a699626 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() diff --git a/pkg/cli/session_get.go b/pkg/cli/session_get.go index df084e36..3cff9058 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" @@ -136,6 +137,30 @@ func buildSessionGetItem(ctx context.Context, db sessionGetStore, overview datab } 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}) @@ -169,20 +197,22 @@ func sessionFromPromptRun(overview database.SessionOverview, run database.Prompt 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)) 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/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..fd0f9faf 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/src/ChatLayer.tsx b/pkg/cli/webapp/src/ChatLayer.tsx index 89c397f8..820e3f2c 100644 --- a/pkg/cli/webapp/src/ChatLayer.tsx +++ b/pkg/cli/webapp/src/ChatLayer.tsx @@ -16,7 +16,7 @@ export function ChatLayer() { <> { - 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..391556bd 100644 --- a/pkg/cli/webapp/src/PromptWorkbench.test.ts +++ b/pkg/cli/webapp/src/PromptWorkbench.test.ts @@ -1,9 +1,5 @@ import { describe, expect, it } from "vitest"; -import { - promptOptions, - runtimeModelsPayload, - runtimeRowsFromPrompt, -} from "./promptWorkbenchHelpers"; +import { promptOptions } from "./promptWorkbenchHelpers"; function prompt( name: string, @@ -88,63 +84,3 @@ 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", - }, - ]); - }); - - it("preserves model options when declared rows become a run override", () => { - expect( - runtimeModelsPayload( - [ - { - model: "gemini-3.5-flash", - backend: "gemini", - effort: "high", - temperature: 0, - noCache: true, - fallbacks: [{ model: "gemini-3-flash" }], - }, - ], - [], - ), - ).toEqual([ - { - model: "gemini-3.5-flash", - backend: "gemini", - effort: "high", - temperature: 0, - noCache: true, - fallbacks: [{ model: "gemini-3-flash" }], - }, - ]); - }); -}); diff --git a/pkg/cli/webapp/src/PromptWorkbench.tsx b/pkg/cli/webapp/src/PromptWorkbench.tsx index b4d9c5c1..7f5cba62 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"; @@ -46,8 +41,6 @@ import { apiClient } from "./api"; 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 +53,25 @@ import { requiredOperation, resolvePromptOps, unwrapResponse, + type PromptDetail, type PromptSourceFilter, type PromptSummary, } from "./promptData"; import type { PromptSchemaKind } from "./promptSchemaSource"; +import { promptOptions } from "./promptWorkbenchHelpers"; import { - normalizeRuntimeModel, - promptOptions, - runtimeModelsPayload, - runtimeRowsFromPrompt, -} from "./promptWorkbenchHelpers"; + 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 +101,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 +120,7 @@ const SCRATCH_PROMPT: PromptDetail = { relPath: "scratch.prompt", writable: false, content: "", + run: EMPTY_RUN_REQUEST, }; const AGENT_TOOLS = [ @@ -220,11 +213,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 +225,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 +233,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 +256,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 +268,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 +284,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 +337,7 @@ function usePromptWorkbenchView({ undefined, () => initialPromptDetailState(), ); - const [createOpen, setCreateOpen] = useState(false); + const [writeMode, setWriteMode] = useState(); const listQuery = useQuery({ queryKey: [ @@ -415,10 +393,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 +401,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 +426,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 +436,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 +449,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 +469,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 +493,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 +539,170 @@ 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} + 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)} + /> + + 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,14 +893,10 @@ function PromptDetailPane({ draft, onDraftChange, onSchemaValidityChange, - variables, variablesValid, - onVariablesChange, onVariablesValidityChange, - runtime, - onRuntimeChange, - runtimeRows, - onRuntimeRowsChange, + runRequest, + onRunRequestChange, models, promptSchema, tools, @@ -946,14 +922,10 @@ 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[]; promptSchema?: PromptSchemaDoc; tools: ToolMeta[]; @@ -1010,16 +982,15 @@ 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 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 +1029,12 @@ function PromptDetailPane({
- } + value={runRequest} + onChange={onRunRequestChange} families={runtimeFamilies} models={promptSelectableModels(models)} tools={tools} secretSelector={CAPTAIN_SECRET_SELECTOR} - variables={variables} - onVariablesChange={onVariablesChange} onVariablesValidityChange={onVariablesValidityChange} enableAttachments {...(permissionCatalog ? { permissionCatalog } : {})} @@ -1111,7 +1071,6 @@ function PromptDetailPane({ disabled={ !runEnabled || !promptReady || - Boolean(runtimeRowsError) || (!schema && !variablesValid) } onClick={onRun} @@ -1146,8 +1105,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 +1152,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" }, @@ -1579,67 +1336,12 @@ 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; + }) => ( +