diff --git a/METRICS.md b/METRICS.md index cebae606c..b97ad614b 100644 --- a/METRICS.md +++ b/METRICS.md @@ -12,6 +12,18 @@ cargo build --release --features metrics The `[metrics]` config block is always parsed (so config validation works) but has no effect without the feature. +## Rig Alignment Telemetry Status + +`[defaults.rig_alignment]` now emits dedicated `spacebot_rig_*` Prometheus metrics. + +These metrics complement the existing request/tool metrics rather than replacing them: + +- Existing request/tool metrics still cover end-to-end execution paths (`spacebot_llm_requests_total`, duration histograms, tool counters). +- `spacebot_rig_request_semantics_total` records whether `tool_choice` / `output_schema` were applied, shadowed, rejected, or unsupported. +- `spacebot_rig_stream_sessions_total` and `spacebot_rig_time_to_first_delta_ms` separate native vs synthetic streaming behavior and expose first-delta latency. +- `spacebot_rig_tool_concurrency_total` records the request-level concurrency chosen for builtin workers. +- `spacebot_rig_drift_detected_total` records semantics requests that the custom Rig bridge could not apply cleanly. + ## Metric Inventory All metrics are prefixed with `spacebot_`. The registry uses a private `prometheus::Registry` (not the default global one) to avoid conflicts with other libraries. @@ -60,6 +72,50 @@ All metrics are prefixed with `spacebot_`. The registry uses a private `promethe **Cardinality:** 1 series. +#### `spacebot_rig_request_semantics_total` + +| Field | Value | +|-------|-------| +| Type | `IntCounterVec` | +| Labels | `process`, `provider`, `field`, `decision` | +| Instrumented in | `src/llm/model.rs` — `SpacebotModel::evaluate_request_semantics()` | +| Description | Request-semantic decisions for `tool_choice` and `output_schema`. `decision` is one of `applied`, `shadowed`, `rejected`, or `unsupported`. | + +**Cardinality:** `processes × providers × 2 fields × 4 decisions`. + +#### `spacebot_rig_stream_sessions_total` + +| Field | Value | +|-------|-------| +| Type | `IntCounterVec` | +| Labels | `process`, `provider`, `mode` | +| Instrumented in | `src/llm/model.rs` — `SpacebotModel::stream()` | +| Description | Streaming session starts. `mode` is `native` for provider-backed streaming and `synthetic` when Spacebot wraps a non-streaming completion into Rig streaming events. | + +**Cardinality:** `processes × providers × 2 modes`. + +#### `spacebot_rig_tool_concurrency_total` + +| Field | Value | +|-------|-------| +| Type | `IntCounterVec` | +| Labels | `worker_type`, `concurrency` | +| Instrumented in | `src/agent/worker.rs` — builtin worker startup | +| Description | Request-level tool concurrency chosen for worker prompts after allowlist validation and runtime clamping. | + +**Cardinality:** `worker_types × concurrency values`. + +#### `spacebot_rig_drift_detected_total` + +| Field | Value | +|-------|-------| +| Type | `IntCounterVec` | +| Labels | `surface` | +| Instrumented in | `src/llm/model.rs` — `SpacebotModel::evaluate_request_semantics()` | +| Description | Counts semantic requests that the custom Rig bridge could not apply cleanly and therefore shadowed, rejected, or treated as unsupported. | + +**Cardinality:** one series per drift surface (currently `tool_choice`, `output_schema`). + #### `spacebot_llm_tokens_total` | Field | Value | @@ -120,6 +176,18 @@ All metrics are prefixed with `spacebot_`. The registry uses a private `promethe **Cardinality:** Same as `spacebot_llm_requests_total` (per-bucket overhead is fixed, not per-series). +#### `spacebot_rig_time_to_first_delta_ms` + +| Field | Value | +|-------|-------| +| Type | `HistogramVec` | +| Labels | `process`, `provider` | +| Buckets | 25, 50, 100, 250, 500, 1000, 1500, 2500, 3500, 5000, 10000, 30000 | +| Instrumented in | `src/llm/model.rs` — instrumented streaming result wrapper | +| Description | Milliseconds from stream start until the first non-empty text delta is emitted. Covers both native and synthetic streaming paths. | + +**Cardinality:** `processes × providers`. + #### `spacebot_tool_call_duration_seconds` | Field | Value | diff --git a/README.md b/README.md index eb7087898..5c8a674f6 100644 --- a/README.md +++ b/README.md @@ -197,6 +197,57 @@ channel = "my-provider/my-model" Additional built-in providers include **Kilo Gateway**, **OpenCode Go**, **NVIDIA**, **MiniMax**, **Moonshot AI (Kimi)**, and **Z.AI Coding Plan** — configure with `kilo_key`, `opencode_go_key`, `nvidia_key`, `minimax_key`, `moonshot_key`, or `zai_coding_plan_key` in `[llm]`. +### Rig Alignment Controls + +Spacebot includes a `defaults.rig_alignment` block for request-semantics forwarding, channel/cortex streaming, and read-only worker tool concurrency. + +- **Request semantics mode** — `off`, `shadow`, or `enforced` for `tool_choice` and `output_schema` forwarding. +- **Provider allowlists** — defaults allow `openai`, `openai-chatgpt`, and `anthropic` for `tool_choice` and `output_schema`. +- **Streaming toggles** — `channel_streaming` and `cortex_chat_streaming` are opt-in and disabled by default. +- **Worker concurrency guardrails** — concurrency only applies when the exposed tool set is fully read-only and allowlisted; otherwise Spacebot forces sequential execution. + +```toml +[defaults.rig_alignment] +request_semantics_mode = "shadow" # "off" | "shadow" | "enforced" +channel_streaming = false +cortex_chat_streaming = false +worker_read_only_tool_concurrency = 1 # standard workers still stay sequential unless the exposed surface is fully read-only +worker_max_tool_concurrency = 4 +worker_read_only_tool_allowlist = [ + "read_skill", + "spacebot_docs", + "memory_recall", + "channel_recall", + "web_search", + "worker_inspect", + "email_search", + "config_inspect" +] +output_schema_provider_allowlist = ["openai", "openai-chatgpt", "anthropic"] +tool_choice_provider_allowlist = ["openai", "openai-chatgpt", "anthropic"] +``` + +In `enforced` mode, unsupported semantics fail closed; in `shadow` mode they are skipped with structured decision logs. + +The default worker surface still exposes mutating tools such as `shell`, `file`, `exec`, `task_update`, and `set_status`, so `worker_read_only_tool_concurrency > 1` only takes effect for specialized worker surfaces that expose solely allowlisted read-only tools. + +Rig-alignment observability is exposed through dedicated `spacebot_rig_*` Prometheus metrics documented in [METRICS.md](./METRICS.md), plus structured logs for request-semantic decisions, stream mode, first-delta latency, worker concurrency, and terminal delivery reason. + +For rollout, use the staged runbook in [`docs/content/docs/(configuration)/config.mdx`](./docs/content/docs/(configuration)/config.mdx): start with `request_semantics_mode = "shadow"`, then enforce in staging, then canary `channel_streaming`, `cortex_chat_streaming`, and worker concurrency one step at a time. + +When bumping Rig, rerun the rig-alignment parity suite before merging: + +```bash +cargo test forwards_ -- --nocapture +cargo test rejects_unsupported_semantics_enforced -- --nocapture +cargo test channel_streaming_reply_terminal -- --nocapture +cargo test cortex_chat_streaming_persists_final -- --nocapture +cargo test adapter_streaming_throttle -- --nocapture +cargo test worker_concurrency_allowlist -- --nocapture +cargo test worker_concurrency_reduces_wall_clock -- --nocapture +cargo test rig_parity -- --nocapture +``` + ### Skills Extensible skill system integrated with [skills.sh](https://skills.sh): diff --git a/docs/content/docs/(configuration)/config.mdx b/docs/content/docs/(configuration)/config.mdx index 1193fa241..dedf11b31 100644 --- a/docs/content/docs/(configuration)/config.mdx +++ b/docs/content/docs/(configuration)/config.mdx @@ -81,6 +81,26 @@ coding = "anthropic/claude-sonnet-4-20250514" [defaults.routing.fallbacks] "anthropic/claude-sonnet-4-20250514" = ["anthropic/claude-haiku-4.5-20250514"] +# Rig request-semantics and streaming/concurrency alignment controls. +[defaults.rig_alignment] +request_semantics_mode = "shadow" # "off" | "shadow" | "enforced" +channel_streaming = false # stream channel model deltas through adapter placeholders +cortex_chat_streaming = false # stream cortex chat SSE `text_delta` events +worker_read_only_tool_concurrency = 1 # standard workers still stay sequential unless the exposed surface is fully read-only +worker_max_tool_concurrency = 4 # hard cap for per-request tool concurrency +worker_read_only_tool_allowlist = [ + "read_skill", + "spacebot_docs", + "memory_recall", + "channel_recall", + "web_search", + "worker_inspect", + "email_search", + "config_inspect" +] +output_schema_provider_allowlist = ["openai", "openai-chatgpt", "anthropic"] +tool_choice_provider_allowlist = ["openai", "openai-chatgpt", "anthropic"] + # Context compaction thresholds (fraction of context_window). [defaults.compaction] background_threshold = 0.80 # background summarization @@ -249,6 +269,7 @@ Most config values are hot-reloaded when their files change. Spacebot watches `c | `max_concurrent_branches` | Yes | Next branch spawn checks new limit | | Browser config | Yes | Next worker spawn uses new config | | Warmup config | Yes | Next warmup pass uses new values | +| Rig-alignment config | Yes | Next LLM request/worker turn reads new semantics, streaming, and tool-concurrency settings | | Identity files (SOUL.md, etc.) | Yes | Next channel message renders new identity | | Skills (SKILL.md files) | Yes | Next message / worker spawn sees new skills | | Bindings | Yes | Next message routes using new bindings | @@ -489,6 +510,138 @@ Map of model names to ordered fallback chains. Used when the primary model retur "anthropic/claude-sonnet-4-20250514" = ["anthropic/claude-haiku-4.5-20250514"] ``` +### `[defaults.rig_alignment]` + +Controls how Spacebot forwards Rig request semantics (`tool_choice`, `output_schema`) to providers, and whether channel/cortex streaming plus read-only worker tool concurrency are enabled. + +| Key | Type | Default | Description | +|-----|------|---------|-------------| +| `request_semantics_mode` | string | `"shadow"` | `off`: never forward `tool_choice`/`output_schema`; `shadow`: apply only when supported; `enforced`: fail unsupported requests explicitly | +| `channel_streaming` | bool | false | Enables channel `stream_prompt(...)` execution and adapter placeholder streaming | +| `cortex_chat_streaming` | bool | false | Enables cortex-chat `stream_prompt(...)` and SSE `text_delta` events | +| `worker_read_only_tool_concurrency` | integer | 1 | Requested per-request tool concurrency for workers (clamped to `worker_max_tool_concurrency`; standard workers still stay sequential unless their exposed surface is fully read-only) | +| `worker_max_tool_concurrency` | integer | 4 | Hard upper bound for worker request-level tool concurrency | +| `worker_read_only_tool_allowlist` | string[] | read-only defaults | Tools that are eligible for concurrent execution; any non-allowlisted exposed tool forces sequential execution | +| `output_schema_provider_allowlist` | string[] | `["openai","openai-chatgpt","anthropic"]` | Provider IDs allowed to receive structured output schema fields | +| `tool_choice_provider_allowlist` | string[] | `["openai","openai-chatgpt","anthropic"]` | Provider IDs allowed to receive tool-choice fields | + +`request_semantics_mode` behavior: + +- `off`: disables forwarding of both request semantics fields. +- `shadow`: evaluates support per provider and request shape, applies only supported fields, and logs the decision. +- `enforced`: same evaluation, but unsupported requested fields return an explicit provider error. + +Worker concurrency safety behavior: + +- Concurrency remains sequential unless **all** exposed worker tools for that request are allowlisted. +- If the configured allowlist itself contains mutable tools, Spacebot warns and forces sequential execution. +- The default worker surface still exposes mutating tools such as `shell`, `file`, `exec`, `task_update`, and `set_status`, so `worker_read_only_tool_concurrency > 1` only takes effect for specialized worker surfaces that expose solely allowlisted read-only tools. + +Rig-alignment observability: + +- `spacebot_rig_request_semantics_total{process,provider,field,decision}` +- `spacebot_rig_stream_sessions_total{process,provider,mode}` +- `spacebot_rig_time_to_first_delta_ms{process,provider}` +- `spacebot_rig_tool_concurrency_total{worker_type,concurrency}` +- `spacebot_rig_drift_detected_total{surface}` + +## Rig Alignment Rollout Runbook + +Use config-first rollout and rollback. Do not enable multiple new knobs at once. + +### Stage 1: Semantics Shadow Mode + +- Enable: + `request_semantics_mode = "shadow"` +- Evidence: + `cargo test forwards_ -- --nocapture` + `cargo test rejects_unsupported_semantics_enforced -- --nocapture` + `rg "request semantic decision" logs debug.log -g '*.log'` +- Rollback trigger: + supported providers start logging `decision=shadowed` for requests that should apply cleanly, or provider errors increase unexpectedly. +- Rollback: + `request_semantics_mode = "off"` + +### Stage 2: Semantics Enforced In Staging + +- Enable: + `request_semantics_mode = "enforced"` +- Evidence: + `cargo test rig_parity -- --nocapture` + `rg "request semantic decision|request semantics rejected" logs debug.log -g '*.log'` +- Rollback trigger: + staging traffic hits explicit compatibility errors for expected production routes. +- Rollback: + `request_semantics_mode = "shadow"` + +### Stage 3: Channel Streaming Canary + +- Enable: + `channel_streaming = true` +- Evidence: + `cargo test channel_streaming_reply_terminal -- --nocapture` + `cargo test adapter_streaming_throttle -- --nocapture` + `rg "stream_mode|time_to_first_delta_ms|terminal_reason" logs debug.log -g '*.log'` +- Rollback trigger: + missing final replies, placeholder cleanup failures, or adapter edit throttling regressions. +- Rollback: + `channel_streaming = false` + +### Stage 4: Cortex Chat Streaming Canary + +- Enable: + `cortex_chat_streaming = true` +- Evidence: + `cargo test cortex_chat_streaming_persists_final -- --nocapture` + `rg "stream_mode|time_to_first_delta_ms" logs debug.log -g '*.log'` +- Rollback trigger: + partial SSE events stop arriving or final assistant responses stop persisting. +- Rollback: + `cortex_chat_streaming = false` + +### Stage 5: Worker Concurrency Canary + +- Enable: + `worker_read_only_tool_concurrency = 2` + Keep `worker_max_tool_concurrency = 4` unless a smaller cap is required. +- Evidence: + `cargo test worker_concurrency_allowlist -- --nocapture` + `cargo test worker_concurrency_reduces_wall_clock -- --nocapture` + `rg "worker tool concurrency resolved|forcing sequential execution" logs debug.log -g '*.log'` +- Rollback trigger: + any mutable or mixed tool surface reaches concurrency > 1, or tool error rates regress during canary. +- Rollback: + `worker_read_only_tool_concurrency = 1` + +## Rig Upgrade Checklist + +Run this parity suite on every Rig upgrade before merge: + +```bash +cargo test forwards_ -- --nocapture +cargo test rejects_unsupported_semantics_enforced -- --nocapture +cargo test channel_streaming_reply_terminal -- --nocapture +cargo test cortex_chat_streaming_persists_final -- --nocapture +cargo test adapter_streaming_throttle -- --nocapture +cargo test worker_concurrency_allowlist -- --nocapture +cargo test worker_concurrency_reduces_wall_clock -- --nocapture +cargo test rig_parity -- --nocapture +just preflight +just gate-pr +``` + +```toml +[defaults.rig_alignment] +request_semantics_mode = "enforced" +channel_streaming = true +cortex_chat_streaming = true +worker_read_only_tool_concurrency = 3 +worker_max_tool_concurrency = 4 +worker_read_only_tool_allowlist = ["read_skill", "spacebot_docs", "web_search"] +output_schema_provider_allowlist = ["openai", "openai-chatgpt", "anthropic"] +tool_choice_provider_allowlist = ["openai", "openai-chatgpt", "anthropic"] +``` + ### `[defaults.compaction]` | Key | Type | Default | Description | diff --git a/interface/src/api/client.ts b/interface/src/api/client.ts index d2c15719e..87a7cc3fd 100644 --- a/interface/src/api/client.ts +++ b/interface/src/api/client.ts @@ -48,6 +48,12 @@ export interface OutboundMessageDeltaEvent { aggregated_text: string; } +export interface OutboundStreamEndEvent { + type: "outbound_stream_end"; + agent_id: string; + channel_id: string; +} + export interface TypingStateEvent { type: "typing_state"; agent_id: string; @@ -150,6 +156,7 @@ export type ApiEvent = | InboundMessageEvent | OutboundMessageEvent | OutboundMessageDeltaEvent + | OutboundStreamEndEvent | TypingStateEvent | WorkerStartedEvent | WorkerStatusEvent @@ -535,6 +542,9 @@ export interface CortexChatMessagesResponse { export type CortexChatSSEEvent = | { type: "thinking" } + | { type: "text_delta"; text_delta: string } + | { type: "tool_started"; tool: string } + | { type: "tool_completed"; tool: string; result_preview: string } | { type: "done"; full_text: string } | { type: "error"; message: string }; diff --git a/interface/src/hooks/useChannelLiveState.ts b/interface/src/hooks/useChannelLiveState.ts index 40f73bc7b..f9091db90 100644 --- a/interface/src/hooks/useChannelLiveState.ts +++ b/interface/src/hooks/useChannelLiveState.ts @@ -7,6 +7,7 @@ import { type InboundMessageEvent, type OutboundMessageDeltaEvent, type OutboundMessageEvent, + type OutboundStreamEndEvent, type TimelineItem, type ToolCompletedEvent, type ToolStartedEvent, @@ -102,6 +103,65 @@ function assistantMessageItem( }; } +export function applyOutboundMessageDelta( + state: ChannelLiveState, + event: OutboundMessageDeltaEvent, +): ChannelLiveState { + const deltaText = event.text_delta; + const aggregatedText = event.aggregated_text; + if (!deltaText && !aggregatedText) { + return state; + } + + const streamMessageId = state.streamingMessageId; + if (streamMessageId) { + const streamIndex = state.timeline.findIndex( + (item) => item.type === "message" && item.id === streamMessageId, + ); + + if (streamIndex >= 0) { + const timeline = [...state.timeline]; + const streamItem = timeline[streamIndex]; + if (streamItem.type === "message") { + timeline[streamIndex] = { + ...streamItem, + content: deltaText + ? streamItem.content + deltaText + : aggregatedText, + }; + } + return { ...state, timeline }; + } + } + + const messageId = `stream-${generateId()}`; + const initialContent = deltaText || aggregatedText; + return { + ...state, + timeline: [ + ...state.timeline, + assistantMessageItem(messageId, event.agent_id, initialContent), + ], + streamingMessageId: messageId, + }; +} + +export function applyOutboundStreamEnd(state: ChannelLiveState): ChannelLiveState { + const streamMessageId = state.streamingMessageId; + if (!streamMessageId) { + return { ...state, isTyping: false }; + } + + return { + ...state, + timeline: state.timeline.filter( + (item) => item.type !== "message" || item.id !== streamMessageId, + ), + streamingMessageId: null, + isTyping: false, + }; +} + /** * Manages all live channel state from SSE events, message history loading, * and status snapshot fetching. Returns the state map and SSE event handlers. @@ -302,65 +362,26 @@ export function useChannelLiveState(channels: ChannelInfo[]) { const event = data as OutboundMessageDeltaEvent; setLiveStates((prev) => { const existing = getOrCreate(prev, event.channel_id); - const streamMessageId = existing.streamingMessageId; - - if (streamMessageId) { - const streamIndex = existing.timeline.findIndex( - (item) => item.type === "message" && item.id === streamMessageId, - ); - - if (streamIndex >= 0) { - const timeline = [...existing.timeline]; - const streamItem = timeline[streamIndex]; - if (streamItem.type === "message") { - timeline[streamIndex] = { - ...streamItem, - content: event.aggregated_text, - }; - } - return { - ...prev, - [event.channel_id]: { ...existing, timeline }, - }; - } - - const messageId = `stream-${generateId()}`; - return { - ...prev, - [event.channel_id]: { - ...existing, - timeline: [ - ...existing.timeline, - assistantMessageItem( - messageId, - event.agent_id, - event.aggregated_text, - ), - ], - streamingMessageId: messageId, - }, - }; - } - - const messageId = `stream-${generateId()}`; return { ...prev, [event.channel_id]: { - ...existing, - timeline: [ - ...existing.timeline, - assistantMessageItem( - messageId, - event.agent_id, - event.aggregated_text, - ), - ], - streamingMessageId: messageId, + ...applyOutboundMessageDelta(existing, event), }, }; }); }, []); + const handleOutboundStreamEnd = useCallback((data: unknown) => { + const event = data as OutboundStreamEndEvent; + setLiveStates((prev) => { + const existing = getOrCreate(prev, event.channel_id); + return { + ...prev, + [event.channel_id]: applyOutboundStreamEnd(existing), + }; + }); + }, []); + const handleTypingState = useCallback((data: unknown) => { const event = data as TypingStateEvent; setLiveStates((prev) => { @@ -806,6 +827,7 @@ export function useChannelLiveState(channels: ChannelInfo[]) { inbound_message: handleInboundMessage, outbound_message: handleOutboundMessage, outbound_message_delta: handleOutboundMessageDelta, + outbound_stream_end: handleOutboundStreamEnd, typing_state: handleTypingState, worker_started: handleWorkerStarted, worker_status: handleWorkerStatus, diff --git a/interface/tests/useChannelLiveState.test.ts b/interface/tests/useChannelLiveState.test.ts new file mode 100644 index 000000000..d93389895 --- /dev/null +++ b/interface/tests/useChannelLiveState.test.ts @@ -0,0 +1,94 @@ +import { expect, test } from "bun:test"; + +import type { ChannelLiveState } from "../src/hooks/useChannelLiveState"; + +(globalThis as typeof globalThis & { + window?: { __SPACEBOT_BASE_PATH?: string }; +}).window = { + __SPACEBOT_BASE_PATH: "", +}; + +const { applyOutboundMessageDelta, applyOutboundStreamEnd } = await import( + "../src/hooks/useChannelLiveState" +); + +function baseLiveState(): ChannelLiveState { + return { + isTyping: true, + timeline: [], + workers: {}, + branches: {}, + streamingMessageId: null, + historyLoaded: true, + hasMore: false, + loadingMore: false, + }; +} + +test("outbound stream end clears the active placeholder before the next stream", () => { + const afterFirstDelta = applyOutboundMessageDelta(baseLiveState(), { + type: "outbound_message_delta", + agent_id: "agent-1", + channel_id: "channel-1", + text_delta: "partial", + aggregated_text: "partial", + }); + + expect(afterFirstDelta.streamingMessageId).not.toBeNull(); + expect(afterFirstDelta.timeline).toHaveLength(1); + + const firstStreamMessageId = afterFirstDelta.streamingMessageId; + const afterStreamEnd = applyOutboundStreamEnd(afterFirstDelta); + + expect(afterStreamEnd.streamingMessageId).toBeNull(); + expect(afterStreamEnd.isTyping).toBe(false); + expect(afterStreamEnd.timeline).toHaveLength(0); + + const afterNextDelta = applyOutboundMessageDelta(afterStreamEnd, { + type: "outbound_message_delta", + agent_id: "agent-1", + channel_id: "channel-1", + text_delta: "new reply", + aggregated_text: "new reply", + }); + + expect(afterNextDelta.streamingMessageId).not.toBeNull(); + expect(afterNextDelta.streamingMessageId).not.toBe(firstStreamMessageId); + expect(afterNextDelta.timeline).toHaveLength(1); + expect(afterNextDelta.timeline[0]).toMatchObject({ + type: "message", + content: "new reply", + }); +}); + +test("aggregated outbound deltas replace placeholder content instead of appending", () => { + const afterFirstDelta = applyOutboundMessageDelta(baseLiveState(), { + type: "outbound_message_delta", + agent_id: "agent-1", + channel_id: "channel-1", + text_delta: "partial", + aggregated_text: "partial", + }); + + const afterAggregatedOnlyUpdate = applyOutboundMessageDelta(afterFirstDelta, { + type: "outbound_message_delta", + agent_id: "agent-1", + channel_id: "channel-1", + text_delta: "", + aggregated_text: "partial reply", + }); + + expect(afterAggregatedOnlyUpdate.timeline).toHaveLength(1); + expect(afterAggregatedOnlyUpdate.timeline[0]).toMatchObject({ + type: "message", + content: "partial reply", + }); +}); + +test("outbound stream end without a placeholder still clears typing", () => { + const updated = applyOutboundStreamEnd(baseLiveState()); + + expect(updated.streamingMessageId).toBeNull(); + expect(updated.isTyping).toBe(false); + expect(updated.timeline).toHaveLength(0); +}); diff --git a/rig-spacebot-alignment-plan.md b/rig-spacebot-alignment-plan.md new file mode 100644 index 000000000..46c0b57f0 --- /dev/null +++ b/rig-spacebot-alignment-plan.md @@ -0,0 +1,1023 @@ +# Plan: Rig-Spacebot Alignment And Leverage + +```yaml +plan_contract: + plan_id: rig-spacebot-alignment-2026-03-05 + generated_at: 2026-03-05 + owner: + primary: spacebot-core + supporting: + - llm-runtime + - agent-orchestration + - messaging-adapters + - observability + - qa-docs + handoff_state: implemented-and-verified + risk_level: high + assumptions: + - Spacebot must keep its custom CompletionModel and LlmManager-based provider stack. + - Spacebot must not adopt Rig built-in provider clients, RAG/vector-store integrations, Agent-as-Tool, or Pipeline system. + - Existing per-process architecture (channel, branch, worker, compactor, cortex) remains the core execution model. + - New behavior must be rollout-gated and backward compatible by default. + - No historical migration files may be edited; any new persisted schema change requires a new migration. + open_questions: [] +``` + +## Goals + +- Fully align Spacebot's custom Rig integration with the parts of Rig 0.31 that materially improve correctness, latency, observability, and maintainability. +- Implement faithful support for Rig request semantics currently dropped by `SpacebotModel`, specifically `tool_choice` and `output_schema`. +- Introduce native Rig streaming into user-facing paths where it provides real time-to-first-token improvement without violating current safety guarantees. +- Add controlled, opt-in request-level tool concurrency for clearly read-only worker tool sets. +- Add parity and drift-detection tests so future Rig upgrades do not silently break Spacebot's custom integration layer. +- Preserve Spacebot's architectural intent: delegation-first multi-process orchestration, custom routing, custom provider stack, and process-scoped tool isolation. + +## Non-goals + +- Do not replace `SpacebotModel` with Rig built-in provider clients. +- Do not adopt Rig Agent-as-Tool composition for core channel/branch/worker orchestration. +- Do not adopt Rig vector store or dynamic tool retrieval as the primary memory/tool architecture. +- Do not replace current memory system, task system, or channel history model. +- Do not enable broad concurrent tool execution for side-effecting or secret-handling tools. +- Do not require database migrations unless a later implementation decision introduces new persisted state. + +## Executive Summary + +Spacebot already uses Rig deeply, but primarily as a low-level substrate: `CompletionModel`, `AgentBuilder`, `PromptHook`, `ToolServer`, Rig message/history types, and Rig's multi-turn tool loop. The main gaps are not architectural absence; they are semantic mismatch and underused Rig capabilities: + +1. `SpacebotModel` does not currently forward Rig `tool_choice` and `output_schema` request semantics. +2. App-level process loops use non-streaming `prompt(...)` paths even though the model implements Rig `stream()`. +3. Request-level `with_tool_concurrency(...)` is unused, even where worker tools are independent and read-only. +4. Drift detection is weak: Spacebot's custom Rig bridge can silently diverge from Rig's expected request surface on future upgrades. + +The recommended implementation sequence is: + +1. Request semantic parity +2. Streaming adoption in channel and cortex-chat +3. Safe read-only tool concurrency +4. Drift detection, instrumentation, and documentation hardening + +This sequence maximizes correctness first, then user-visible latency wins, then performance wins, while preserving Spacebot's repo-level architecture constraints. + +## Existing Architecture Snapshot + +### Relevant Existing Components + +- `src/llm/model.rs` + - `SpacebotModel` + - custom request serialization for OpenAI-compatible, Responses API, Gemini-compatible, and Anthropic + - custom streaming parsing and aggregation +- `src/llm/manager.rs` + - provider config resolution, OAuth refresh, HTTP client ownership, cooldown tracking +- `src/llm/routing.rs` + - process-type routing and fallback selection +- `src/hooks/spacebot.rs` + - `PromptHook` implementation + - leak detection + - tool nudge behavior + - status event emission +- `src/agent/channel.rs` + - main user-facing prompt loop + - per-turn tool add/remove + - history reconciliation +- `src/agent/worker.rs` + - segmented multi-turn worker loop + - tool nudge retry + - compaction recovery +- `src/agent/cortex_chat.rs` + - interactive admin/cortex chat path +- `src/tools.rs` + - branch/worker/cortex tool server construction +- `src/messaging/*.rs` + - adapter-specific streaming edit/render behavior already exists at outbound transport layer + +### Architectural Boundaries + +- Rig boundary: + - Rig owns agent loop primitives, request builders, hook interfaces, tool server protocol, message model, and streaming interfaces. +- Spacebot boundary: + - Spacebot owns provider transport, model routing, secrets handling, message persistence, task orchestration, memory system, and adapter delivery semantics. +- Provider boundary: + - Provider-specific serialization must happen in `SpacebotModel` and `src/llm/anthropic/*`, not in higher-level agents. +- Messaging boundary: + - End-user streaming behavior must remain adapter-aware and respect platform-specific edit rate limits. + +### Architectural Invariants + +- `SpacebotModel` remains the single translation layer from Rig `CompletionRequest` to provider payloads. +- Channel, branch, worker, compactor, and cortex remain separate process types with separate prompts and tool surfaces. +- No end-user visible text may bypass leak detection or reply-tool safety checks. +- Channel history is not committed until a terminal condition is reached and history reconciliation completes. +- Worker tool concurrency may only apply to explicitly allowlisted read-only tools. +- Unsupported provider features must never fail silently when the caller requested strict semantics. + +## Target Architecture + +### New Capability Surface + +The implementation introduces four logical capabilities: + +1. Request Semantic Parity Layer +2. Streaming Execution Layer +3. Read-Only Tool Concurrency Layer +4. Drift Detection And Observability Layer + +These are capabilities, not a new top-level subsystem. They fit into the current module boundaries. + +### Components + +#### 1. Request Semantic Parity Layer + +Primary files: + +- `src/llm/model.rs` +- `src/llm/anthropic/params.rs` +- `src/agent/cortex.rs` +- `src/config/*` + +Responsibilities: + +- Carry Rig `CompletionRequest.tool_choice` into provider payloads when supported. +- Carry Rig `CompletionRequest.output_schema` into provider payloads when supported. +- Explicitly classify unsupported combinations. +- Enforce rollout-mode behavior (`off`, `shadow`, `enforced`). + +#### 2. Streaming Execution Layer + +Primary files: + +- `src/agent/channel.rs` +- `src/agent/cortex_chat.rs` +- `src/hooks/spacebot.rs` +- `src/messaging/slack.rs` +- `src/messaging/discord.rs` +- `src/messaging/telegram.rs` +- `src/api/state.rs` + +Responsibilities: + +- Switch selected paths from `prompt(...)` to Rig streaming request flow. +- Preserve existing hook semantics for deltas, tool calls, terminal responses, and cancellation. +- Reconcile final streamed history with existing persistence and reply semantics. +- Preserve adapter-specific rate limiting and edit chunking. + +#### 3. Read-Only Tool Concurrency Layer + +Primary files: + +- `src/agent/worker.rs` +- `src/tools.rs` +- selected `src/tools/*.rs` +- `src/config/*` + +Responsibilities: + +- Define a concurrency-safe tool allowlist. +- Apply `with_tool_concurrency(...)` only for eligible worker prompts. +- Preserve ordering and side-effect guarantees for non-allowlisted tools. + +#### 4. Drift Detection And Observability Layer + +Primary files: + +- `src/llm/model.rs` +- `src/hooks/spacebot.rs` +- `src/agent/channel.rs` +- `METRICS.md` +- `README.md` +- `docs/content/docs/(configuration)/config.mdx` +- `tests/*` + +Responsibilities: + +- Detect when Rig request fields are being dropped by Spacebot's custom bridge. +- Emit metrics for semantic application, streaming latency, concurrency usage, and unsupported-provider decisions. +- Provide tests that fail when Rig request surface and Spacebot behavior diverge. + +## Data Flow + +### Flow A: Request Semantic Parity + +1. Process builds Rig agent with `AgentBuilder`. +2. Caller invokes `prompt(...)` or `prompt_typed(...)`. +3. Rig builds `CompletionRequest`. +4. `SpacebotModel::completion(...)` or `SpacebotModel::stream(...)` receives request. +5. New semantic-parity helper inspects: + - `tool_choice` + - `output_schema` + - `process_type` + - rollout mode + - provider capability +6. Provider serializer applies, shadows, or rejects unsupported fields. +7. Metrics/logs record applied vs unsupported vs shadowed decisions. +8. Response returns through existing hook and agent loop. + +### Flow B: Channel Streaming + +1. Channel builds per-turn tool server. +2. Channel creates Rig agent. +3. Under feature flag, channel uses Rig streaming prompt request instead of one-shot prompt. +4. Streaming hook emits text deltas and tool-call deltas. +5. Adapter updates in-flight message according to adapter edit cadence. +6. Terminal condition occurs: + - final text response + - `reply` tool succeeds + - cancellation + - max turns +7. Final response is reconciled into channel history using existing history-commit rules. +8. Optional fallback text behavior remains preserved for no-reply cases. + +### Flow C: Worker Read-Only Tool Concurrency + +1. Worker selects task prompt. +2. Worker computes eligible concurrency profile from tool surface plus config. +3. If eligible: + - issue request with `with_tool_concurrency(n)` +4. Rig executes independent tool calls concurrently. +5. Hook still receives each tool start/result event. +6. Existing status/leak/error handling remains authoritative. +7. Worker loop continues or stops under existing segment/max-turn logic. + +## Proposed Data Model And Schemas + +### Config Schema + +Add a new runtime-config group for Rig alignment behavior. + +```toml +[defaults.rig_alignment] +request_semantics_mode = "shadow" # off | shadow | enforced +channel_streaming = false +cortex_chat_streaming = false +worker_read_only_tool_concurrency = 1 +worker_max_tool_concurrency = 4 +worker_read_only_tool_allowlist = [ + "read_skill", + "spacebot_docs", + "memory_recall", + "channel_recall", + "web_search", + "worker_inspect", + "email_search", + "config_inspect" +] +output_schema_provider_allowlist = [ + "openai", + "openai-chatgpt", + "anthropic" +] +tool_choice_provider_allowlist = [ + "openai", + "openai-chatgpt", + "anthropic" +] +``` + +### Rust Config Types + +Add to `src/config/types.rs`, `src/config/toml_schema.rs`, `src/config/load.rs`, and `src/config/runtime.rs`: + +```rust +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum RigSemanticsMode { + Off, + Shadow, + Enforced, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RigAlignmentConfig { + pub request_semantics_mode: RigSemanticsMode, + pub channel_streaming: bool, + pub cortex_chat_streaming: bool, + pub worker_read_only_tool_concurrency: usize, + pub worker_max_tool_concurrency: usize, + pub worker_read_only_tool_allowlist: Vec, + pub output_schema_provider_allowlist: Vec, + pub tool_choice_provider_allowlist: Vec, +} +``` + +### Semantic Decision Model + +Add an internal non-persisted decision struct in `src/llm/model.rs`: + +```rust +struct RequestSemanticDecision { + apply_tool_choice: bool, + apply_output_schema: bool, + rejected_reason: Option, + shadow_only: bool, +} +``` + +### Metrics Schema + +Add or extend metrics with the following label sets: + +- `spacebot_rig_request_semantics_total{process,provider,field,decision}` + - `field`: `tool_choice` | `output_schema` + - `decision`: `applied` | `shadowed` | `rejected` | `unsupported` +- `spacebot_rig_stream_sessions_total{process,provider,mode}` + - `mode`: `native` | `synthetic` +- `spacebot_rig_time_to_first_delta_ms_bucket{process,provider}` +- `spacebot_rig_tool_concurrency_total{worker_type,concurrency}` +- `spacebot_rig_drift_detected_total{surface}` + +## Security And Privacy Model + +### Threats + +#### Threat 1: Secret leakage through streamed deltas + +Risk: + +- Streaming emits partial text earlier than final reply-tool gating. +- Secrets may appear in deltas before final reconciliation. + +Mitigations: + +- Route all user-visible delta emission through the same leak scrub path as final output. +- Do not stream raw tool args or tool results to end users. +- Preserve `reply`-tool blocking behavior and reply-result leak checks in `SpacebotHook`. +- Cap delta payload lengths before event fan-out. + +#### Threat 2: Unsafe concurrent tool execution + +Risk: + +- Concurrent side effects can reorder writes, statuses, or task updates. + +Mitigations: + +- Concurrency applies only to explicit read-only allowlist. +- `set_status`, `file`, `exec`, `shell`, `secret_set`, `task_update`, `browser`, and all write-capable tools remain sequential. +- Add a validator test that fails if a non-allowlisted tool is marked concurrency-safe without an explicit review. + +#### Threat 3: Structured output contract silently not enforced + +Risk: + +- `prompt_typed(...)` appears safe but `SpacebotModel` drops `output_schema`. + +Mitigations: + +- In `enforced` mode, reject `output_schema` requests for unsupported providers. +- Add serializer tests that assert the schema is present in outgoing requests. +- Emit parity metrics for every structured-output request. + +#### Threat 4: Provider-specific tool choice mismatch + +Risk: + +- `ToolChoice::Required` or `Specific` may be requested but ignored. + +Mitigations: + +- Capability gating per provider/API type. +- Fail closed in `enforced` mode for unsupported mappings. +- Shadow logs during rollout before enabling strict mode globally. + +### Secrets Handling + +- Continue to source provider keys only through `LlmManager` and the existing secret store. +- Do not add any new secret propagation path. +- Do not log raw output schema contents if they may encode user data; log schema name/hash instead. +- Continue scrubbing tool outputs before SSE/dashboard emission. +- Preserve current worker/branch secret-scrub behavior in `SpacebotHook`. + +## Performance Targets + +### Request Semantic Parity Targets + +- Added serialization overhead for `tool_choice` / `output_schema` handling: + - p50 <= 1 ms per request + - p95 <= 5 ms per request +- No measurable increase in provider HTTP request size above: + - +3 KB p50 + - +15 KB p95 for structured output requests + +### Streaming Targets + +- Channel time-to-first-visible-delta on native-streaming providers: + - p50 <= 1.5 s + - p95 <= 3.5 s +- Synthetic-stream providers must not regress full-response latency by >5%. +- Adapter edit cadence: + - Slack/Discord: no more than 2 updates/second/message + - Telegram: no more than 1 update/second/message + +### Tool Concurrency Targets + +- For a worker turn with 3 eligible read-only tools, wall-clock completion time should improve by >=30% versus sequential baseline in test harness. +- Maximum request-level tool concurrency: + - default: 1 + - rollout cap: 4 +- No increase in tool error rate >2 percentage points during canary. + +## Instrumentation Plan + +### Logs + +Emit structured logs for: + +- request semantic decisions +- streaming session start/stop +- time-to-first-delta +- native vs synthetic stream mode +- concurrency-enabled worker turns +- rejected unsupported semantics +- history reconciliation result after streamed completion + +Required log keys: + +- `agent_id` +- `process_type` +- `provider` +- `model` +- `request_semantics_mode` +- `tool_choice_requested` +- `output_schema_requested` +- `streaming_enabled` +- `stream_mode` +- `tool_concurrency` +- `terminal_reason` + +### Metrics + +Use the metrics listed in the metrics schema section. + +### Tracing + +- Continue using Rig span model for prompt/tool execution. +- Add span fields: + - `spacebot.rig.tool_choice_applied` + - `spacebot.rig.output_schema_applied` + - `spacebot.rig.stream_mode` + - `spacebot.rig.tool_concurrency` + +### Baseline Measurement Before Rollout + +Collect a 7-day baseline in current production or staging for: + +- channel response latency +- time to first outbound message edit +- worker tool error rate +- provider distribution +- MCP tool count distribution per worker + +## Error Handling, Retries, And No-Silent-Failure Policy + +### Policy + +- No newly introduced request semantic may be dropped silently. +- Every unsupported semantic must produce one of: + - explicit error + - warning + metric in shadow mode + - structured status event +- Every rollout-gated behavior must be observable in logs and metrics. + +### Retry Rules + +- Preserve existing model retry and fallback logic in `SpacebotModel`. +- Preserve worker overflow retry and segment retry behavior. +- Preserve tool-nudge retry behavior. +- Streaming adapter message-edit failures: + - retry up to 2 times per update burst + - then downgrade to final-response-only delivery for that message + - log adapter-specific warning with message key + +### Unsupported Capability Behavior + +#### `output_schema` + +- `off`: ignore request and log nothing extra +- `shadow`: do not apply request, emit warning + metric +- `enforced`: return explicit provider compatibility error + +#### `tool_choice` + +- `off`: ignore request and log nothing extra +- `shadow`: do not apply request, emit warning + metric +- `enforced`: return explicit provider compatibility error + +### No Silent Failure Rules + +- No `Result` may be dropped in new code paths unless the existing repo policy explicitly allows best-effort event sends. +- Every new fallback path must: + - increment a metric + - emit a structured log + - preserve root-cause context + +## Orchestration And State Semantics + +### Execution States + +Apply the following logical execution state model to channel streaming and worker concurrency flows. + +| State | Description | Entry Condition | Exit Condition | +|---|---|---|---| +| `Idle` | No prompt running | process available | prompt starts | +| `BuildingRequest` | Rig request assembly in progress | prompt accepted | request built or rejected | +| `ValidatingSemantics` | capability gating for `tool_choice` / `output_schema` | request built | semantic decision made | +| `AwaitingProvider` | provider request dispatched | request sent | first response chunk or full response | +| `StreamingText` | text delta stream active | first delta received | tool call, terminal text, cancel, or error | +| `ExecutingTools` | one or more tool calls executing | tool call received | all tool results returned or failure | +| `AwaitingContinuation` | multi-turn waiting for next model step | tool results committed | next provider call or terminal state | +| `ReconcilingHistory` | final history normalization | terminal response received | history persisted | +| `Completed` | successful terminal condition | final text or reply delivered | none | +| `Cancelled` | prompt cancelled by hook/user/system | cancel reason received | none | +| `Failed` | unrecoverable error | error returned after retries | none | + +### State Transitions + +- `Idle -> BuildingRequest` +- `BuildingRequest -> ValidatingSemantics` +- `ValidatingSemantics -> Failed` if enforced unsupported semantics +- `ValidatingSemantics -> AwaitingProvider` +- `AwaitingProvider -> StreamingText` if streaming deltas start +- `AwaitingProvider -> ExecutingTools` if tool call arrives before text +- `AwaitingProvider -> ReconcilingHistory` if full final response returns +- `StreamingText -> ExecutingTools` if tool call arrives mid-stream +- `StreamingText -> ReconcilingHistory` on final terminal text +- `ExecutingTools -> AwaitingContinuation` when all tool results return +- `AwaitingContinuation -> AwaitingProvider` on next turn +- `ReconcilingHistory -> Completed|Cancelled|Failed` + +### Stop Conditions + +- Channel: + - successful `reply` tool result + - explicit `skip` + - user/system cancellation + - unrecoverable provider error + - max-turn exhaustion +- Worker: + - successful final response after outcome signal + - max segments reached + - unrecoverable provider error + - cancellation + +### Continue Conditions + +- Continue only when: + - tool results are fully captured + - current prompt has not exceeded max-turn or max-segment bounds + - no enforced semantic compatibility failure occurred + - no leak guard terminated the loop + +### Retry And Backoff Rules + +- Provider retries: + - keep existing exponential backoff and fallback chain logic +- Streaming message-edit retries: + - 2 retries max per burst + - no blind infinite retry +- Concurrency downgrade: + - if a concurrency-enabled worker turn hits allowlist violation or ordering assertion, rerun that prompt sequentially once and emit metric + +### Reconciliation Loop For Drift + +Implement a CI and runtime drift loop: + +1. CI drift tests: + - verify outgoing payload includes or rejects all expected Rig request fields + - fail when serializer behavior diverges +2. Runtime drift metrics: + - count requested vs applied semantics + - count synthetic-stream fallbacks +3. Upgrade reconciliation: + - when Rig version changes, run targeted parity suite + - if new `CompletionRequest` fields appear in Rig, add explicit decision handling before upgrade is considered complete + +## Implementation Phases + +## Phase 0: Baseline And Feature Flag Scaffolding + +### Deliverable + +- New config surface exists but defaults preserve current behavior. +- Baseline latency/error metrics collection is available. + +### Scope + +- `src/config/types.rs` +- `src/config/toml_schema.rs` +- `src/config/load.rs` +- `src/config/runtime.rs` +- `METRICS.md` + +## Phase 1: Request Semantic Parity + +### Deliverable + +- `tool_choice` and `output_schema` are explicitly handled in `SpacebotModel`. +- Unsupported combinations produce visible behavior according to mode. + +### Scope + +- `src/llm/model.rs` +- `src/llm/anthropic/params.rs` +- tests in `src/llm/model.rs` and `tests/` + +## Phase 2: Streaming Channel And Cortex Chat + +### Deliverable + +- Channel and cortex-chat can use Rig streaming under feature flags. +- Final history reconciliation is correct and adapter-safe. + +### Scope + +- `src/agent/channel.rs` +- `src/agent/cortex_chat.rs` +- `src/hooks/spacebot.rs` +- `src/messaging/slack.rs` +- `src/messaging/discord.rs` +- `src/messaging/telegram.rs` +- `src/api/state.rs` + +## Phase 3: Worker Read-Only Tool Concurrency + +### Deliverable + +- Eligible worker prompts can run read-only tools concurrently with bounded safety. + +### Scope + +- `src/agent/worker.rs` +- `src/tools.rs` +- relevant tool modules + +## Phase 4: Drift Detection, Docs, And Rollout Hardening + +### Deliverable + +- Validation matrix passes. +- Docs and operational runbooks updated. + +### Scope + +- `README.md` +- `METRICS.md` +- `docs/content/docs/(configuration)/config.mdx` +- `tests/*` + +## Dependency-Aware Task Graph + +| Task ID | Task | Depends On | +|---|---|---| +| T1 | Add rollout config types, defaults, load/runtime plumbing | none | +| T2 | Add parity decision helpers and provider capability gating | T1 | +| T3 | Implement `tool_choice` forwarding for OpenAI-compatible, Responses, and Anthropic | T2 | +| T4 | Implement `output_schema` forwarding for OpenAI-compatible, Responses, and Anthropic | T2 | +| T5 | Add parity metrics, logs, and serializer unit tests | T3, T4 | +| T6 | Add channel streaming execution path behind feature flag | T1, T5 | +| T7 | Add cortex-chat streaming execution path behind feature flag | T1, T5 | +| T8 | Harden adapter streaming behavior and terminal reconciliation | T6, T7 | +| T9 | Define read-only tool allowlist and concurrency config validation | T1 | +| T10 | Implement worker request-level concurrency gating | T9 | +| T11 | Add worker concurrency tests and downgrade-on-violation behavior | T10 | +| T12 | Add drift-reconciliation suite and upgrade checklist | T5, T8, T11 | +| T13 | Update docs, metrics docs, and rollout runbook | T12 | +| T14 | Run staged rollout validation and final gates | T13 | + +## Task Ownership And File Scope + +| Task ID | Owner | File Scope | Merge Conflict Notes | +|---|---|---|---| +| T1 | llm-runtime | `src/config/types.rs`, `src/config/toml_schema.rs`, `src/config/load.rs`, `src/config/runtime.rs` | Keep config-only changes isolated from agent logic | +| T2 | llm-runtime | `src/llm/model.rs`, `src/llm/anthropic/params.rs` | Avoid mixing serializer changes with streaming changes | +| T3 | llm-runtime | `src/llm/model.rs`, `src/llm/anthropic/params.rs` | Same write set as T2/T4; execute serially | +| T4 | llm-runtime | `src/llm/model.rs`, `src/llm/anthropic/params.rs` | Same write set as T2/T3; execute serially | +| T5 | observability + qa | `src/llm/model.rs`, `tests/*`, `METRICS.md` | Keep tests in separate commits if possible | +| T6 | agent-orchestration | `src/agent/channel.rs`, `src/hooks/spacebot.rs` | Hot spot with T8; branch carefully | +| T7 | agent-orchestration | `src/agent/cortex_chat.rs`, `src/hooks/spacebot.rs` | Coordinate with T6 on hook changes | +| T8 | messaging-adapters | `src/messaging/slack.rs`, `src/messaging/discord.rs`, `src/messaging/telegram.rs`, `src/api/state.rs` | Separate by adapter if parallelizing | +| T9 | tools + security | `src/tools.rs`, selected `src/tools/*.rs` | Do not broaden allowlist without review | +| T10 | agent-orchestration | `src/agent/worker.rs` | Isolate from channel changes | +| T11 | qa | `src/agent/worker.rs`, `tests/*` | Can stack on T10 | +| T12 | qa + llm-runtime | `tests/*`, `README.md` | Minimal overlap if docs staged later | +| T13 | qa-docs | `README.md`, `METRICS.md`, `docs/content/docs/(configuration)/config.mdx` | Low conflict if delayed to end | +| T14 | release/qa | no code changes required | Verification-only | + +## Detailed Tasks + +### T1: Add Rollout Config Plumbing + +- Add `RigAlignmentConfig` and `RigSemanticsMode`. +- Expose runtime accessors via `RuntimeConfig`. +- Default all new behavior to existing-safe values: + - `request_semantics_mode = shadow` + - streaming disabled + - concurrency = 1 + +Acceptance criteria: + +- Config parses with and without new block. +- Old configs continue to load unchanged. +- Runtime snapshot surfaces the new config values. + +### T2: Add Request Capability Decision Helpers + +- Add helper functions that evaluate: + - provider/API type + - `tool_choice` present + - `output_schema` present + - rollout mode +- Centralize decision logic so every provider path uses the same rules. + +Acceptance criteria: + +- No serializer path decides independently without using shared helper logic. +- Unsupported combinations produce deterministic decision outcomes. + +### T3: Implement `tool_choice` Forwarding + +- OpenAI chat-completions payloads: + - map Rig `ToolChoice` to provider JSON +- OpenAI Responses payloads: + - map Rig `ToolChoice` to provider JSON +- Anthropic payloads: + - map Rig `ToolChoice` to provider JSON when supported + +Acceptance criteria: + +- `ToolChoice::Required` and `ToolChoice::Specific` are observable in serialized requests where supported. +- Unsupported providers hit shadow or enforced behavior explicitly. + +### T4: Implement `output_schema` Forwarding + +- OpenAI chat-completions: + - map to `response_format` equivalent where supported +- OpenAI Responses: + - map to structured output format +- Anthropic: + - map to `output_config` JSON-schema shape where supported + +Acceptance criteria: + +- `prompt_typed(...)` is contractually meaningful on supported providers. +- Unsupported providers fail closed in enforced mode. + +### T5: Add Parity Metrics And Tests + +- Add serializer tests covering: + - tool choice present + - output schema present + - unsupported provider behavior +- Emit semantic parity metrics and logs. + +Acceptance criteria: + +- Every decision path has a unit test. +- Metrics fire for both applied and rejected cases. + +### T6: Add Channel Streaming Path + +- Add feature-flagged streaming branch in channel execution. +- Reuse existing hook delta events and history reconciliation. +- Preserve retrigger behavior and malformed tool syntax recovery semantics. + +Acceptance criteria: + +- Channel streaming can be toggled off without code-path drift. +- Reply-tool terminal behavior remains authoritative. + +### T7: Add Cortex Chat Streaming Path + +- Replace one-shot prompt path with streaming request under feature flag. +- Continue emitting SSE-compatible events through existing event channel. + +Acceptance criteria: + +- Cortex chat streams partial text and still persists final assistant response. +- Timeout behavior remains explicit and logged. + +### T8: Adapter Streaming Hardening + +- Validate per-adapter edit throttling. +- Ensure terminal final response replaces or finalizes in-flight stream cleanly. +- Ensure failure to edit stream degrades gracefully to final-response delivery. + +Acceptance criteria: + +- No adapter exceeds platform edit cadence. +- Stream failure does not lose final response. + +### T9: Read-Only Tool Allowlist And Validation + +- Define allowlist in config defaults. +- Add validator/helper that rejects concurrency for any non-allowlisted tool. +- Add comments documenting why each allowlisted tool is safe. + +Acceptance criteria: + +- No write-capable tool is concurrency-enabled. +- Test suite fails on allowlist regression. + +### T10: Worker Request-Level Concurrency + +- Apply `with_tool_concurrency(...)` only when: + - worker config permits it + - all exposed tools for the request are allowlisted or the specific request class is concurrency-safe + +Acceptance criteria: + +- Default worker behavior remains sequential. +- Read-only workloads can opt into bounded concurrency. + +### T11: Worker Concurrency Tests And Safety Downgrade + +- Add tests for: + - concurrent read-only tools + - sequential fallback on violation + - metrics/logging on downgrade + +Acceptance criteria: + +- Violation path is deterministic and observable. +- Sequential fallback occurs at most once per request. + +### T12: Drift-Reconciliation Suite + +- Add focused regression suite guarding: + - request semantics parity + - stream history reconciliation + - adapter finalization + - concurrency allowlist +- Add upgrade checklist entry for future Rig version bumps. + +Acceptance criteria: + +- Rig upgrade PRs have a targeted parity suite to run. +- Drift is detected by CI before merge. + +### T13: Documentation And Runbook + +- Update user-facing config docs. +- Update metrics documentation. +- Update README section describing how Spacebot uses Rig and which features are intentionally not adopted. + +Acceptance criteria: + +- Docs match config keys and rollout modes. +- No new behavior is undocumented. + +### T14: Rollout Validation + +- Run staged rollout: + 1. semantics shadow mode + 2. semantics enforced in staging + 3. channel streaming canary + 4. cortex-chat streaming canary + 5. worker concurrency canary + +Acceptance criteria: + +- Each stage has explicit rollback trigger and evidence commands. + +## Test Strategy + +### Unit Tests + +Targets: + +- serializer mapping functions +- capability gating helpers +- allowlist classification +- history reconciliation helpers +- streaming terminal-state helpers + +Logging expectations: + +- each negative-path test asserts at least one structured warning or metric increment + +### Integration Tests + +Targets: + +- channel streaming with tool call then final reply +- cortex-chat streaming SSE events +- worker concurrency with mock read-only tools +- unsupported provider semantics in `shadow` and `enforced` modes + +Logging expectations: + +- `agent_id`, `process_type`, `provider`, and `terminal_reason` appear in emitted traces/logs + +### End-To-End Tests + +Targets: + +- Slack/Discord/Telegram end-to-end streaming edit path with synthetic inbound message fixture +- final response persistence after stream completion +- concurrency-safe worker scenario with multiple read-only tool calls + +Logging expectations: + +- no end-user visible secret leakage in logs +- adapter degrade-to-final path is logged when edit retries fail + +## Validation Matrix + +| Req ID | Requirement | PASS Evidence Command | FAIL Signal | +|---|---|---|---| +| R1 | Config is backward compatible | `cargo test config::load -- --nocapture` | any old-config fixture fails to load | +| R2 | `tool_choice` is forwarded on supported providers | `cargo test forwards_tool_choice -- --nocapture` | serialized payload missing provider-specific tool choice | +| R3 | `output_schema` is forwarded on supported providers | `cargo test forwards_output_schema -- --nocapture` | serialized payload missing schema format field | +| R4 | Unsupported semantics are explicit in enforced mode | `cargo test rejects_unsupported_semantics_enforced -- --nocapture` | request succeeds silently | +| R5 | Channel streaming preserves terminal reply semantics | `cargo test channel_streaming_reply_terminal -- --nocapture` | reply delivered but loop continues or history corrupts | +| R6 | Cortex chat streaming emits partials and persists final response | `cargo test cortex_chat_streaming_persists_final -- --nocapture` | SSE partials absent or final message not stored | +| R7 | Adapter streaming obeys edit cadence | `cargo test adapter_streaming_throttle -- --nocapture` | edit timestamps violate throttle assertions | +| R8 | Worker concurrency only applies to allowlisted tools | `cargo test worker_concurrency_allowlist -- --nocapture` | non-allowlisted tool runs concurrently | +| R9 | Worker concurrency improves eligible workloads | `cargo test worker_concurrency_reduces_wall_clock -- --nocapture` | no measurable improvement or regression beyond threshold | +| R10 | Drift suite guards custom Rig bridge | `cargo test rig_parity -- --nocapture` | any parity regression test fails | +| R11 | Docs and gates are clean | `just preflight && just gate-pr` | any gate red | + +## Rollout Plan + +### Feature Flags + +- `request_semantics_mode` +- `channel_streaming` +- `cortex_chat_streaming` +- `worker_read_only_tool_concurrency` + +### Rollout Sequence + +1. Deploy config plumbing only, all runtime behavior effectively unchanged. +2. Enable `request_semantics_mode = shadow` in staging. +3. Review metrics for unsupported combinations and payload changes. +4. Enable `request_semantics_mode = enforced` in staging. +5. Canary `channel_streaming` for one adapter and one provider family with native streaming. +6. Expand channel streaming to remaining adapters if metrics hold. +7. Enable `cortex_chat_streaming`. +8. Enable worker read-only concurrency for one worker type or task class. + +### Rollback + +- Rollback mechanism is config-first: + - set `request_semantics_mode = off` + - set streaming flags to `false` + - set concurrency back to `1` +- No database rollback expected if implementation remains config-only. +- If persisted schema is later introduced, rollback must use a new forward migration or a runtime compatibility gate, never edit historical migration files. + +### Backwards Compatibility + +- Old config files remain valid. +- Old providers without structured output support remain operable in `shadow` or `off`. +- Existing non-streaming behavior remains default until flags are enabled. + +## Operational Risks And Mitigations + +| Risk | Impact | Mitigation | +|---|---|---| +| Provider rejects new `tool_choice` or schema field | request failures | shadow mode first, provider allowlist, serializer tests | +| Streaming exposes secret-like partial text | privacy leak | scrub deltas before user-visible delivery | +| History reconciliation diverges between streaming and non-streaming | broken transcript state | shared terminal reconciliation helper and parity tests | +| Concurrency reorders side effects | task corruption | strict allowlist and downgrade-to-sequential | +| Synthetic streams do not improve UX | wasted complexity | measure native vs synthetic separately and keep synthetic rollout optional | +| Rig upgrade changes request shape | silent drift | explicit parity suite and upgrade checklist | + +## Commands For Implementation And Verification + +### Targeted Verification Commands + +```bash +cargo test forwards_tool_choice -- --nocapture +cargo test forwards_output_schema -- --nocapture +cargo test rejects_unsupported_semantics_enforced -- --nocapture +cargo test channel_streaming_reply_terminal -- --nocapture +cargo test cortex_chat_streaming_persists_final -- --nocapture +cargo test adapter_streaming_throttle -- --nocapture +cargo test worker_concurrency_allowlist -- --nocapture +cargo test worker_concurrency_reduces_wall_clock -- --nocapture +cargo test rig_parity -- --nocapture +just preflight +just gate-pr +``` + +### Logging Review Commands + +```bash +rg "request_semantics_mode|output_schema_requested|tool_choice_requested|stream_mode|tool_concurrency" target debug.log logs -g '*.log' +``` + +## Definition Of Done + +- [x] `RigAlignmentConfig` exists, parses, hot-reloads, and defaults to backward-compatible behavior. +- [x] `SpacebotModel` explicitly handles `tool_choice` for supported providers and explicitly rejects or shadows unsupported cases. +- [x] `SpacebotModel` explicitly handles `output_schema` for supported providers and explicitly rejects or shadows unsupported cases. +- [x] `prompt_typed(...)` is contractually meaningful on supported providers and covered by tests. +- [x] Channel streaming is feature-gated, adapter-safe, and preserves terminal reply semantics. +- [x] Cortex-chat streaming is feature-gated and persists final assistant output after streamed partials. +- [x] Worker request-level tool concurrency is bounded, allowlisted, and disabled by default. +- [x] No new path silently drops a semantic request, retry exhaustion, or adapter streaming failure. +- [x] Metrics and logs exist for semantic parity, streaming mode, time to first delta, and concurrency usage. +- [x] Drift-reconciliation tests fail when custom Rig bridge behavior diverges from expected request semantics. +- [x] User-facing config and metrics docs are updated in existing documentation files. +- [x] `just preflight` passes. +- [x] `just gate-pr` passes. +- [x] All targeted validation commands in this document pass with non-zero failure on regression. diff --git a/src/agent/branch.rs b/src/agent/branch.rs index 001eab7ea..76ca16123 100644 --- a/src/agent/branch.rs +++ b/src/agent/branch.rs @@ -92,6 +92,7 @@ impl Branch { let model_name = routing.resolve(ProcessType::Branch, None).to_string(); let model = SpacebotModel::make(&self.deps.llm_manager, &model_name) .with_context(&*self.deps.agent_id, "branch") + .with_rig_alignment((**self.deps.runtime_config.rig_alignment.load()).clone()) .with_routing((**routing).clone()); let agent = AgentBuilder::new(model) diff --git a/src/agent/channel.rs b/src/agent/channel.rs index 113e691f1..e2580d0f3 100644 --- a/src/agent/channel.rs +++ b/src/agent/channel.rs @@ -22,14 +22,17 @@ use crate::{ AgentDeps, BranchId, ChannelId, InboundMessage, OutboundResponse, ProcessEvent, ProcessId, ProcessType, WorkerId, }; -use rig::agent::AgentBuilder; -use rig::completion::CompletionModel; +use futures::StreamExt as _; +use rig::agent::MultiTurnStreamItem; +use rig::agent::{AgentBuilder, HookAction, PromptHook, ToolCallHookAction}; +use rig::completion::{CompletionModel, GetTokenUsage, PromptError}; use rig::message::UserContent; use rig::one_or_many::OneOrMany; +use rig::streaming::StreamingPrompt; use rig::tool::server::ToolServer; use std::collections::HashMap; use std::collections::HashSet; -use std::sync::{Arc, Weak}; +use std::sync::{Arc, Mutex, Weak}; use tokio::sync::broadcast; use tokio::sync::{RwLock, mpsc}; @@ -74,6 +77,306 @@ fn should_flush_coalesce_buffer_for_event(event: &ProcessEvent) -> bool { ) } +fn streaming_error_to_prompt_result( + error: rig::agent::StreamingError, + fallback_history: Vec, +) -> ( + std::result::Result, + Vec, +) { + match error { + rig::agent::StreamingError::Prompt(error) => { + let history = match error.as_ref() { + PromptError::MaxTurnsError { chat_history, .. } + | PromptError::PromptCancelled { chat_history, .. } => (**chat_history).clone(), + _ => fallback_history, + }; + (Err(*error), history) + } + rig::agent::StreamingError::Completion(error) => { + (Err(PromptError::CompletionError(error)), fallback_history) + } + rig::agent::StreamingError::Tool(error) => { + (Err(PromptError::ToolError(error)), fallback_history) + } + } +} + +fn missing_stream_final_response_error() -> PromptError { + PromptError::CompletionError(rig::completion::CompletionError::ProviderError( + "stream ended without final response".to_string(), + )) +} + +const STREAM_BLOCK_PREFIX_BYTES: usize = 64; +const STREAM_LEAK_TAIL_BYTES: usize = 256; + +fn push_prefix_bytes(buffer: &mut String, text: &str, max_bytes: usize) { + if buffer.len() >= max_bytes || text.is_empty() { + return; + } + + let remaining = max_bytes - buffer.len(); + if text.len() <= remaining { + buffer.push_str(text); + return; + } + + let mut end = remaining; + while end > 0 && !text.is_char_boundary(end) { + end -= 1; + } + buffer.push_str(&text[..end]); +} + +fn trim_to_last_bytes(buffer: &mut String, max_bytes: usize) { + if buffer.len() <= max_bytes { + return; + } + + let mut start = buffer.len() - max_bytes; + while start < buffer.len() && !buffer.is_char_boundary(start) { + start += 1; + } + buffer.drain(..start); +} + +#[derive(Default)] +struct StreamForwardingState { + leading_text: String, + leak_tail: String, + blocked: bool, +} + +impl StreamForwardingState { + fn should_forward_delta(&mut self, text_delta: &str) -> bool { + if text_delta.is_empty() || self.blocked { + return false; + } + + push_prefix_bytes( + &mut self.leading_text, + text_delta, + STREAM_BLOCK_PREFIX_BYTES, + ); + if crate::tools::should_block_user_visible_text(&self.leading_text) { + self.blocked = true; + return false; + } + + self.leak_tail.push_str(text_delta); + trim_to_last_bytes(&mut self.leak_tail, STREAM_LEAK_TAIL_BYTES); + + if let Some(leak) = crate::secrets::scrub::scan_for_leaks(&self.leak_tail) { + tracing::warn!( + leak_prefix = %&leak[..leak.len().min(8)], + "blocked streamed output matching secret pattern" + ); + self.blocked = true; + return false; + } + + true + } +} + +fn should_forward_stream_delta(state: &mut StreamForwardingState, text_delta: &str) -> bool { + if text_delta.is_empty() { + return false; + } + + state.should_forward_delta(text_delta) +} + +#[derive(Clone)] +struct ChannelStreamingHook { + spacebot_hook: SpacebotHook, + response_tx: mpsc::Sender, + forwarding_state: Arc>, +} + +impl ChannelStreamingHook { + fn new(spacebot_hook: SpacebotHook, response_tx: mpsc::Sender) -> Self { + Self { + spacebot_hook, + response_tx, + forwarding_state: Arc::new(Mutex::new(StreamForwardingState::default())), + } + } +} + +impl PromptHook for ChannelStreamingHook +where + M: CompletionModel, +{ + async fn on_completion_call( + &self, + prompt: &rig::message::Message, + history: &[rig::message::Message], + ) -> HookAction { + >::on_completion_call(&self.spacebot_hook, prompt, history) + .await + } + + async fn on_completion_response( + &self, + prompt: &rig::message::Message, + response: &rig::completion::CompletionResponse, + ) -> HookAction { + >::on_completion_response( + &self.spacebot_hook, + prompt, + response, + ) + .await + } + + async fn on_tool_call( + &self, + tool_name: &str, + tool_call_id: Option, + internal_call_id: &str, + args: &str, + ) -> ToolCallHookAction { + >::on_tool_call( + &self.spacebot_hook, + tool_name, + tool_call_id, + internal_call_id, + args, + ) + .await + } + + async fn on_tool_result( + &self, + tool_name: &str, + tool_call_id: Option, + internal_call_id: &str, + args: &str, + result: &str, + ) -> HookAction { + >::on_tool_result( + &self.spacebot_hook, + tool_name, + tool_call_id, + internal_call_id, + args, + result, + ) + .await + } + + async fn on_text_delta(&self, text_delta: &str, aggregated_text: &str) -> HookAction { + let action = >::on_text_delta( + &self.spacebot_hook, + text_delta, + aggregated_text, + ) + .await; + if !matches!(action, HookAction::Continue) { + return action; + } + + let should_forward = { + let mut state = self.forwarding_state.lock().unwrap_or_else(|poisoned| { + tracing::warn!("channel streaming forwarding state poisoned; recovering"); + poisoned.into_inner() + }); + should_forward_stream_delta(&mut state, text_delta) + }; + + if should_forward { + match self + .response_tx + .try_send(OutboundResponse::StreamChunk(text_delta.to_string())) + { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + tracing::debug!( + "dropping streamed text delta because outbound response channel is full" + ); + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + tracing::debug!( + "dropping streamed text delta because outbound response channel is closed" + ); + } + } + } + + HookAction::Continue + } + + async fn on_tool_call_delta( + &self, + tool_call_id: &str, + internal_call_id: &str, + tool_name: Option<&str>, + tool_call_delta: &str, + ) -> HookAction { + >::on_tool_call_delta( + &self.spacebot_hook, + tool_call_id, + internal_call_id, + tool_name, + tool_call_delta, + ) + .await + } + + async fn on_stream_completion_response_finish( + &self, + prompt: &rig::message::Message, + response: &::StreamingResponse, + ) -> HookAction { + >::on_stream_completion_response_finish( + &self.spacebot_hook, + prompt, + response, + ) + .await + } +} + +async fn prompt_once_streaming( + agent: &rig::agent::Agent, + hook: P, + history: Vec, + prompt: &str, +) -> ( + std::result::Result, + Vec, +) +where + M: CompletionModel + 'static, + M::StreamingResponse: GetTokenUsage, + P: rig::agent::PromptHook + 'static, +{ + let fallback_history = history.clone(); + let mut stream = agent + .stream_prompt(prompt) + .with_hook(hook) + .with_history(history) + .await; + + let mut final_history = fallback_history; + while let Some(item) = stream.next().await { + match item { + Ok(MultiTurnStreamItem::FinalResponse(final_response)) => { + if let Some(history) = final_response.history() { + final_history = history.to_vec(); + } + return (Ok(final_response.response().to_string()), final_history); + } + Ok(_) => {} + Err(error) => return streaming_error_to_prompt_result(error, final_history), + } + } + + (Err(missing_stream_final_response_error()), final_history) +} + /// Shared state that channel tools need to act on the channel. /// /// Wrapped in Arc and passed to tools (branch, spawn_worker, route, cancel) @@ -1805,6 +2108,7 @@ impl Channel { let rc = &self.deps.runtime_config; let routing = rc.routing.load(); + let rig_alignment = rc.rig_alignment.load(); let max_turns = if is_retrigger { RETRIGGER_MAX_TURNS } else { @@ -1813,6 +2117,7 @@ impl Channel { let model_name = routing.resolve(ProcessType::Channel, None); let model = SpacebotModel::make(&self.deps.llm_manager, model_name) .with_context(&*self.deps.agent_id, "channel") + .with_rig_alignment((**rig_alignment).clone()) .with_routing((**routing).clone()); let agent = AgentBuilder::new(model) @@ -1868,7 +2173,15 @@ impl Channel { }; let history_len_before = history.len(); - let mut result = self.hook.prompt_once(&agent, &mut history, user_text).await; + let streaming_enabled = rig_alignment.channel_streaming; + let (mut result, mut history) = if streaming_enabled { + let streaming_hook = + ChannelStreamingHook::new(self.hook.clone(), self.response_tx.clone()); + prompt_once_streaming(&agent, streaming_hook, history, user_text).await + } else { + let result = self.hook.prompt_once(&agent, &mut history, user_text).await; + (result, history) + }; // If the LLM responded with text that looks like tool call syntax, it failed // to use the tool calling API. Inject a correction and retry a couple @@ -1891,10 +2204,19 @@ impl Channel { let prompt_engine = self.deps.runtime_config.prompts.load(); let correction = prompt_engine.render_system_tool_syntax_correction()?; - result = self - .hook - .prompt_once(&agent, &mut history, &correction) - .await; + if streaming_enabled { + let streaming_hook = + ChannelStreamingHook::new(self.hook.clone(), self.response_tx.clone()); + let (next_result, next_history) = + prompt_once_streaming(&agent, streaming_hook, history, &correction).await; + result = next_result; + history = next_history; + } else { + result = self + .hook + .prompt_once(&agent, &mut history, &correction) + .await; + } } let retrigger_reply_preserved = { @@ -2589,9 +2911,19 @@ impl Channel { #[cfg(test)] mod tests { - use super::{recv_channel_event, should_process_event_for_channel}; + use super::{ + ChannelStreamingHook, StreamForwardingState, missing_stream_final_response_error, + recv_channel_event, should_forward_stream_delta, should_process_event_for_channel, + streaming_error_to_prompt_result, + }; + use crate::hooks::SpacebotHook; + use crate::llm::SpacebotModel; use crate::memory::MemoryType; - use crate::{AgentId, ChannelId, ProcessEvent, ProcessId}; + use crate::{AgentId, ChannelId, ProcessEvent, ProcessId, ProcessType}; + use rig::agent::{HookAction, PromptHook}; + use rig::completion::PromptError; + use rig::message::{Message, UserContent}; + use rig::one_or_many::OneOrMany; use std::sync::Arc; #[tokio::test] @@ -2689,6 +3021,124 @@ mod tests { assert!(should_process_event_for_channel(&event, &channel_id)); } + #[test] + fn channel_streaming_reply_terminal() { + let fallback_history = vec![Message::User { + content: OneOrMany::one(UserContent::text("fallback-history")), + }]; + let terminal_history = vec![ + Message::User { + content: OneOrMany::one(UserContent::text("real-user-prompt")), + }, + Message::Assistant { + id: None, + content: OneOrMany::one(rig::message::AssistantContent::tool_call( + "call_1", + "reply", + serde_json::json!({ "content": "terminal reply" }), + )), + }, + ]; + + let streaming_error = + rig::agent::StreamingError::Prompt(Box::new(PromptError::PromptCancelled { + chat_history: Box::new(terminal_history.clone()), + reason: "reply delivered".to_string(), + })); + + let (result, history) = streaming_error_to_prompt_result(streaming_error, fallback_history); + + assert!( + matches!( + result, + Err(PromptError::PromptCancelled { reason, .. }) if reason == "reply delivered" + ), + "streaming terminal reply should stay terminal and preserve PromptCancelled semantics" + ); + assert_eq!( + history, terminal_history, + "streaming terminal reply should use terminal chat history, not fallback history" + ); + } + + #[test] + fn missing_stream_final_response_is_an_error() { + let error = missing_stream_final_response_error(); + assert!(matches!( + error, + rig::completion::PromptError::CompletionError( + rig::completion::CompletionError::ProviderError(_) + ) + )); + assert!( + error + .to_string() + .contains("stream ended without final response") + ); + } + + #[test] + fn streamed_deltas_block_secrets_and_structured_tool_syntax() { + let mut structured = StreamForwardingState::default(); + assert!(!should_forward_stream_delta( + &mut structured, + "{\"content\":\"hello\"}" + )); + + let mut leak = StreamForwardingState::default(); + assert!(should_forward_stream_delta( + &mut leak, + "token sk-ant-abc1234567890" + )); + assert!(!should_forward_stream_delta(&mut leak, "1234567890")); + + let mut normal = StreamForwardingState::default(); + assert!(should_forward_stream_delta( + &mut normal, + "normal assistant " + )); + assert!(should_forward_stream_delta(&mut normal, "reply")); + } + + #[tokio::test] + async fn channel_streaming_hook_drops_when_response_channel_is_full() { + let (response_tx, mut response_rx) = tokio::sync::mpsc::channel(1); + response_tx + .send(crate::OutboundResponse::Text("already queued".to_string())) + .await + .expect("prefill response queue"); + + let (event_tx, _event_rx) = tokio::sync::broadcast::channel(8); + let hook = ChannelStreamingHook::new( + SpacebotHook::new( + Arc::::from("agent"), + ProcessId::Channel(Arc::::from("channel")), + ProcessType::Channel, + Some(Arc::::from("channel")), + event_tx, + ), + response_tx, + ); + + let action = >::on_text_delta( + &hook, + "partial", + "partial response", + ) + .await; + assert!(matches!(action, HookAction::Continue)); + + let queued = response_rx + .recv() + .await + .expect("prefilled message should remain"); + assert!(matches!(queued, crate::OutboundResponse::Text(_))); + assert!( + response_rx.try_recv().is_err(), + "streaming hook should drop deltas instead of blocking when the outbound queue is full" + ); + } + #[test] fn worker_complete_event_matches_own_channel() { let channel_id: ChannelId = Arc::from("channel-a"); diff --git a/src/agent/compactor.rs b/src/agent/compactor.rs index 04567da89..7bdb176a4 100644 --- a/src/agent/compactor.rs +++ b/src/agent/compactor.rs @@ -224,6 +224,7 @@ async fn run_compaction( let model_name = routing.resolve(ProcessType::Compactor, None).to_string(); let model = SpacebotModel::make(&deps.llm_manager, &model_name) .with_context(&*deps.agent_id, "compactor") + .with_rig_alignment((**deps.runtime_config.rig_alignment.load()).clone()) .with_routing((**routing).clone()); // Give the compaction worker memory_save so it can directly persist memories diff --git a/src/agent/cortex.rs b/src/agent/cortex.rs index 62c593ef3..948409c07 100644 --- a/src/agent/cortex.rs +++ b/src/agent/cortex.rs @@ -2066,6 +2066,7 @@ pub async fn generate_bulletin(deps: &AgentDeps, logger: &CortexLogger) -> bool let model_name = routing.resolve(ProcessType::Cortex, None).to_string(); let model = SpacebotModel::make(&deps.llm_manager, &model_name) .with_context(&*deps.agent_id, "cortex") + .with_rig_alignment((**deps.runtime_config.rig_alignment.load()).clone()) .with_routing((**routing).clone()); // No tools needed — the LLM just synthesizes the pre-gathered data. @@ -2251,6 +2252,7 @@ async fn generate_profile(deps: &AgentDeps, logger: &CortexLogger) { let model_name = routing.resolve(ProcessType::Cortex, None).to_string(); let model = SpacebotModel::make(&deps.llm_manager, &model_name) .with_context(&*deps.agent_id, "cortex") + .with_rig_alignment((**deps.runtime_config.rig_alignment.load()).clone()) .with_routing((**routing).clone()); let agent = AgentBuilder::new(model) diff --git a/src/agent/cortex_chat.rs b/src/agent/cortex_chat.rs index 575bd85cd..de028e214 100644 --- a/src/agent/cortex_chat.rs +++ b/src/agent/cortex_chat.rs @@ -10,8 +10,13 @@ use crate::hooks::SpacebotHook; use crate::llm::SpacebotModel; use crate::{AgentDeps, ProcessId, ProcessType}; +use futures::StreamExt as _; +use rig::agent::MultiTurnStreamItem; use rig::agent::{AgentBuilder, HookAction, PromptHook, ToolCallHookAction}; -use rig::completion::{AssistantContent, CompletionModel, CompletionResponse, Message, Prompt}; +use rig::completion::{ + AssistantContent, CompletionModel, CompletionResponse, GetTokenUsage, Message, Prompt, +}; +use rig::streaming::StreamingPrompt; use rig::tool::server::ToolServerHandle; use serde::Serialize; use sqlx::SqlitePool; @@ -38,6 +43,8 @@ pub struct CortexChatMessage { pub enum CortexChatEvent { /// The cortex is processing (before LLM response). Thinking, + /// A partial text delta arrived from the model stream. + TextDelta { text_delta: String }, /// A tool call started. ToolStarted { tool: String }, /// A tool call completed. @@ -79,6 +86,14 @@ impl CortexChatHook { async fn send(&self, event: CortexChatEvent) { let _ = self.event_tx.send(event).await; } + + fn try_send(&self, event: CortexChatEvent) { + if let Err(error) = self.event_tx.try_send(event) + && matches!(error, tokio::sync::mpsc::error::TrySendError::Full(_)) + { + tracing::debug!("dropping cortex chat streaming event because the SSE channel is full"); + } + } } fn try_acquire_send_lock( @@ -97,13 +112,106 @@ async fn persist_and_emit_cortex_chat_error( channel_ref: Option<&str>, message: String, ) { - let _ = store + if let Err(error) = store .save_message(thread_id, "assistant", &message, channel_ref) + .await + { + tracing::error!(%error, %thread_id, "failed to persist cortex chat error message"); + } + if let Err(error) = event_tx.send(CortexChatEvent::Error { message }).await { + tracing::debug!(%error, %thread_id, "failed to emit cortex chat error event"); + } +} + +async fn persist_and_emit_cortex_chat_done( + store: &CortexChatStore, + event_tx: &mpsc::Sender, + thread_id: &str, + channel_ref: Option<&str>, + response: String, +) { + if let Err(error) = store + .save_message(thread_id, "assistant", &response, channel_ref) + .await + { + tracing::error!(%error, %thread_id, "failed to persist cortex chat response"); + } + if let Err(error) = event_tx + .send(CortexChatEvent::Done { + full_text: response, + }) + .await + { + tracing::debug!(%error, %thread_id, "failed to emit cortex chat done event"); + } +} + +fn streaming_error_to_prompt_error( + error: rig::agent::StreamingError, +) -> rig::completion::PromptError { + match error { + rig::agent::StreamingError::Prompt(error) => *error, + rig::agent::StreamingError::Completion(error) => { + rig::completion::PromptError::CompletionError(error) + } + rig::agent::StreamingError::Tool(error) => rig::completion::PromptError::ToolError(error), + } +} + +fn missing_stream_final_response_error() -> rig::completion::PromptError { + rig::completion::PromptError::CompletionError(rig::completion::CompletionError::ProviderError( + "stream ended without final response".to_string(), + )) +} + +async fn prompt_once_streaming( + agent: &rig::agent::Agent, + hook: P, + history: Vec, + prompt: &str, +) -> std::result::Result +where + M: CompletionModel + 'static, + M::StreamingResponse: GetTokenUsage, + P: rig::agent::PromptHook + 'static, +{ + let mut stream = agent + .stream_prompt(prompt) + .with_hook(hook) + .with_history(history) .await; - let _ = event_tx.send(CortexChatEvent::Error { message }).await; + + while let Some(item) = stream.next().await { + match item { + Ok(MultiTurnStreamItem::FinalResponse(final_response)) => { + return Ok(final_response.response().to_string()); + } + Ok(_) => {} + Err(error) => return Err(streaming_error_to_prompt_error(error)), + } + } + + Err(missing_stream_final_response_error()) } impl PromptHook for CortexChatHook { + async fn on_text_delta(&self, text_delta: &str, aggregated_text: &str) -> HookAction { + let action = >::on_text_delta( + &self.spacebot_hook, + text_delta, + aggregated_text, + ) + .await; + if !matches!(action, HookAction::Continue) { + return action; + } + + self.try_send(CortexChatEvent::TextDelta { + text_delta: text_delta.to_string(), + }); + HookAction::Continue + } + async fn on_tool_call( &self, tool_name: &str, @@ -326,6 +434,7 @@ impl CortexChatSession { let model_name = routing.resolve(ProcessType::Cortex, None).to_string(); let model = SpacebotModel::make(&self.deps.llm_manager, &model_name) .with_context(self.deps.agent_id.as_ref(), "cortex") + .with_rig_alignment((**self.deps.runtime_config.rig_alignment.load()).clone()) .with_routing(routing.as_ref().clone()); let agent = AgentBuilder::new(model) @@ -357,29 +466,39 @@ impl CortexChatSession { .branch_timeout_secs .max(1), ); + let streaming_enabled = self + .deps + .runtime_config + .rig_alignment + .load() + .cortex_chat_streaming; tokio::spawn(async move { let _send_guard = send_guard; let channel_ref = channel_context_id.as_deref(); - let prompt_result = tokio::time::timeout( - prompt_timeout, - agent - .prompt(&user_text) - .with_hook(hook.clone()) - .with_history(&mut history), - ) + let prompt_result = tokio::time::timeout(prompt_timeout, async move { + if streaming_enabled { + prompt_once_streaming(&agent, hook.clone(), history, &user_text).await + } else { + agent + .prompt(&user_text) + .with_hook(hook.clone()) + .with_history(&mut history) + .await + } + }) .await; match prompt_result { Ok(Ok(response)) => { - let _ = store - .save_message(&thread_id, "assistant", &response, channel_ref) - .await; - let _ = event_tx - .send(CortexChatEvent::Done { - full_text: response, - }) - .await; + persist_and_emit_cortex_chat_done( + &store, + &event_tx, + &thread_id, + channel_ref, + response, + ) + .await; } Ok(Err(error)) => { let error_text = format!("Cortex chat error: {error}"); @@ -514,7 +633,16 @@ impl CortexChatSession { #[cfg(test)] mod tests { - use super::{CortexChatSendError, try_acquire_send_lock}; + use super::{ + CortexChatEvent, CortexChatHook, CortexChatSendError, CortexChatStore, + missing_stream_final_response_error, persist_and_emit_cortex_chat_done, + try_acquire_send_lock, + }; + use crate::hooks::SpacebotHook; + use crate::llm::SpacebotModel; + use crate::{ProcessId, ProcessType}; + use rig::agent::{HookAction, PromptHook}; + use sqlx::SqlitePool; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; @@ -562,4 +690,143 @@ mod tests { "single-flight lock should be released after timeout path" ); } + + #[tokio::test] + async fn cortex_chat_streaming_persists_final() { + let pool = SqlitePool::connect("sqlite::memory:") + .await + .expect("in-memory sqlite should connect"); + sqlx::query( + "CREATE TABLE cortex_chat_messages ( \ + id TEXT PRIMARY KEY, \ + thread_id TEXT NOT NULL, \ + role TEXT NOT NULL, \ + content TEXT NOT NULL, \ + channel_context TEXT, \ + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP \ + )", + ) + .execute(&pool) + .await + .expect("cortex chat table should be created"); + + let store = CortexChatStore::new(pool); + let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(8); + let (process_event_tx, _process_event_rx) = tokio::sync::broadcast::channel(8); + + let spacebot_hook = SpacebotHook::new( + Arc::from("agent"), + ProcessId::Worker(uuid::Uuid::new_v4()), + ProcessType::Cortex, + None, + process_event_tx, + ); + let hook = CortexChatHook::new(event_tx.clone(), spacebot_hook); + + let hook_action = >::on_text_delta( + &hook, + "partial ", + "partial response", + ) + .await; + assert!( + matches!(hook_action, HookAction::Continue), + "streaming text deltas should continue" + ); + + let final_response = "final response".to_string(); + persist_and_emit_cortex_chat_done(&store, &event_tx, "thread-1", None, final_response) + .await; + + let first_event = tokio::time::timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("should receive first event") + .expect("first event should exist"); + match first_event { + CortexChatEvent::TextDelta { text_delta } => { + assert_eq!(text_delta, "partial "); + } + other => panic!("expected TextDelta event, got {other:?}"), + } + + let second_event = tokio::time::timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("should receive done event") + .expect("done event should exist"); + match second_event { + CortexChatEvent::Done { full_text } => { + assert_eq!(full_text, "final response"); + } + other => panic!("expected Done event, got {other:?}"), + } + + let history = store + .load_history("thread-1", 10) + .await + .expect("history should load"); + assert_eq!( + history.len(), + 1, + "assistant final response should be persisted" + ); + assert_eq!(history[0].role, "assistant"); + assert_eq!(history[0].content, "final response"); + } + + #[tokio::test] + async fn text_delta_backpressure_does_not_timeout_prompt() { + let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(1); + event_tx + .send(CortexChatEvent::Thinking) + .await + .expect("prefill event channel"); + let (process_event_tx, _process_event_rx) = tokio::sync::broadcast::channel(8); + + let spacebot_hook = SpacebotHook::new( + Arc::from("agent"), + ProcessId::Worker(uuid::Uuid::new_v4()), + ProcessType::Cortex, + None, + process_event_tx, + ); + let hook = CortexChatHook::new(event_tx, spacebot_hook); + + let action = tokio::time::timeout( + Duration::from_millis(20), + >::on_text_delta( + &hook, + "partial ", + "partial response", + ), + ) + .await + .expect("text delta should not block on a full SSE channel"); + + assert!(matches!(action, HookAction::Continue)); + let first_event = event_rx + .recv() + .await + .expect("prefilled event should remain"); + assert!(matches!(first_event, CortexChatEvent::Thinking)); + assert!( + event_rx.try_recv().is_err(), + "full SSE channel should drop text deltas instead of backpressuring the prompt" + ); + } + + #[test] + fn missing_stream_final_response_is_an_error() { + let error = missing_stream_final_response_error(); + assert!(matches!( + error, + rig::completion::PromptError::CompletionError( + rig::completion::CompletionError::ProviderError(_) + ) + )); + assert!( + error + .to_string() + .contains("stream ended without final response") + ); + } } diff --git a/src/agent/ingestion.rs b/src/agent/ingestion.rs index 9ffb7eb78..1f5fb6509 100644 --- a/src/agent/ingestion.rs +++ b/src/agent/ingestion.rs @@ -469,6 +469,7 @@ async fn process_chunk( let model_name = routing.resolve(ProcessType::Branch, None).to_string(); let model = SpacebotModel::make(&deps.llm_manager, &model_name) .with_context(&*deps.agent_id, "branch") + .with_rig_alignment((**deps.runtime_config.rig_alignment.load()).clone()) .with_routing((**routing).clone()); let conversation_logger = diff --git a/src/agent/worker.rs b/src/agent/worker.rs index b0f40706d..8ddfd2f31 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -1,7 +1,7 @@ //! Worker: Independent task execution process. use crate::agent::compactor::estimate_history_tokens; -use crate::config::BrowserConfig; +use crate::config::{BrowserConfig, RigAlignmentConfig}; use crate::error::Result; use crate::hooks::{SpacebotHook, ToolNudgePolicy}; use crate::llm::SpacebotModel; @@ -9,11 +9,14 @@ use crate::llm::routing::is_context_overflow_error; use crate::{AgentDeps, ChannelId, ProcessId, ProcessType, WorkerId}; use rig::agent::AgentBuilder; use rig::completion::CompletionModel; +use std::collections::HashSet; use std::fmt::Write as _; use std::path::PathBuf; use tokio::sync::{mpsc, watch}; use uuid::Uuid; +static WORKER_CONCURRENCY_SURFACE_WARNING: std::sync::Once = std::sync::Once::new(); + /// How many turns per segment before we check context and potentially compact. const TURNS_PER_SEGMENT: usize = 25; @@ -251,6 +254,31 @@ impl Worker { tracing::info!(worker_id = %self.id, task = %self.task, "worker starting"); let mcp_tools = self.deps.mcp_manager.get_tools().await; + let worker_tool_names = crate::tools::list_worker_tool_names( + self.browser_config.enabled, + self.brave_search_key.is_some(), + self.deps.runtime_config.secrets.load().as_ref().is_some(), + &mcp_tools, + ); + let tool_concurrency = resolve_worker_tool_concurrency( + self.deps.runtime_config.rig_alignment.load().as_ref(), + &worker_tool_names, + ); + tracing::info!( + worker_id = %self.id, + worker_type = "builtin", + tool_concurrency, + "worker tool concurrency resolved" + ); + #[cfg(feature = "metrics")] + { + let metrics = crate::telemetry::Metrics::global(); + let concurrency_label = tool_concurrency.to_string(); + metrics + .rig_tool_concurrency_total + .with_label_values(&["builtin", &concurrency_label]) + .inc(); + } // Create per-worker ToolServer with task tools let worker_tool_server = crate::tools::create_worker_tool_server( @@ -272,6 +300,7 @@ impl Worker { let model_name = routing.resolve(ProcessType::Worker, None).to_string(); let model = SpacebotModel::make(&self.deps.llm_manager, &model_name) .with_context(&*self.deps.agent_id, "worker") + .with_rig_alignment((**self.deps.runtime_config.rig_alignment.load()).clone()) .with_routing((**routing).clone()); let agent = AgentBuilder::new(model) @@ -313,7 +342,12 @@ impl Worker { match self .hook - .prompt_with_tool_nudge_retry(&agent, &mut history, &prompt) + .prompt_with_tool_nudge_retry_with_concurrency( + &agent, + &mut history, + &prompt, + tool_concurrency, + ) .await { Ok(response) => { @@ -424,7 +458,12 @@ impl Worker { let follow_up_result: std::result::Result = loop { match follow_up_hook - .prompt_once(&agent, &mut history, &follow_up_prompt) + .prompt_once_with_tool_concurrency( + &agent, + &mut history, + &follow_up_prompt, + tool_concurrency, + ) .await { Ok(response) => break Ok(response), @@ -880,3 +919,353 @@ fn build_worker_recap(messages: &[rig::message::Message]) -> String { recap } } + +fn resolve_worker_tool_concurrency( + rig_alignment: &RigAlignmentConfig, + worker_tool_names: &[String], +) -> usize { + let max_tool_concurrency = rig_alignment.worker_max_tool_concurrency.max(1); + let configured = rig_alignment + .worker_read_only_tool_concurrency + .max(1) + .min(max_tool_concurrency); + + if configured <= 1 { + return 1; + } + + let allowlist: HashSet<&str> = rig_alignment + .worker_read_only_tool_allowlist + .iter() + .map(String::as_str) + .collect(); + let unsafe_allowlist_entries: Vec<&str> = rig_alignment + .worker_read_only_tool_allowlist + .iter() + .map(String::as_str) + .filter(|tool_name| !is_read_only_worker_tool(tool_name)) + .collect(); + + if !unsafe_allowlist_entries.is_empty() { + tracing::warn!( + ?unsafe_allowlist_entries, + "worker read-only concurrency allowlist contains mutable tools; forcing sequential execution" + ); + return 1; + } + + let non_allowlisted_tools: Vec<&str> = worker_tool_names + .iter() + .map(String::as_str) + .filter(|tool_name| !allowlist.contains(tool_name)) + .collect(); + if !non_allowlisted_tools.is_empty() { + WORKER_CONCURRENCY_SURFACE_WARNING.call_once(|| { + tracing::warn!( + configured, + ?non_allowlisted_tools, + "worker read-only concurrency only applies to fully allowlisted read-only tool surfaces; forcing sequential execution" + ); + }); + return 1; + } + + configured +} + +fn is_read_only_worker_tool(tool_name: &str) -> bool { + crate::config::is_builtin_read_only_worker_tool(tool_name) +} + +#[cfg(test)] +mod tests { + use super::{RigAlignmentConfig, resolve_worker_tool_concurrency}; + use crate::hooks::SpacebotHook; + use crate::{ProcessId, ProcessType}; + use rig::OneOrMany; + use rig::agent::AgentBuilder; + use rig::completion::{ + CompletionError, CompletionModel, CompletionRequest, CompletionResponse, ToolDefinition, + Usage, + }; + use rig::message::{AssistantContent, Message, UserContent}; + use rig::streaming::StreamingCompletionResponse; + use rig::tool::Tool; + use rig::tool::server::ToolServer; + use schemars::JsonSchema; + use serde::{Deserialize, Serialize}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::{Duration, Instant}; + + #[test] + fn worker_concurrency_allowlist_rejects_default_worker_surface() { + let tool_names = vec![ + "shell".to_string(), + "file".to_string(), + "exec".to_string(), + "task_update".to_string(), + "set_status".to_string(), + "read_skill".to_string(), + "web_search".to_string(), + ]; + + assert_eq!( + resolve_worker_tool_concurrency(&RigAlignmentConfig::default(), &tool_names), + 1 + ); + } + + #[test] + fn worker_concurrency_allowlist_accepts_read_only_surface() { + let rig_alignment = RigAlignmentConfig { + worker_read_only_tool_concurrency: 3, + ..RigAlignmentConfig::default() + }; + let tool_names = vec!["read_skill".to_string(), "web_search".to_string()]; + + assert_eq!( + resolve_worker_tool_concurrency(&rig_alignment, &tool_names), + 3 + ); + } + + #[test] + fn worker_concurrency_allowlist_rejects_unknown_tool_names() { + let mut rig_alignment = RigAlignmentConfig { + worker_read_only_tool_concurrency: 3, + ..RigAlignmentConfig::default() + }; + rig_alignment + .worker_read_only_tool_allowlist + .push("custom_mutating_tool".to_string()); + let tool_names = vec!["custom_mutating_tool".to_string()]; + + assert_eq!( + resolve_worker_tool_concurrency(&rig_alignment, &tool_names), + 1 + ); + } + + #[derive(Debug, Deserialize, JsonSchema)] + struct TimedProbeArgs { + delay_ms: u64, + label: String, + } + + #[derive(Debug, Serialize)] + struct TimedProbeOutput { + label: String, + } + + #[derive(Debug, thiserror::Error)] + #[error("timed probe failed: {0}")] + struct TimedProbeError(String); + + #[derive(Clone)] + struct TimedProbeTool { + active_calls: Arc, + max_active_calls: Arc, + } + + impl TimedProbeTool { + fn new(active_calls: Arc, max_active_calls: Arc) -> Self { + Self { + active_calls, + max_active_calls, + } + } + } + + impl Tool for TimedProbeTool { + const NAME: &'static str = "timed_probe"; + + type Error = TimedProbeError; + type Args = TimedProbeArgs; + type Output = TimedProbeOutput; + + async fn definition(&self, _prompt: String) -> ToolDefinition { + ToolDefinition { + name: Self::NAME.to_string(), + description: "Sleep for a fixed duration and report probe completion".to_string(), + parameters: serde_json::json!({ + "type": "object", + "properties": { + "delay_ms": { + "type": "integer", + "minimum": 1 + }, + "label": { + "type": "string" + } + }, + "required": ["delay_ms", "label"] + }), + } + } + + async fn call(&self, args: Self::Args) -> Result { + let active = self.active_calls.fetch_add(1, Ordering::SeqCst) + 1; + self.max_active_calls.fetch_max(active, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(args.delay_ms)).await; + self.active_calls.fetch_sub(1, Ordering::SeqCst); + + Ok(TimedProbeOutput { label: args.label }) + } + } + + #[derive(Clone, Default)] + struct PromptToolConcurrencyModel; + + impl PromptToolConcurrencyModel { + fn tool_result_count(request: &CompletionRequest) -> usize { + request + .chat_history + .iter() + .filter_map(|message| match message { + Message::User { content } => Some( + content + .iter() + .filter(|item| matches!(item, UserContent::ToolResult(_))) + .count(), + ), + _ => None, + }) + .sum() + } + + fn completion_response(choice: OneOrMany) -> CompletionResponse<()> { + CompletionResponse { + choice, + message_id: None, + usage: Usage::default(), + raw_response: (), + } + } + } + + impl CompletionModel for PromptToolConcurrencyModel { + type Response = (); + type StreamingResponse = (); + type Client = (); + + fn make(_client: &Self::Client, _model: impl Into) -> Self { + Self + } + + async fn completion( + &self, + request: CompletionRequest, + ) -> Result, CompletionError> { + match Self::tool_result_count(&request) { + 0 => Ok(Self::completion_response( + OneOrMany::many(vec![ + AssistantContent::tool_call( + "probe_call_1", + TimedProbeTool::NAME, + serde_json::json!({ + "delay_ms": 150, + "label": "probe-1", + }), + ), + AssistantContent::tool_call( + "probe_call_2", + TimedProbeTool::NAME, + serde_json::json!({ + "delay_ms": 150, + "label": "probe-2", + }), + ), + ]) + .expect("tool call batch should be non-empty"), + )), + 2 => Ok(Self::completion_response(OneOrMany::one( + AssistantContent::text("done"), + ))), + observed => Err(CompletionError::ProviderError(format!( + "expected two tool results before final response, observed {observed}" + ))), + } + } + + async fn stream( + &self, + _request: CompletionRequest, + ) -> Result, CompletionError> { + Err(CompletionError::ProviderError( + "streaming not implemented for PromptToolConcurrencyModel".to_string(), + )) + } + } + + async fn measure_prompt_tool_concurrency(concurrency: usize) -> (Duration, usize) { + let active_calls = Arc::new(AtomicUsize::new(0)); + let max_active_calls = Arc::new(AtomicUsize::new(0)); + let tool_server = ToolServer::new() + .tool(TimedProbeTool::new( + active_calls.clone(), + max_active_calls.clone(), + )) + .run(); + let agent = AgentBuilder::new(PromptToolConcurrencyModel) + .default_max_turns(4) + .tool_server_handle(tool_server) + .build(); + let (event_tx, _event_rx) = tokio::sync::broadcast::channel(8); + let hook = SpacebotHook::new( + Arc::::from("agent"), + ProcessId::Worker(uuid::Uuid::new_v4()), + ProcessType::Worker, + None, + event_tx, + ); + let mut history = Vec::new(); + let start = Instant::now(); + let result = hook + .prompt_once_with_tool_concurrency(&agent, &mut history, "run the probes", concurrency) + .await + .expect("prompt should complete after tool calls"); + + assert_eq!(result, "done"); + (start.elapsed(), max_active_calls.load(Ordering::SeqCst)) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn worker_concurrency_reduces_wall_clock() { + let rig_alignment = RigAlignmentConfig { + worker_read_only_tool_concurrency: 2, + worker_max_tool_concurrency: 2, + ..RigAlignmentConfig::default() + }; + + let allowlisted_tools = vec!["read_skill".to_string(), "web_search".to_string()]; + let mixed_surface_tools = vec!["read_skill".to_string(), "shell".to_string()]; + + let parallel_concurrency = + resolve_worker_tool_concurrency(&rig_alignment, &allowlisted_tools); + let sequential_concurrency = + resolve_worker_tool_concurrency(&rig_alignment, &mixed_surface_tools); + + assert_eq!(parallel_concurrency, 2); + assert_eq!(sequential_concurrency, 1); + + let deterministic_margin = Duration::from_millis(70); + let (sequential_elapsed, sequential_max_active) = + measure_prompt_tool_concurrency(sequential_concurrency).await; + let (parallel_elapsed, parallel_max_active) = + measure_prompt_tool_concurrency(parallel_concurrency).await; + + assert!( + parallel_elapsed + deterministic_margin < sequential_elapsed, + "expected allowlisted concurrency to reduce wall clock (sequential={sequential_elapsed:?}, parallel={parallel_elapsed:?}, margin={deterministic_margin:?})" + ); + assert_eq!( + sequential_max_active, 1, + "sequential prompt path should never overlap tool execution" + ); + assert_eq!( + parallel_max_active, 2, + "parallel prompt path should overlap both tool calls" + ); + } +} diff --git a/src/api/agents.rs b/src/api/agents.rs index 2fe72ddf0..92d24b919 100644 --- a/src/api/agents.rs +++ b/src/api/agents.rs @@ -567,6 +567,7 @@ pub(super) async fn create_agent( warmup: None, browser: None, channel: None, + rig_alignment: None, mcp: None, brave_search_key: None, cron_timezone: None, diff --git a/src/api/cortex.rs b/src/api/cortex.rs index 1ae192a37..69aceb3eb 100644 --- a/src/api/cortex.rs +++ b/src/api/cortex.rs @@ -108,6 +108,7 @@ pub(super) async fn cortex_chat_messages( /// /// The stream emits: /// - `thinking` — cortex is processing +/// - `text_delta` — partial text from the model stream /// - `tool_started` — a tool call began /// - `tool_completed` — a tool call finished (with result preview) /// - `done` — full response text @@ -146,6 +147,7 @@ pub(super) async fn cortex_chat_send( while let Some(event) = event_rx.recv().await { let event_name = match &event { CortexChatEvent::Thinking => "thinking", + CortexChatEvent::TextDelta { .. } => "text_delta", CortexChatEvent::ToolStarted { .. } => "tool_started", CortexChatEvent::ToolCompleted { .. } => "tool_completed", CortexChatEvent::Done { .. } => "done", diff --git a/src/api/state.rs b/src/api/state.rs index 0f97f5d6e..b849b9f47 100644 --- a/src/api/state.rs +++ b/src/api/state.rs @@ -148,6 +148,11 @@ pub enum ApiEvent { text_delta: String, aggregated_text: String, }, + /// Terminal signal for a streamed assistant message with no final text payload. + OutboundStreamEnd { + agent_id: String, + channel_id: String, + }, /// A worker was started. WorkerStarted { agent_id: String, diff --git a/src/api/system.rs b/src/api/system.rs index 2891caa9e..7f1f3111a 100644 --- a/src/api/system.rs +++ b/src/api/system.rs @@ -85,6 +85,7 @@ pub(super) async fn events_sse( ApiEvent::InboundMessage { .. } => "inbound_message", ApiEvent::OutboundMessage { .. } => "outbound_message", ApiEvent::OutboundMessageDelta { .. } => "outbound_message_delta", + ApiEvent::OutboundStreamEnd { .. } => "outbound_stream_end", ApiEvent::TypingState { .. } => "typing_state", ApiEvent::WorkerStarted { .. } => "worker_started", ApiEvent::WorkerStatusUpdate { .. } => "worker_status", diff --git a/src/config/load.rs b/src/config/load.rs index f78657ecb..dd214b94c 100644 --- a/src/config/load.rs +++ b/src/config/load.rs @@ -14,10 +14,10 @@ use super::{ CoalesceConfig, CompactionConfig, Config, CortexConfig, CronDef, DefaultsConfig, DiscordConfig, DiscordInstanceConfig, EmailConfig, EmailInstanceConfig, GroupDef, HumanDef, IngestionConfig, LinkDef, LlmConfig, McpServerConfig, McpTransport, MemoryPersistenceConfig, MessagingConfig, - MetricsConfig, OpenCodeConfig, ProviderConfig, SlackCommandConfig, SlackConfig, - SlackInstanceConfig, TelegramConfig, TelegramInstanceConfig, TelemetryConfig, TwitchConfig, - TwitchInstanceConfig, WarmupConfig, WebhookConfig, normalize_adapter, - validate_named_messaging_adapters, + MetricsConfig, OpenCodeConfig, ProviderConfig, RigAlignmentConfig, RigSemanticsMode, + SlackCommandConfig, SlackConfig, SlackInstanceConfig, TelegramConfig, TelegramInstanceConfig, + TelemetryConfig, TwitchConfig, TwitchInstanceConfig, WarmupConfig, WebhookConfig, + normalize_adapter, validate_named_messaging_adapters, }; use crate::error::{ConfigError, Result}; @@ -126,6 +126,59 @@ fn parse_close_policy(value: Option<&str>) -> Option { } } +fn parse_rig_semantics_mode(value: Option<&str>, default: RigSemanticsMode) -> RigSemanticsMode { + match value.map(str::trim).filter(|value| !value.is_empty()) { + Some("off") => RigSemanticsMode::Off, + Some("shadow") => RigSemanticsMode::Shadow, + Some("enforced") => RigSemanticsMode::Enforced, + Some(other) => { + tracing::warn!( + value = other, + "unknown rig_alignment.request_semantics_mode, expected one of: off, shadow, enforced" + ); + default + } + None => default, + } +} + +fn resolve_rig_alignment( + overrides: Option, + defaults: RigAlignmentConfig, +) -> RigAlignmentConfig { + let Some(overrides) = overrides else { + return defaults; + }; + + RigAlignmentConfig { + request_semantics_mode: parse_rig_semantics_mode( + overrides.request_semantics_mode.as_deref(), + defaults.request_semantics_mode, + ), + channel_streaming: overrides + .channel_streaming + .unwrap_or(defaults.channel_streaming), + cortex_chat_streaming: overrides + .cortex_chat_streaming + .unwrap_or(defaults.cortex_chat_streaming), + worker_read_only_tool_concurrency: overrides + .worker_read_only_tool_concurrency + .unwrap_or(defaults.worker_read_only_tool_concurrency), + worker_max_tool_concurrency: overrides + .worker_max_tool_concurrency + .unwrap_or(defaults.worker_max_tool_concurrency), + worker_read_only_tool_allowlist: overrides + .worker_read_only_tool_allowlist + .unwrap_or(defaults.worker_read_only_tool_allowlist), + output_schema_provider_allowlist: overrides + .output_schema_provider_allowlist + .unwrap_or(defaults.output_schema_provider_allowlist), + tool_choice_provider_allowlist: overrides + .tool_choice_provider_allowlist + .unwrap_or(defaults.tool_choice_provider_allowlist), + } +} + impl CortexConfig { fn resolve(overrides: TomlCortexConfig, defaults: CortexConfig) -> CortexConfig { CortexConfig { @@ -793,6 +846,7 @@ impl Config { warmup: None, browser: None, channel: None, + rig_alignment: None, mcp: None, brave_search_key: None, cron_timezone: None, @@ -1316,6 +1370,10 @@ impl Config { }; let defaults = DefaultsConfig { routing: resolve_routing(toml.defaults.routing, &base_routing), + rig_alignment: resolve_rig_alignment( + toml.defaults.rig_alignment, + base_defaults.rig_alignment.clone(), + ), max_concurrent_branches: toml .defaults .max_concurrent_branches @@ -1609,6 +1667,9 @@ impl Config { .listen_only_mode .map(|listen_only_mode| ChannelConfig { listen_only_mode }) }), + rig_alignment: a.rig_alignment.map(|rig_alignment| { + resolve_rig_alignment(Some(rig_alignment), defaults.rig_alignment.clone()) + }), mcp: match a.mcp { Some(mcp_servers) => Some( mcp_servers @@ -1648,6 +1709,7 @@ impl Config { warmup: None, browser: None, channel: None, + rig_alignment: None, mcp: None, brave_search_key: None, cron_timezone: None, diff --git a/src/config/runtime.rs b/src/config/runtime.rs index 47dfb70df..a33d585d9 100644 --- a/src/config/runtime.rs +++ b/src/config/runtime.rs @@ -6,7 +6,8 @@ use arc_swap::ArcSwap; use super::{ BrowserConfig, ChannelConfig, CoalesceConfig, CompactionConfig, Config, CortexConfig, DefaultsConfig, IngestionConfig, McpServerConfig, MemoryPersistenceConfig, OpenCodeConfig, - ResolvedAgentConfig, WarmupConfig, WarmupStatus, WorkReadiness, evaluate_work_readiness, + ResolvedAgentConfig, RigAlignmentConfig, WarmupConfig, WarmupStatus, WorkReadiness, + evaluate_work_readiness, }; use crate::llm::routing::RoutingConfig; use crate::tools::browser::SharedBrowserHandle; @@ -27,6 +28,7 @@ pub struct RuntimeConfig { pub coalesce: ArcSwap, pub ingestion: ArcSwap, pub channel_config: ArcSwap, + pub rig_alignment: ArcSwap, pub max_turns: ArcSwap, pub branch_max_turns: ArcSwap, pub context_window: ArcSwap, @@ -103,6 +105,7 @@ impl RuntimeConfig { coalesce: ArcSwap::from_pointee(agent_config.coalesce), ingestion: ArcSwap::from_pointee(agent_config.ingestion), channel_config: ArcSwap::from_pointee(agent_config.channel), + rig_alignment: ArcSwap::from_pointee(agent_config.rig_alignment.clone()), max_turns: ArcSwap::from_pointee(agent_config.max_turns), branch_max_turns: ArcSwap::from_pointee(agent_config.branch_max_turns), context_window: ArcSwap::from_pointee(agent_config.context_window), @@ -242,6 +245,8 @@ impl RuntimeConfig { .unwrap_or(current.as_ref().listen_only_mode); Arc::new(next) }); + self.rig_alignment + .store(Arc::new(resolved.rig_alignment.clone())); self.max_turns.store(Arc::new(resolved.max_turns)); self.branch_max_turns .store(Arc::new(resolved.branch_max_turns)); diff --git a/src/config/toml_schema.rs b/src/config/toml_schema.rs index b1f3041c0..bc29562c4 100644 --- a/src/config/toml_schema.rs +++ b/src/config/toml_schema.rs @@ -265,6 +265,7 @@ impl<'de> Deserialize<'de> for TomlLlmConfig { #[derive(Deserialize, Default)] pub(super) struct TomlDefaultsConfig { pub(super) routing: Option, + pub(super) rig_alignment: Option, pub(super) max_concurrent_branches: Option, pub(super) max_concurrent_workers: Option, pub(super) max_turns: Option, @@ -306,6 +307,18 @@ pub(super) struct TomlRoutingConfig { pub(super) fallbacks: Option>>, } +#[derive(Deserialize, Clone)] +pub(super) struct TomlRigAlignmentConfig { + pub(super) request_semantics_mode: Option, + pub(super) channel_streaming: Option, + pub(super) cortex_chat_streaming: Option, + pub(super) worker_read_only_tool_concurrency: Option, + pub(super) worker_max_tool_concurrency: Option, + pub(super) worker_read_only_tool_allowlist: Option>, + pub(super) output_schema_provider_allowlist: Option>, + pub(super) tool_choice_provider_allowlist: Option>, +} + #[derive(Deserialize)] pub(super) struct TomlMemoryPersistenceConfig { pub(super) enabled: Option, @@ -435,6 +448,7 @@ pub(super) struct TomlAgentConfig { pub(super) warmup: Option, pub(super) browser: Option, pub(super) channel: Option, + pub(super) rig_alignment: Option, pub(super) mcp: Option>, pub(super) brave_search_key: Option, pub(super) cron_timezone: Option, diff --git a/src/config/types.rs b/src/config/types.rs index 3588b2618..b96a35d91 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -499,6 +499,7 @@ impl SystemSecrets for LlmConfig { #[derive(Clone)] pub struct DefaultsConfig { pub routing: RoutingConfig, + pub rig_alignment: RigAlignmentConfig, pub max_concurrent_branches: usize, pub max_concurrent_workers: usize, pub max_turns: usize, @@ -530,6 +531,7 @@ impl std::fmt::Debug for DefaultsConfig { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("DefaultsConfig") .field("routing", &self.routing) + .field("rig_alignment", &self.rig_alignment) .field("max_concurrent_branches", &self.max_concurrent_branches) .field("max_concurrent_workers", &self.max_concurrent_workers) .field("max_turns", &self.max_turns) @@ -775,6 +777,74 @@ pub struct ChannelConfig { pub listen_only_mode: bool, } +/// Rollout mode for request-semantics parity behavior. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RigSemanticsMode { + Off, + Shadow, + Enforced, +} + +/// Runtime configuration for Rig-alignment behavior. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RigAlignmentConfig { + pub request_semantics_mode: RigSemanticsMode, + pub channel_streaming: bool, + pub cortex_chat_streaming: bool, + pub worker_read_only_tool_concurrency: usize, + pub worker_max_tool_concurrency: usize, + pub worker_read_only_tool_allowlist: Vec, + pub output_schema_provider_allowlist: Vec, + pub tool_choice_provider_allowlist: Vec, +} + +const BUILTIN_READ_ONLY_WORKER_TOOLS: &[&str] = &[ + "read_skill", + "spacebot_docs", + "memory_recall", + "channel_recall", + "web_search", + "worker_inspect", + "email_search", + "config_inspect", +]; + +pub fn is_builtin_read_only_worker_tool(tool_name: &str) -> bool { + BUILTIN_READ_ONLY_WORKER_TOOLS.contains(&tool_name) +} + +pub fn default_worker_read_only_tool_allowlist() -> Vec { + BUILTIN_READ_ONLY_WORKER_TOOLS + .iter() + .copied() + .map(str::to_string) + .collect() +} + +impl Default for RigAlignmentConfig { + fn default() -> Self { + Self { + request_semantics_mode: RigSemanticsMode::Shadow, + channel_streaming: false, + cortex_chat_streaming: false, + worker_read_only_tool_concurrency: 1, + worker_max_tool_concurrency: 4, + worker_read_only_tool_allowlist: default_worker_read_only_tool_allowlist(), + output_schema_provider_allowlist: vec![ + "openai".to_string(), + "openai-chatgpt".to_string(), + "anthropic".to_string(), + ], + tool_choice_provider_allowlist: vec![ + "openai".to_string(), + "openai-chatgpt".to_string(), + "anthropic".to_string(), + ], + } + } +} + /// OpenCode subprocess worker configuration. #[derive(Debug, Clone, PartialEq, Eq)] pub struct OpenCodeConfig { @@ -1003,6 +1073,7 @@ pub struct AgentConfig { pub warmup: Option, pub browser: Option, pub channel: Option, + pub rig_alignment: Option, pub mcp: Option>, /// Per-agent Brave Search API key override. None inherits from defaults. pub brave_search_key: Option, @@ -1059,6 +1130,7 @@ pub struct ResolvedAgentConfig { pub warmup: WarmupConfig, pub browser: BrowserConfig, pub channel: ChannelConfig, + pub rig_alignment: RigAlignmentConfig, pub mcp: Vec, pub brave_search_key: Option, pub cron_timezone: Option, @@ -1074,6 +1146,7 @@ impl Default for DefaultsConfig { fn default() -> Self { Self { routing: RoutingConfig::default(), + rig_alignment: RigAlignmentConfig::default(), max_concurrent_branches: 5, max_concurrent_workers: 5, max_turns: 5, @@ -1151,6 +1224,10 @@ impl AgentConfig { .clone() .unwrap_or_else(|| defaults.browser.clone()), channel: self.channel.unwrap_or(defaults.channel), + rig_alignment: self + .rig_alignment + .clone() + .unwrap_or_else(|| defaults.rig_alignment.clone()), mcp: resolve_mcp_configs(&defaults.mcp, self.mcp.as_deref()), brave_search_key: self .brave_search_key diff --git a/src/hooks/spacebot.rs b/src/hooks/spacebot.rs index 9cc7bace9..52385af63 100644 --- a/src/hooks/spacebot.rs +++ b/src/hooks/spacebot.rs @@ -113,6 +113,22 @@ impl SpacebotHook { history: &mut Vec, prompt: &str, ) -> std::result::Result + where + M: CompletionModel, + { + self.prompt_with_tool_nudge_retry_with_concurrency(agent, history, prompt, 1) + .await + } + + /// Prompt an agent with bounded hook-driven tool nudge retries and request-level + /// tool concurrency. + pub async fn prompt_with_tool_nudge_retry_with_concurrency( + &self, + agent: &rig::agent::Agent, + history: &mut Vec, + prompt: &str, + tool_concurrency: usize, + ) -> std::result::Result where M: CompletionModel, { @@ -125,11 +141,20 @@ impl SpacebotHook { loop { let history_len_before_attempt = history.len(); - let result = agent - .prompt(current_prompt.as_ref()) - .with_history(history) - .with_hook(self.clone()) - .await; + let result = if tool_concurrency > 1 { + agent + .prompt(current_prompt.as_ref()) + .with_history(history) + .with_hook(self.clone()) + .with_tool_concurrency(tool_concurrency) + .await + } else { + agent + .prompt(current_prompt.as_ref()) + .with_history(history) + .with_hook(self.clone()) + .await + }; match &result { Err(PromptError::PromptCancelled { reason, .. }) @@ -225,16 +250,40 @@ impl SpacebotHook { history: &mut Vec, prompt: &str, ) -> std::result::Result + where + M: CompletionModel, + { + self.prompt_once_with_tool_concurrency(agent, history, prompt, 1) + .await + } + + /// Prompt once with a per-request tool concurrency setting. + pub async fn prompt_once_with_tool_concurrency( + &self, + agent: &rig::agent::Agent, + history: &mut Vec, + prompt: &str, + tool_concurrency: usize, + ) -> std::result::Result where M: CompletionModel, { self.reset_tool_nudge_state(); self.set_tool_nudge_request_active(false); - agent - .prompt(prompt) - .with_history(history) - .with_hook(self.clone()) - .await + if tool_concurrency > 1 { + agent + .prompt(prompt) + .with_history(history) + .with_hook(self.clone()) + .with_tool_concurrency(tool_concurrency) + .await + } else { + agent + .prompt(prompt) + .with_history(history) + .with_hook(self.clone()) + .await + } } /// Send a status update event. @@ -465,12 +514,14 @@ where if self.process_type == ProcessType::Channel && let Some(channel_id) = self.channel_id.clone() { + let scrubbed_text_delta = crate::secrets::scrub::scrub_leaks(text_delta); + let scrubbed_aggregated_text = crate::secrets::scrub::scrub_leaks(aggregated_text); let event = ProcessEvent::TextDelta { agent_id: self.agent_id.clone(), process_id: self.process_id.clone(), channel_id: Some(channel_id), - text_delta: text_delta.to_string(), - aggregated_text: aggregated_text.to_string(), + text_delta: scrubbed_text_delta, + aggregated_text: scrubbed_aggregated_text, }; self.event_tx.send(event).ok(); } @@ -503,7 +554,8 @@ where } // Send event without blocking. Truncate args to keep broadcast payloads bounded. - let capped_args = crate::tools::truncate_output(args, 2_000); + let scrubbed_args = crate::secrets::scrub::scrub_leaks(args); + let capped_args = crate::tools::truncate_output(&scrubbed_args, 2_000); let event = ProcessEvent::ToolStarted { agent_id: self.agent_id.clone(), process_id: self.process_id.clone(), @@ -1121,7 +1173,8 @@ mod tests { ); let action = - >::on_text_delta(&hook, "hi", "hi").await; + >::on_text_delta(&hook, "hi", "partial hi") + .await; assert!(matches!(action, HookAction::Continue)); let event = event_rx.recv().await.expect("text delta event"); @@ -1131,7 +1184,65 @@ mod tests { ref text_delta, ref aggregated_text, .. - } if text_delta == "hi" && aggregated_text == "hi" + } if text_delta == "hi" && aggregated_text == "partial hi" + )); + } + + #[tokio::test] + async fn channel_text_delta_scrubs_secret_patterns() { + let (event_tx, mut event_rx) = tokio::sync::broadcast::channel(8); + let hook = SpacebotHook::new( + std::sync::Arc::::from("agent"), + ProcessId::Channel(std::sync::Arc::::from("channel")), + ProcessType::Channel, + Some(std::sync::Arc::::from("channel")), + event_tx, + ); + + let action = >::on_text_delta( + &hook, + "sk-ant-api03-abc1234567890123456789012345678901234567", + "token sk-ant-api03-abc1234567890123456789012345678901234567", + ) + .await; + assert!(matches!(action, HookAction::Continue)); + + let event = event_rx.recv().await.expect("text delta event"); + assert!(matches!( + event, + ProcessEvent::TextDelta { + ref text_delta, + ref aggregated_text, + .. + } if !text_delta.contains("sk-ant-") && !aggregated_text.contains("sk-ant-") + )); + } + + #[tokio::test] + async fn tool_started_event_scrubs_secret_patterns() { + let (event_tx, mut event_rx) = tokio::sync::broadcast::channel(8); + let hook = SpacebotHook::new( + std::sync::Arc::::from("agent"), + ProcessId::Worker(uuid::Uuid::new_v4()), + ProcessType::Worker, + None, + event_tx, + ); + + let action = >::on_tool_call( + &hook, + "shell", + None, + "internal-call", + r#"{"command":"export API_KEY=sk-ant-api03-abc1234567890123456789012345678901234567"}"#, + ) + .await; + assert!(matches!(action, rig::agent::ToolCallHookAction::Continue)); + + let event = event_rx.recv().await.expect("tool started event"); + assert!(matches!( + event, + ProcessEvent::ToolStarted { ref args, .. } if !args.contains("sk-ant-") )); } } diff --git a/src/lib.rs b/src/lib.rs index 3f20573f8..1245f230f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -608,6 +608,8 @@ pub enum OutboundResponse { post_at: i64, }, StreamStart, + /// Incremental streamed text delta. Adapters that edit placeholders should + /// accumulate deltas locally and decide when to flush visible updates. StreamChunk(String), StreamEnd, Status(StatusUpdate), diff --git a/src/llm/anthropic.rs b/src/llm/anthropic.rs index be7e39eb4..0dec916ed 100644 --- a/src/llm/anthropic.rs +++ b/src/llm/anthropic.rs @@ -7,5 +7,5 @@ pub mod tools; pub use auth::{AnthropicAuthPath, apply_auth_headers, detect_auth_path}; pub use cache::{CacheRetention, get_cache_control, resolve_cache_retention}; -pub use params::build_anthropic_request; +pub use params::{AnthropicRequestOptions, build_anthropic_request}; pub use tools::{from_claude_code_name, to_claude_code_name}; diff --git a/src/llm/anthropic/params.rs b/src/llm/anthropic/params.rs index 11dcda8ef..fce7e5c25 100644 --- a/src/llm/anthropic/params.rs +++ b/src/llm/anthropic/params.rs @@ -5,7 +5,8 @@ use super::cache; use super::tools; use reqwest::RequestBuilder; -use rig::completion::CompletionRequest; +use rig::completion::{CompletionError, CompletionRequest}; +use rig::message::ToolChoice; const CLAUDE_CODE_SYSTEM_PREAMBLE: &str = "You are Claude Code, Anthropic's official CLI for Claude."; @@ -19,6 +20,13 @@ pub struct AnthropicRequest { pub original_tools: Vec<(String, String)>, } +/// Request-time behavior toggles for Anthropic request assembly. +pub struct AnthropicRequestOptions { + pub force_bearer: bool, + pub apply_tool_choice: bool, + pub apply_output_schema: bool, +} + /// Adaptive thinking is only available on 4.6-generation models. fn supports_adaptive_thinking(model_id: &str) -> bool { model_id.contains("opus-4-6") @@ -58,8 +66,13 @@ pub fn build_anthropic_request( model_name: &str, request: &CompletionRequest, thinking_effort: &str, - force_bearer: bool, -) -> AnthropicRequest { + options: AnthropicRequestOptions, +) -> Result { + let AnthropicRequestOptions { + force_bearer, + apply_tool_choice, + apply_output_schema, + } = options; let is_oauth = auth::detect_auth_path(api_key, force_bearer) == AnthropicAuthPath::OAuthToken; let adaptive_thinking = supports_adaptive_thinking(model_name); let retention = cache::resolve_cache_retention(None); @@ -77,6 +90,19 @@ pub fn build_anthropic_request( body["messages"] = serde_json::json!(messages); let original_tools = build_tools(&mut body, request, is_oauth, &cache_control); + if apply_tool_choice && let Some(tool_choice) = request.tool_choice.as_ref() { + body["tool_choice"] = map_anthropic_tool_choice(tool_choice)?; + } + if apply_output_schema && let Some(output_schema) = request.output_schema.as_ref() { + let mut schema = output_schema.clone().to_value(); + sanitize_anthropic_schema(&mut schema); + body["output_config"] = serde_json::json!({ + "format": { + "type": "json_schema", + "schema": schema, + } + }); + } if let Some(temperature) = request.temperature { body["temperature"] = serde_json::json!(temperature); @@ -94,7 +120,11 @@ pub fn build_anthropic_request( } } }; - body["output_config"] = serde_json::json!({ "effort": effort }); + if body["output_config"].is_object() { + body["output_config"]["effort"] = serde_json::json!(effort); + } else { + body["output_config"] = serde_json::json!({ "effort": effort }); + } } let builder = http_client @@ -105,11 +135,11 @@ pub fn build_anthropic_request( let (builder, auth_path) = auth::apply_auth_headers(builder, api_key, false, force_bearer); let builder = builder.json(&body); - AnthropicRequest { + Ok(AnthropicRequest { builder, auth_path, original_tools, - } + }) } fn build_system_prompt( @@ -147,6 +177,66 @@ fn build_system_prompt( } } +fn map_anthropic_tool_choice( + tool_choice: &ToolChoice, +) -> Result { + let mapped = match tool_choice { + ToolChoice::Auto => serde_json::json!({ "type": "auto" }), + ToolChoice::None => serde_json::json!({ "type": "none" }), + ToolChoice::Required => serde_json::json!({ "type": "any" }), + ToolChoice::Specific { function_names } => { + let Some(function_name) = function_names.first() else { + return Err(CompletionError::ProviderError( + "tool_choice specific requires at least one function name".to_string(), + )); + }; + if function_names.len() > 1 { + return Err(CompletionError::ProviderError( + "Anthropic only supports one specific tool name".to_string(), + )); + } + serde_json::json!({ + "type": "tool", + "name": function_name, + }) + } + }; + Ok(mapped) +} + +fn sanitize_anthropic_schema(schema: &mut serde_json::Value) { + use serde_json::Value; + + if let Value::Object(object) = schema { + let is_object_schema = object.get("type") == Some(&Value::String("object".to_string())) + || object.contains_key("properties"); + + if is_object_schema && !object.contains_key("additionalProperties") { + object.insert("additionalProperties".to_string(), Value::Bool(false)); + } + + for key in ["$defs", "properties"] { + if let Some(Value::Object(property_map)) = object.get_mut(key) { + for property in property_map.values_mut() { + sanitize_anthropic_schema(property); + } + } + } + + if let Some(items) = object.get_mut("items") { + sanitize_anthropic_schema(items); + } + + for key in ["anyOf", "allOf", "oneOf"] { + if let Some(Value::Array(variants)) = object.get_mut(key) { + for variant in variants { + sanitize_anthropic_schema(variant); + } + } + } + } +} + /// Build tool definitions, optionally normalizing names. Returns the original /// tool (name, description) pairs for reverse-mapping on response. fn build_tools( @@ -201,6 +291,7 @@ fn build_tools( #[cfg(test)] mod tests { use super::*; + use schemars::JsonSchema; #[test] fn adaptive_thinking_detected_for_4_6_models() { @@ -218,4 +309,19 @@ mod tests { assert!(!supports_adaptive_thinking("claude-opus-4-0")); assert!(!supports_adaptive_thinking("gpt-4o")); } + + #[allow(dead_code)] + #[derive(JsonSchema)] + struct AnthropicOptionalSchema { + answer: String, + optional_note: Option, + } + + #[test] + fn anthropic_schema_sanitizer_preserves_optional_fields() { + let mut schema = schemars::schema_for!(AnthropicOptionalSchema).to_value(); + sanitize_anthropic_schema(&mut schema); + + assert_eq!(schema["required"], serde_json::json!(["answer"])); + } } diff --git a/src/llm/model.rs b/src/llm/model.rs index c74a76c4d..daea1618c 100644 --- a/src/llm/model.rs +++ b/src/llm/model.rs @@ -1,6 +1,6 @@ //! SpacebotModel: Custom CompletionModel implementation that routes through LlmManager. -use crate::config::{ApiType, ProviderConfig}; +use crate::config::{ApiType, ProviderConfig, RigAlignmentConfig, RigSemanticsMode}; use crate::llm::manager::LlmManager; use crate::llm::routing::{ self, MAX_FALLBACK_ATTEMPTS, MAX_RETRIES_PER_MODEL, RETRY_BASE_DELAY_MS, RoutingConfig, @@ -9,11 +9,13 @@ use crate::llm::routing::{ use futures::StreamExt as _; use rig::completion::{self, CompletionError, CompletionModel, CompletionRequest, GetTokenUsage}; use rig::message::{ - AssistantContent, DocumentSourceKind, Image, Message, MimeType, Text, ToolCall, ToolFunction, - UserContent, + AssistantContent, DocumentSourceKind, Image, Message, MimeType, Text, ToolCall, ToolChoice, + ToolFunction, UserContent, }; use rig::one_or_many::OneOrMany; -use rig::streaming::{RawStreamingChoice, RawStreamingToolCall, StreamingCompletionResponse}; +use rig::streaming::{ + RawStreamingChoice, RawStreamingToolCall, StreamingCompletionResponse, StreamingResult, +}; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; use std::sync::Arc; @@ -39,6 +41,12 @@ impl GetTokenUsage for RawStreamingResponse { } } +#[derive(Debug)] +struct RequestSemanticDecision { + apply_tool_choice: bool, + apply_output_schema: bool, +} + /// Custom completion model that routes through LlmManager. /// /// Optionally holds a RoutingConfig for fallback behavior. When present, @@ -52,6 +60,7 @@ pub struct SpacebotModel { routing: Option, agent_id: Option, process_type: Option, + rig_alignment: Option, } impl SpacebotModel { @@ -82,6 +91,170 @@ impl SpacebotModel { self } + /// Attach Rig-alignment runtime config for request semantic behavior. + pub fn with_rig_alignment(mut self, rig_alignment: RigAlignmentConfig) -> Self { + self.rig_alignment = Some(rig_alignment); + self + } + + fn rig_alignment(&self) -> RigAlignmentConfig { + self.rig_alignment.clone().unwrap_or_default() + } + + fn process_label(&self) -> &str { + self.process_type.as_deref().unwrap_or("unknown") + } + + fn record_request_semantic_decision(&self, field: &'static str, decision: &'static str) { + #[cfg(feature = "metrics")] + { + let metrics = crate::telemetry::Metrics::global(); + metrics + .rig_request_semantics_total + .with_label_values(&[self.process_label(), &self.provider, field, decision]) + .inc(); + } + #[cfg(not(feature = "metrics"))] + let _ = (field, decision); + } + + fn record_drift_surface(&self, surface: &'static str) { + #[cfg(feature = "metrics")] + { + let metrics = crate::telemetry::Metrics::global(); + metrics + .rig_drift_detected_total + .with_label_values(&[surface]) + .inc(); + } + #[cfg(not(feature = "metrics"))] + let _ = surface; + } + + fn evaluate_request_semantics( + &self, + request: &CompletionRequest, + ) -> Result { + let rig_alignment = self.rig_alignment(); + let mode = rig_alignment.request_semantics_mode; + let provider = self.provider.as_str(); + + let tool_choice_requested = request.tool_choice.is_some(); + let tool_choice_requires_tool_definitions = matches!( + request.tool_choice, + Some(ToolChoice::Required) | Some(ToolChoice::Specific { .. }) + ); + let output_schema_requested = request.output_schema.is_some(); + + if matches!(mode, RigSemanticsMode::Off) { + if tool_choice_requested { + self.record_request_semantic_decision("tool_choice", "unsupported"); + self.record_drift_surface("tool_choice"); + } + if output_schema_requested { + self.record_request_semantic_decision("output_schema", "unsupported"); + self.record_drift_surface("output_schema"); + } + return Ok(RequestSemanticDecision { + apply_tool_choice: false, + apply_output_schema: false, + }); + } + + let tool_choice_supported = if !tool_choice_requested { + true + } else if tool_choice_requires_tool_definitions && request.tools.is_empty() { + tracing::warn!( + provider, + mode = ?mode, + field = "tool_choice", + "request semantics rejected: tool_choice requires at least one tool definition" + ); + false + } else { + rig_alignment + .tool_choice_provider_allowlist + .iter() + .any(|allowed| allowed == provider) + }; + let output_schema_supported = !output_schema_requested + || rig_alignment + .output_schema_provider_allowlist + .iter() + .any(|allowed| allowed == provider); + + if tool_choice_requested { + let decision = if tool_choice_supported { + "applied" + } else if matches!(mode, RigSemanticsMode::Enforced) { + "rejected" + } else { + "shadowed" + }; + self.record_request_semantic_decision("tool_choice", decision); + if decision != "applied" { + self.record_drift_surface("tool_choice"); + } + tracing::info!( + agent_id = self.agent_id.as_deref().unwrap_or("unknown"), + process_type = self.process_label(), + provider, + model = %self.full_model_name, + request_semantics_mode = ?mode, + tool_choice_requested, + output_schema_requested, + applied = tool_choice_supported, + decision, + field = "tool_choice", + "request semantic decision" + ); + } + if output_schema_requested { + let decision = if output_schema_supported { + "applied" + } else if matches!(mode, RigSemanticsMode::Enforced) { + "rejected" + } else { + "shadowed" + }; + self.record_request_semantic_decision("output_schema", decision); + if decision != "applied" { + self.record_drift_surface("output_schema"); + } + tracing::info!( + agent_id = self.agent_id.as_deref().unwrap_or("unknown"), + process_type = self.process_label(), + provider, + model = %self.full_model_name, + request_semantics_mode = ?mode, + tool_choice_requested, + output_schema_requested, + applied = output_schema_supported, + decision, + field = "output_schema", + "request semantic decision" + ); + } + + if matches!(mode, RigSemanticsMode::Enforced) { + if tool_choice_requested && !tool_choice_supported { + return Err(CompletionError::ProviderError(format!( + "request semantics rejected: provider '{provider}' does not support tool_choice in enforced mode" + ))); + } + if output_schema_requested && !output_schema_supported { + return Err(CompletionError::ProviderError(format!( + "request semantics rejected: provider '{provider}' does not support output_schema in enforced mode" + ))); + } + } + + Ok(RequestSemanticDecision { + apply_tool_choice: tool_choice_requested && tool_choice_supported, + apply_output_schema: output_schema_requested && output_schema_supported, + }) + } + async fn provider_config_for_current_model(&self) -> Result { let provider_id = self .full_model_name @@ -184,7 +357,13 @@ impl SpacebotModel { let model = if model_name == self.full_model_name { self.clone() } else { - SpacebotModel::make(&self.llm_manager, model_name) + let mut model = SpacebotModel::make(&self.llm_manager, model_name); + if let Some(rig_alignment) = self.rig_alignment.clone() { + model = model.with_rig_alignment(rig_alignment); + } + model.agent_id = self.agent_id.clone(); + model.process_type = self.process_type.clone(); + model }; let mut last_error = None; @@ -258,6 +437,7 @@ impl CompletionModel for SpacebotModel { routing: None, agent_id: None, process_type: None, + rig_alignment: None, } } @@ -453,8 +633,9 @@ impl CompletionModel for SpacebotModel { request: CompletionRequest, ) -> Result, CompletionError> { let provider_config = self.provider_config_for_current_model().await?; + let stream_mode = stream_mode_for_api_type(&provider_config.api_type); - match provider_config.api_type { + let stream = match provider_config.api_type { ApiType::OpenAiCompletions => self.stream_openai(request, &provider_config).await, ApiType::OpenAiChatCompletions => { let endpoint = format!( @@ -476,6 +657,7 @@ impl CompletionModel for SpacebotModel { &endpoint, Some(provider_config.api_key.clone()), &headers, + stream_mode, ) .await } @@ -493,18 +675,50 @@ impl CompletionModel for SpacebotModel { ("HTTP-Referer", "https://github.com/spacedriveapp/spacebot"), ("X-Title", "spacebot"), ], + stream_mode, ) .await } ApiType::Gemini => { - self.stream_openai_compatible(request, "Google Gemini", &provider_config) - .await + self.stream_openai_compatible( + request, + "Google Gemini", + &provider_config, + stream_mode, + ) + .await } ApiType::Anthropic | ApiType::OpenAiResponses => { let response = self.attempt_completion(request).await?; - Ok(stream_from_completion_response(response)) + Ok(stream_from_completion_response( + response, + self.process_label().to_string(), + self.provider.clone(), + self.full_model_name.clone(), + stream_mode, + )) } + }?; + + tracing::info!( + agent_id = self.agent_id.as_deref().unwrap_or("unknown"), + process_type = self.process_label(), + provider = self.provider.as_str(), + model = %self.full_model_name, + streaming_enabled = true, + stream_mode, + "starting streaming completion" + ); + #[cfg(feature = "metrics")] + { + let metrics = crate::telemetry::Metrics::global(); + metrics + .rig_stream_sessions_total + .with_label_values(&[self.process_label(), &self.provider, stream_mode]) + .inc(); } + + Ok(stream) } } @@ -514,6 +728,7 @@ impl SpacebotModel { request: CompletionRequest, provider_config: &ProviderConfig, ) -> Result, CompletionError> { + let semantics = self.evaluate_request_semantics(&request)?; let api_key = provider_config.api_key.as_str(); let effort = self @@ -528,8 +743,12 @@ impl SpacebotModel { &self.model_name, &request, effort, - provider_config.use_bearer_auth, - ); + crate::llm::anthropic::AnthropicRequestOptions { + force_bearer: provider_config.use_bearer_auth, + apply_tool_choice: semantics.apply_tool_choice, + apply_output_schema: semantics.apply_output_schema, + }, + )?; let is_oauth = anthropic_request.auth_path == crate::llm::anthropic::AnthropicAuthPath::OAuthToken; @@ -587,6 +806,7 @@ impl SpacebotModel { request: CompletionRequest, provider_config: &ProviderConfig, ) -> Result, CompletionError> { + let semantics = self.evaluate_request_semantics(&request)?; let api_key = provider_config.api_key.as_str(); let provider_label = provider_config .name @@ -637,6 +857,16 @@ impl SpacebotModel { .collect(); body["tools"] = serde_json::json!(tools); } + if semantics.apply_tool_choice + && let Some(tool_choice) = request.tool_choice.as_ref() + { + body["tool_choice"] = map_openai_tool_choice(tool_choice)?; + } + if semantics.apply_output_schema + && let Some(output_schema) = request.output_schema.as_ref() + { + body["response_format"] = map_openai_response_format(output_schema); + } let chat_completions_url = format!( "{}/v1/chat/completions", @@ -677,6 +907,7 @@ impl SpacebotModel { }, body, &provider_label, + "native", ) .await } @@ -686,6 +917,7 @@ impl SpacebotModel { request: CompletionRequest, provider_config: &ProviderConfig, ) -> Result, CompletionError> { + let semantics = self.evaluate_request_semantics(&request)?; let base_url = provider_config.base_url.trim_end_matches('/'); let is_chatgpt_codex = self.provider == "openai-chatgpt"; let responses_url = if is_chatgpt_codex { @@ -745,6 +977,16 @@ impl SpacebotModel { .collect(); body["tools"] = serde_json::json!(tools); } + if semantics.apply_tool_choice + && let Some(tool_choice) = request.tool_choice.as_ref() + { + body["tool_choice"] = map_openai_responses_tool_choice(tool_choice)?; + } + if semantics.apply_output_schema + && let Some(output_schema) = request.output_schema.as_ref() + { + body["text"] = map_openai_responses_text_config(output_schema); + } let openai_account_id = if self.provider == "openai-chatgpt" { self.llm_manager.get_openai_account_id().await @@ -817,7 +1059,12 @@ impl SpacebotModel { provider_config: &ProviderConfig, ) -> Result, CompletionError> { let stream = self - .stream_openai_compatible(request, provider_display_name, provider_config) + .stream_openai_compatible( + request, + provider_display_name, + provider_config, + stream_mode_for_api_type(&provider_config.api_type), + ) .await?; collect_streaming_completion_response(stream).await } @@ -827,7 +1074,9 @@ impl SpacebotModel { request: CompletionRequest, provider_display_name: &str, provider_config: &ProviderConfig, + stream_mode: &'static str, ) -> Result, CompletionError> { + let semantics = self.evaluate_request_semantics(&request)?; let base_url = provider_config.base_url.trim_end_matches('/'); let endpoint_path = match provider_config.api_type { ApiType::OpenAiCompletions | ApiType::OpenAiResponses => "/v1/chat/completions", @@ -888,6 +1137,16 @@ impl SpacebotModel { .collect(); body["tools"] = serde_json::json!(tools); } + if semantics.apply_tool_choice + && let Some(tool_choice) = request.tool_choice.as_ref() + { + body["tool_choice"] = map_openai_tool_choice(tool_choice)?; + } + if semantics.apply_output_schema + && let Some(output_schema) = request.output_schema.as_ref() + { + body["response_format"] = map_openai_response_format(output_schema); + } let http_client = self.llm_manager.http_client().clone(); let auth_header = format!("Bearer {api_key}"); @@ -901,6 +1160,7 @@ impl SpacebotModel { }, body, provider_display_name, + stream_mode, ) .await } @@ -926,6 +1186,7 @@ impl SpacebotModel { endpoint, api_key, extra_headers, + "native", ) .await?; collect_streaming_completion_response(stream).await @@ -938,7 +1199,9 @@ impl SpacebotModel { endpoint: &str, api_key: Option, extra_headers: &[(&str, &str)], + stream_mode: &'static str, ) -> Result, CompletionError> { + let semantics = self.evaluate_request_semantics(&request)?; let mut messages = Vec::new(); if let Some(preamble) = &request.preamble { @@ -981,6 +1244,16 @@ impl SpacebotModel { .collect(); body["tools"] = serde_json::json!(tools); } + if semantics.apply_tool_choice + && let Some(tool_choice) = request.tool_choice.as_ref() + { + body["tool_choice"] = map_openai_tool_choice(tool_choice)?; + } + if semantics.apply_output_schema + && let Some(output_schema) = request.output_schema.as_ref() + { + body["response_format"] = map_openai_response_format(output_schema); + } let http_client = self.llm_manager.http_client().clone(); let endpoint = endpoint.to_string(); @@ -1008,6 +1281,7 @@ impl SpacebotModel { }, body, provider_display_name, + stream_mode, ) .await } @@ -1017,6 +1291,7 @@ impl SpacebotModel { mut build_request: F, request_body: serde_json::Value, provider_label: &str, + stream_mode: &'static str, ) -> Result, CompletionError> where F: FnMut(&serde_json::Value) -> reqwest::RequestBuilder, @@ -1194,11 +1469,72 @@ impl SpacebotModel { })); }; - Ok(StreamingCompletionResponse::stream(Box::pin(stream))) + Ok(StreamingCompletionResponse::stream( + instrument_stream_result( + Box::pin(stream), + self.process_label().to_string(), + self.provider.clone(), + self.full_model_name.clone(), + stream_mode, + ), + )) } } // --- Helpers --- +fn stream_mode_for_api_type(api_type: &ApiType) -> &'static str { + match api_type { + ApiType::Anthropic | ApiType::OpenAiResponses => "synthetic", + ApiType::OpenAiCompletions + | ApiType::OpenAiChatCompletions + | ApiType::KiloGateway + | ApiType::Gemini => "native", + } +} + +fn instrument_stream_result( + mut inner: StreamingResult, + process_type: String, + provider: String, + model: String, + stream_mode: &'static str, +) -> StreamingResult { + let stream = async_stream::stream! { + let started_at = std::time::Instant::now(); + let mut first_text_delta_observed = false; + + while let Some(item) = inner.next().await { + if !first_text_delta_observed + && let Ok(RawStreamingChoice::Message(text)) = &item + && !text.is_empty() + { + first_text_delta_observed = true; + let elapsed_ms = started_at.elapsed().as_secs_f64() * 1000.0; + tracing::info!( + process_type = process_type.as_str(), + provider = provider.as_str(), + model = model.as_str(), + stream_mode, + time_to_first_delta_ms = elapsed_ms, + "stream produced first text delta" + ); + #[cfg(feature = "metrics")] + { + let metrics = crate::telemetry::Metrics::global(); + metrics + .rig_time_to_first_delta_ms + .with_label_values(&[process_type.as_str(), provider.as_str()]) + .observe(elapsed_ms); + } + } + + yield item; + } + }; + + Box::pin(stream) +} + /// Reverse-map Claude Code canonical tool names back to the original names /// from the request's tool definitions. fn reverse_map_tool_names( @@ -1528,6 +1864,190 @@ fn truncate_body(body: &str) -> &str { } } +fn schema_name(output_schema: &schemars::Schema) -> String { + const FALLBACK_SCHEMA_NAME: &str = "response_schema"; + const MAX_SCHEMA_NAME_LEN: usize = 64; + + let title = output_schema + .as_object() + .and_then(|schema_object| schema_object.get("title")) + .and_then(serde_json::Value::as_str) + .unwrap_or(FALLBACK_SCHEMA_NAME) + .trim(); + + let mut normalized = String::with_capacity(title.len().min(MAX_SCHEMA_NAME_LEN)); + for character in title.chars() { + if normalized.len() >= MAX_SCHEMA_NAME_LEN { + break; + } + + if character.is_ascii_alphanumeric() || character == '_' || character == '-' { + normalized.push(character); + } else if !normalized.ends_with('_') { + normalized.push('_'); + } + } + + let normalized = normalized.trim_matches('_'); + if normalized.is_empty() { + FALLBACK_SCHEMA_NAME.to_string() + } else { + normalized.to_string() + } +} + +fn sanitize_openai_schema(schema: &mut serde_json::Value) { + use serde_json::Value; + + if let Value::Object(object) = schema { + if object.contains_key("$ref") { + if let Some(defs) = object.get_mut("$defs") + && let Value::Object(defs_object) = defs + { + for schema in defs_object.values_mut() { + sanitize_openai_schema(schema); + } + } + object.retain(|key, _| key == "$ref" || key == "$defs"); + return; + } + + let is_object_schema = object.get("type") == Some(&Value::String("object".to_string())) + || object.contains_key("properties"); + + if is_object_schema && !object.contains_key("additionalProperties") { + object.insert("additionalProperties".to_string(), Value::Bool(false)); + } + + if let Some(defs) = object.get_mut("$defs") + && let Value::Object(defs_object) = defs + { + for schema in defs_object.values_mut() { + sanitize_openai_schema(schema); + } + } + + if let Some(properties) = object.get_mut("properties") + && let Value::Object(property_map) = properties + { + for schema in property_map.values_mut() { + sanitize_openai_schema(schema); + } + } + + if let Some(items) = object.get_mut("items") { + sanitize_openai_schema(items); + } + + if let Some(one_of) = object.remove("oneOf") { + match object.get_mut("anyOf") { + Some(Value::Array(existing)) => { + if let Value::Array(mut incoming) = one_of { + existing.append(&mut incoming); + } + } + _ => { + object.insert("anyOf".to_string(), one_of); + } + } + } + + for key in ["anyOf", "oneOf", "allOf"] { + if let Some(variants) = object.get_mut(key) + && let Value::Array(variant_array) = variants + { + for variant in variant_array { + sanitize_openai_schema(variant); + } + } + } + } +} + +fn sanitized_openai_schema_value(output_schema: &schemars::Schema) -> serde_json::Value { + let mut schema = output_schema.clone().to_value(); + sanitize_openai_schema(&mut schema); + schema +} + +fn map_openai_response_format(output_schema: &schemars::Schema) -> serde_json::Value { + serde_json::json!({ + "type": "json_schema", + "json_schema": { + "name": schema_name(output_schema), + "strict": true, + "schema": sanitized_openai_schema_value(output_schema), + } + }) +} + +fn map_openai_responses_text_config(output_schema: &schemars::Schema) -> serde_json::Value { + serde_json::json!({ + "format": { + "type": "json_schema", + "name": schema_name(output_schema), + "strict": true, + "schema": sanitized_openai_schema_value(output_schema), + } + }) +} + +fn map_openai_tool_choice(tool_choice: &ToolChoice) -> Result { + let mapped = match tool_choice { + ToolChoice::Auto => serde_json::json!("auto"), + ToolChoice::None => serde_json::json!("none"), + ToolChoice::Required => serde_json::json!("required"), + ToolChoice::Specific { function_names } => { + let Some(function_name) = function_names.first() else { + return Err(CompletionError::ProviderError( + "tool_choice specific requires at least one function name".to_string(), + )); + }; + if function_names.len() > 1 { + return Err(CompletionError::ProviderError( + "OpenAI-compatible tool_choice supports exactly one specific function name" + .to_string(), + )); + } + serde_json::json!({ + "type": "function", + "function": { + "name": function_name, + } + }) + } + }; + Ok(mapped) +} + +fn map_openai_responses_tool_choice( + tool_choice: &ToolChoice, +) -> Result { + let mapped = match tool_choice { + ToolChoice::Auto => serde_json::json!("auto"), + ToolChoice::None => serde_json::json!("none"), + ToolChoice::Required => serde_json::json!("required"), + ToolChoice::Specific { function_names } => { + let Some(function_name) = function_names.first() else { + return Err(CompletionError::ProviderError( + "tool_choice specific requires at least one function name".to_string(), + )); + }; + if function_names.len() > 1 { + return Err(CompletionError::ProviderError( + "Responses API tool_choice supports exactly one specific function name" + .to_string(), + )); + } + serde_json::json!({ + "type": "function", + "name": function_name, + }) + } + }; + Ok(mapped) +} + fn with_streaming_enabled(request_body: &serde_json::Value) -> serde_json::Value { let mut body = request_body.clone(); body["stream"] = serde_json::json!(true); @@ -1564,6 +2084,10 @@ impl Default for OpenAiStreamingToolCall { fn stream_from_completion_response( response: completion::CompletionResponse, + process_type: String, + provider: String, + model: String, + stream_mode: &'static str, ) -> StreamingCompletionResponse { let usage = response.usage; let message_id = response.message_id; @@ -1616,7 +2140,13 @@ fn stream_from_completion_response( })); }; - StreamingCompletionResponse::stream(Box::pin(stream)) + StreamingCompletionResponse::stream(instrument_stream_result( + Box::pin(stream), + process_type, + provider, + model, + stream_mode, + )) } async fn collect_streaming_completion_response( @@ -2722,8 +3252,104 @@ fn remap_model_name_for_api(provider: &str, model_name: &str) -> String { #[cfg(test)] mod tests { use super::*; + use crate::config::LlmConfig; + use crate::llm::manager::LlmManager; use rig::message::Message; - use std::collections::BTreeMap; + use schemars::JsonSchema; + use std::collections::{BTreeMap, HashMap}; + use std::sync::Arc; + + #[allow(dead_code)] + #[derive(Debug, serde::Deserialize, JsonSchema)] + struct TestStructuredResponse { + answer: String, + } + + #[allow(dead_code)] + #[derive(Debug, serde::Deserialize, JsonSchema)] + struct TestStructuredResponseWithOptionalField { + answer: String, + optional_note: Option, + } + + fn build_test_llm_config(provider: &str, credential: &str) -> LlmConfig { + let mut providers = HashMap::new(); + if let Some(provider_config) = crate::config::default_provider_config(provider, credential) + { + providers.insert(provider.to_string(), provider_config); + } + + LlmConfig { + anthropic_key: (provider == "anthropic").then(|| credential.to_string()), + openai_key: (provider == "openai").then(|| credential.to_string()), + openrouter_key: (provider == "openrouter").then(|| credential.to_string()), + kilo_key: (provider == "kilo").then(|| credential.to_string()), + zhipu_key: (provider == "zhipu").then(|| credential.to_string()), + groq_key: (provider == "groq").then(|| credential.to_string()), + together_key: (provider == "together").then(|| credential.to_string()), + fireworks_key: (provider == "fireworks").then(|| credential.to_string()), + deepseek_key: (provider == "deepseek").then(|| credential.to_string()), + xai_key: (provider == "xai").then(|| credential.to_string()), + mistral_key: (provider == "mistral").then(|| credential.to_string()), + gemini_key: (provider == "gemini").then(|| credential.to_string()), + ollama_key: None, + ollama_base_url: (provider == "ollama").then(|| credential.to_string()), + opencode_zen_key: (provider == "opencode-zen").then(|| credential.to_string()), + opencode_go_key: (provider == "opencode-go").then(|| credential.to_string()), + nvidia_key: (provider == "nvidia").then(|| credential.to_string()), + minimax_key: (provider == "minimax").then(|| credential.to_string()), + minimax_cn_key: (provider == "minimax-cn").then(|| credential.to_string()), + moonshot_key: (provider == "moonshot").then(|| credential.to_string()), + zai_coding_plan_key: (provider == "zai-coding-plan").then(|| credential.to_string()), + providers, + } + } + + fn structured_schema() -> schemars::Schema { + schemars::schema_for!(TestStructuredResponse) + } + + fn structured_schema_with_optional_field() -> schemars::Schema { + schemars::schema_for!(TestStructuredResponseWithOptionalField) + } + + async fn build_test_model(provider: &str, rig_alignment: RigAlignmentConfig) -> SpacebotModel { + let llm_manager = Arc::new( + LlmManager::new(build_test_llm_config(provider, "test-key")) + .await + .expect("llm manager"), + ); + SpacebotModel::make(&llm_manager, format!("{provider}/test-model")) + .with_rig_alignment(rig_alignment) + } + + fn request_with_semantics(include_tools: bool) -> CompletionRequest { + CompletionRequest { + model: None, + preamble: None, + chat_history: OneOrMany::one(Message::from("hello")), + documents: Vec::new(), + tools: if include_tools { + vec![rig::completion::ToolDefinition { + name: "worker_inspect".to_string(), + description: "inspect worker state".to_string(), + parameters: serde_json::json!({ + "type": "object", + "properties": {}, + }), + }] + } else { + Vec::new() + }, + temperature: None, + max_tokens: None, + tool_choice: Some(ToolChoice::Specific { + function_names: vec!["worker_inspect".to_string()], + }), + additional_params: None, + output_schema: Some(structured_schema()), + } + } #[test] fn reverse_map_restores_original_tool_names() { @@ -3285,4 +3911,255 @@ mod tests { assert!(msg.contains("Google")); assert!(msg.contains("invalid schema")); } + + #[test] + fn forwards_tool_choice_openai_chat_specific() { + let mapped = map_openai_tool_choice(&ToolChoice::Specific { + function_names: vec!["worker_inspect".to_string()], + }) + .expect("specific tool choice should serialize"); + + assert_eq!(mapped["type"], "function"); + assert_eq!(mapped["function"]["name"], "worker_inspect"); + } + + #[test] + fn forwards_output_schema_openai_responses() { + let mapped = map_openai_responses_text_config(&structured_schema()); + + assert_eq!(mapped["format"]["type"], "json_schema"); + assert_eq!(mapped["format"]["name"], "TestStructuredResponse"); + assert_eq!( + mapped["format"]["schema"]["additionalProperties"], + serde_json::json!(false) + ); + assert_eq!( + mapped["format"]["schema"]["required"], + serde_json::json!(["answer"]) + ); + } + + #[test] + fn openai_schema_sanitizer_preserves_optional_fields() { + let mapped = map_openai_responses_text_config(&structured_schema_with_optional_field()); + + assert_eq!( + mapped["format"]["schema"]["required"], + serde_json::json!(["answer"]) + ); + } + + #[allow(dead_code)] + #[derive(JsonSchema)] + #[schemars( + title = "Test Structured Response (Invalid Chars!) + Extra Long Name Segment 1234567890" + )] + struct InvalidNamedStructuredResponse { + answer: String, + } + + fn invalid_named_schema() -> schemars::Schema { + schemars::schema_for!(InvalidNamedStructuredResponse) + } + + #[allow(dead_code)] + #[derive(JsonSchema)] + #[schemars(title = "!!! ### ???")] + struct EmptyNormalizedSchemaNameResponse { + answer: String, + } + + fn empty_normalized_name_schema() -> schemars::Schema { + schemars::schema_for!(EmptyNormalizedSchemaNameResponse) + } + + #[test] + fn schema_name_normalizes_to_provider_safe_and_truncates() { + let mapped = map_openai_response_format(&invalid_named_schema()); + let name = mapped["json_schema"]["name"] + .as_str() + .expect("schema name should be serialized"); + + assert!(name.len() <= 64); + assert!( + name.chars() + .all(|character| character.is_ascii_alphanumeric() + || character == '_' + || character == '-') + ); + } + + #[test] + fn schema_name_falls_back_when_title_normalizes_empty() { + let mapped = map_openai_response_format(&empty_normalized_name_schema()); + assert_eq!(mapped["json_schema"]["name"], "response_schema"); + } + + #[test] + fn sanitize_openai_schema_keeps_defs_with_refs() { + let mut schema = serde_json::json!({ + "$ref": "#/$defs/Result", + "$defs": { + "Result": { + "type": "object", + "properties": { + "answer": { "type": "string" } + } + } + }, + "description": "discard me" + }); + + sanitize_openai_schema(&mut schema); + + assert_eq!(schema["$ref"], "#/$defs/Result"); + assert!(schema.get("$defs").is_some()); + assert_eq!( + schema["$defs"]["Result"]["additionalProperties"], + serde_json::json!(false) + ); + assert!(schema.get("description").is_none()); + } + + #[tokio::test] + async fn rejects_unsupported_semantics_enforced() { + let model = build_test_model( + "gemini", + RigAlignmentConfig { + request_semantics_mode: RigSemanticsMode::Enforced, + ..RigAlignmentConfig::default() + }, + ) + .await; + let mut request = request_with_semantics(false); + request.tool_choice = None; + + let error = model + .evaluate_request_semantics(&request) + .expect_err("unsupported provider should fail in enforced mode"); + assert!(error.to_string().contains("output_schema")); + } + + #[tokio::test] + async fn rig_parity_off_mode_suppresses_requested_semantics() { + let model = build_test_model( + "openai", + RigAlignmentConfig { + request_semantics_mode: RigSemanticsMode::Off, + tool_choice_provider_allowlist: vec!["openai".to_string()], + output_schema_provider_allowlist: vec!["openai".to_string()], + ..RigAlignmentConfig::default() + }, + ) + .await; + + let decision = model + .evaluate_request_semantics(&request_with_semantics(true)) + .expect("off mode should not error"); + + assert!(!decision.apply_tool_choice); + assert!(!decision.apply_output_schema); + } + + #[tokio::test] + async fn rig_parity_shadow_mode_applies_supported_provider_and_shadows_unsupported_one() { + let supported_model = build_test_model( + "openai", + RigAlignmentConfig { + request_semantics_mode: RigSemanticsMode::Shadow, + tool_choice_provider_allowlist: vec!["openai".to_string()], + output_schema_provider_allowlist: vec!["openai".to_string()], + ..RigAlignmentConfig::default() + }, + ) + .await; + let supported = supported_model + .evaluate_request_semantics(&request_with_semantics(true)) + .expect("supported provider should not error in shadow mode"); + assert!(supported.apply_tool_choice); + assert!(supported.apply_output_schema); + + let unsupported_model = build_test_model( + "gemini", + RigAlignmentConfig { + request_semantics_mode: RigSemanticsMode::Shadow, + tool_choice_provider_allowlist: vec!["openai".to_string()], + output_schema_provider_allowlist: vec!["openai".to_string()], + ..RigAlignmentConfig::default() + }, + ) + .await; + let unsupported = unsupported_model + .evaluate_request_semantics(&request_with_semantics(true)) + .expect("shadow mode should downgrade unsupported semantics"); + assert!(!unsupported.apply_tool_choice); + assert!(!unsupported.apply_output_schema); + } + + #[tokio::test] + async fn rig_parity_enforced_rejects_tool_choice_without_tool_definitions() { + let model = build_test_model( + "openai", + RigAlignmentConfig { + request_semantics_mode: RigSemanticsMode::Enforced, + tool_choice_provider_allowlist: vec!["openai".to_string()], + output_schema_provider_allowlist: vec!["openai".to_string()], + ..RigAlignmentConfig::default() + }, + ) + .await; + + let error = model + .evaluate_request_semantics(&request_with_semantics(false)) + .expect_err("tool_choice without tool definitions should fail in enforced mode"); + assert!(error.to_string().contains("tool_choice")); + } + + #[tokio::test] + async fn rig_parity_enforced_allows_auto_tool_choice_without_tool_definitions() { + let model = build_test_model( + "openai", + RigAlignmentConfig { + request_semantics_mode: RigSemanticsMode::Enforced, + tool_choice_provider_allowlist: vec!["openai".to_string()], + output_schema_provider_allowlist: vec!["openai".to_string()], + ..RigAlignmentConfig::default() + }, + ) + .await; + + let mut request = request_with_semantics(false); + request.tool_choice = Some(ToolChoice::Auto); + request.output_schema = None; + + let decision = model + .evaluate_request_semantics(&request) + .expect("auto tool choice should not require tool definitions"); + assert!(decision.apply_tool_choice); + assert!(!decision.apply_output_schema); + } + + #[tokio::test] + async fn rig_parity_enforced_allows_none_tool_choice_without_tool_definitions() { + let model = build_test_model( + "openai", + RigAlignmentConfig { + request_semantics_mode: RigSemanticsMode::Enforced, + tool_choice_provider_allowlist: vec!["openai".to_string()], + output_schema_provider_allowlist: vec!["openai".to_string()], + ..RigAlignmentConfig::default() + }, + ) + .await; + + let mut request = request_with_semantics(false); + request.tool_choice = Some(ToolChoice::None); + request.output_schema = None; + + let decision = model + .evaluate_request_semantics(&request) + .expect("none tool choice should not require tool definitions"); + assert!(decision.apply_tool_choice); + assert!(!decision.apply_output_schema); + } } diff --git a/src/main.rs b/src/main.rs index ecd26328c..be6a89d15 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1882,7 +1882,7 @@ async fn run( let (response_tx, mut response_rx) = mpsc::channel::(32); // Subscribe to the agent's event bus - let event_rx = agent.deps.event_tx.subscribe(); + let channel_event_rx = agent.deps.event_tx.subscribe(); let channel_id: spacebot::ChannelId = Arc::from(conversation_id.as_str()); @@ -1890,7 +1890,7 @@ async fn run( channel_id, agent.deps.clone(), response_tx, - event_rx, + channel_event_rx, agent.config.screenshot_dir(), agent.config.logs_dir(), ); @@ -1977,68 +1977,165 @@ async fn run( let sse_agent_id = agent_id.to_string(); let sse_channel_id = conversation_id.clone(); let outbound_handle = tokio::spawn(async move { - while let Some(response) = response_rx.recv().await { - // Forward relevant events to SSE clients - match &response { - spacebot::OutboundResponse::Text(text) => { - api_event_tx.send(spacebot::api::ApiEvent::OutboundMessage { - agent_id: sse_agent_id.clone(), - channel_id: sse_channel_id.clone(), - text: text.clone(), - }).ok(); - } - spacebot::OutboundResponse::RichMessage { text, .. } => { - api_event_tx.send(spacebot::api::ApiEvent::OutboundMessage { - agent_id: sse_agent_id.clone(), - channel_id: sse_channel_id.clone(), - text: text.clone(), - }).ok(); - } - spacebot::OutboundResponse::ThreadReply { text, .. } => { - api_event_tx.send(spacebot::api::ApiEvent::OutboundMessage { - agent_id: sse_agent_id.clone(), - channel_id: sse_channel_id.clone(), - text: text.clone(), - }).ok(); - } - spacebot::OutboundResponse::Status(spacebot::StatusUpdate::Thinking) => { - api_event_tx.send(spacebot::api::ApiEvent::TypingState { - agent_id: sse_agent_id.clone(), - channel_id: sse_channel_id.clone(), - is_typing: true, - }).ok(); - } - spacebot::OutboundResponse::Status(spacebot::StatusUpdate::StopTyping) => { - api_event_tx.send(spacebot::api::ApiEvent::TypingState { - agent_id: sse_agent_id.clone(), - channel_id: sse_channel_id.clone(), - is_typing: false, - }).ok(); - } - _ => {} - } - - let current_message = outbound_message.read().await.clone(); - - match response { - spacebot::OutboundResponse::Status(status) => { - if let Err(error) = messaging_for_outbound - .send_status(¤t_message, status) - .await - { - tracing::warn!(%error, "failed to send status update"); + let mut stream_active = false; + let mut stream_message: Option = None; + let mut last_stream_text = String::new(); + + loop { + tokio::select! { + response = response_rx.recv() => { + let Some(response) = response else { + if let Some(cleanup) = pending_stream_cleanup_response( + stream_active, + &last_stream_text, + ) { + if let Some(api_event) = api_event_for_stream_cleanup( + &sse_agent_id, + &sse_channel_id, + &cleanup, + ) { + api_event_tx.send(api_event).ok(); + } + let current_message = match stream_message.as_ref() { + Some(message) => message.clone(), + None => outbound_message.read().await.clone(), + }; + if let Err(error) = messaging_for_outbound + .respond(¤t_message, cleanup) + .await + { + tracing::warn!(%error, "failed to clean up streamed response after outbound channel closed"); + } + } + last_stream_text.clear(); + break; + }; + if let Some(api_event) = api_event_for_outbound_response( + &sse_agent_id, + &sse_channel_id, + &response, + ) { + api_event_tx.send(api_event).ok(); } - } - response => { - tracing::info!( - conversation_id = %outbound_conversation_id, - "routing outbound response to messaging adapter" - ); - if let Err(error) = messaging_for_outbound - .respond(¤t_message, response) - .await - { - tracing::error!(%error, "failed to send outbound response"); + + match response { + spacebot::OutboundResponse::Status(status) => { + if matches!(status, spacebot::StatusUpdate::StopTyping) { + if let Some(cleanup) = pending_stream_cleanup_response( + stream_active, + &last_stream_text, + ) { + if let Some(api_event) = api_event_for_stream_cleanup( + &sse_agent_id, + &sse_channel_id, + &cleanup, + ) { + api_event_tx.send(api_event).ok(); + } + let current_message = match stream_message.as_ref() { + Some(message) => message.clone(), + None => outbound_message.read().await.clone(), + }; + if let Err(error) = messaging_for_outbound + .respond(¤t_message, cleanup) + .await + { + tracing::warn!(%error, "failed to finalize streamed response from buffered text"); + } + } + stream_active = false; + stream_message = None; + last_stream_text.clear(); + } + + let current_message = outbound_message.read().await.clone(); + if let Err(error) = messaging_for_outbound + .send_status(¤t_message, status) + .await + { + tracing::warn!(%error, "failed to send status update"); + } + } + response @ ( + spacebot::OutboundResponse::Text(_) + | spacebot::OutboundResponse::RichMessage { .. } + | spacebot::OutboundResponse::ThreadReply { .. } + ) => { + tracing::info!( + conversation_id = %outbound_conversation_id, + terminal_reason = outbound_terminal_reason(&response), + "routing outbound response to messaging adapter" + ); + let current_message = match stream_message.as_ref() { + Some(message) if stream_active => message.clone(), + _ => outbound_message.read().await.clone(), + }; + if stream_active + && should_end_stream_before_terminal_response(&response) + && let Err(error) = messaging_for_outbound + .respond(¤t_message, spacebot::OutboundResponse::StreamEnd) + .await + { + tracing::warn!(%error, "failed to clear streamed placeholder before terminal response"); + } + if let Err(error) = messaging_for_outbound + .respond(¤t_message, response) + .await + { + tracing::error!(%error, "failed to send outbound response"); + } + stream_active = false; + stream_message = None; + last_stream_text.clear(); + } + spacebot::OutboundResponse::StreamChunk(text) => { + if text.is_empty() { + continue; + } + + last_stream_text.push_str(&text); + let current_message = match stream_message.as_ref() { + Some(message) => message.clone(), + None => { + let message = outbound_message.read().await.clone(); + stream_message = Some(message.clone()); + message + } + }; + if !stream_active { + if let Err(error) = messaging_for_outbound + .respond(¤t_message, spacebot::OutboundResponse::StreamStart) + .await + { + tracing::warn!(%error, "failed to start streamed response"); + continue; + } + stream_active = true; + } + + if let Err(error) = messaging_for_outbound + .respond( + ¤t_message, + spacebot::OutboundResponse::StreamChunk(text), + ) + .await + { + tracing::warn!(%error, "failed to send streamed response chunk"); + } + } + response => { + tracing::info!( + conversation_id = %outbound_conversation_id, + "routing outbound response to messaging adapter" + ); + let current_message = outbound_message.read().await.clone(); + if let Err(error) = messaging_for_outbound + .respond(¤t_message, response) + .await + { + tracing::error!(%error, "failed to send outbound response"); + } + } } } } @@ -3149,9 +3246,115 @@ async fn initialize_agents( Ok(()) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PendingStreamAction { + None, + FinalizeText, + EndStream, +} + +fn pending_stream_action(stream_active: bool, last_stream_text: &str) -> PendingStreamAction { + if !last_stream_text.trim().is_empty() { + PendingStreamAction::FinalizeText + } else if stream_active { + PendingStreamAction::EndStream + } else { + PendingStreamAction::None + } +} + +fn pending_stream_cleanup_response( + stream_active: bool, + last_stream_text: &str, +) -> Option { + match pending_stream_action(stream_active, last_stream_text) { + PendingStreamAction::FinalizeText => Some(spacebot::OutboundResponse::Text( + last_stream_text.to_string(), + )), + PendingStreamAction::EndStream => Some(spacebot::OutboundResponse::StreamEnd), + PendingStreamAction::None => None, + } +} + +fn api_event_for_outbound_response( + agent_id: &str, + channel_id: &str, + response: &spacebot::OutboundResponse, +) -> Option { + match response { + spacebot::OutboundResponse::Text(text) + | spacebot::OutboundResponse::RichMessage { text, .. } + | spacebot::OutboundResponse::ThreadReply { text, .. } => { + Some(spacebot::api::ApiEvent::OutboundMessage { + agent_id: agent_id.to_string(), + channel_id: channel_id.to_string(), + text: text.clone(), + }) + } + spacebot::OutboundResponse::Status(spacebot::StatusUpdate::Thinking) => { + Some(spacebot::api::ApiEvent::TypingState { + agent_id: agent_id.to_string(), + channel_id: channel_id.to_string(), + is_typing: true, + }) + } + spacebot::OutboundResponse::Status(spacebot::StatusUpdate::StopTyping) => { + Some(spacebot::api::ApiEvent::TypingState { + agent_id: agent_id.to_string(), + channel_id: channel_id.to_string(), + is_typing: false, + }) + } + _ => None, + } +} + +fn api_event_for_stream_cleanup( + agent_id: &str, + channel_id: &str, + cleanup: &spacebot::OutboundResponse, +) -> Option { + match cleanup { + spacebot::OutboundResponse::StreamEnd => Some(spacebot::api::ApiEvent::OutboundStreamEnd { + agent_id: agent_id.to_string(), + channel_id: channel_id.to_string(), + }), + _ => api_event_for_outbound_response(agent_id, channel_id, cleanup), + } +} + +fn should_end_stream_before_terminal_response(response: &spacebot::OutboundResponse) -> bool { + matches!( + response, + spacebot::OutboundResponse::RichMessage { .. } + | spacebot::OutboundResponse::ThreadReply { .. } + ) +} + +fn outbound_terminal_reason(response: &spacebot::OutboundResponse) -> &'static str { + match response { + spacebot::OutboundResponse::Text(_) => "text", + spacebot::OutboundResponse::File { .. } => "file", + spacebot::OutboundResponse::Reaction(_) => "reaction", + spacebot::OutboundResponse::RemoveReaction(_) => "remove_reaction", + spacebot::OutboundResponse::Ephemeral { .. } => "ephemeral", + spacebot::OutboundResponse::RichMessage { .. } => "rich_message", + spacebot::OutboundResponse::ScheduledMessage { .. } => "scheduled_message", + spacebot::OutboundResponse::ThreadReply { .. } => "thread_reply", + spacebot::OutboundResponse::Status(_) => "status", + spacebot::OutboundResponse::StreamStart => "stream_start", + spacebot::OutboundResponse::StreamChunk(_) => "stream_chunk", + spacebot::OutboundResponse::StreamEnd => "stream_end", + } +} + #[cfg(test)] mod tests { - use super::wait_for_startup_warmup_tasks; + use super::{ + PendingStreamAction, api_event_for_stream_cleanup, outbound_terminal_reason, + pending_stream_action, pending_stream_cleanup_response, + should_end_stream_before_terminal_response, wait_for_startup_warmup_tasks, + }; use std::future::pending; use std::sync::Arc; use std::time::Duration; @@ -3214,4 +3417,107 @@ mod tests { "startup warmup timeout should return without waiting for non-cooperative task" ); } + + #[test] + fn pending_stream_action_prefers_final_text_even_without_active_stream() { + assert_eq!( + pending_stream_action(false, "final text from failed stream start"), + PendingStreamAction::FinalizeText + ); + assert_eq!( + pending_stream_action(true, " "), + PendingStreamAction::EndStream + ); + assert_eq!( + pending_stream_action(false, " "), + PendingStreamAction::None + ); + } + + #[test] + fn pending_stream_cleanup_response_prefers_buffered_text_on_shutdown() { + assert!(matches!( + pending_stream_cleanup_response(false, "buffered final text"), + Some(spacebot::OutboundResponse::Text(text)) if text == "buffered final text" + )); + assert!(matches!( + pending_stream_cleanup_response(true, " "), + Some(spacebot::OutboundResponse::StreamEnd) + )); + assert!(pending_stream_cleanup_response(false, " ").is_none()); + } + + #[test] + fn stream_cleanup_api_event_uses_terminal_contract() { + assert!(matches!( + api_event_for_stream_cleanup( + "agent-1", + "channel-1", + &spacebot::OutboundResponse::Text("done".to_string()), + ), + Some(spacebot::api::ApiEvent::OutboundMessage { + agent_id, + channel_id, + text, + }) if agent_id == "agent-1" && channel_id == "channel-1" && text == "done" + )); + assert!(matches!( + api_event_for_stream_cleanup( + "agent-1", + "channel-1", + &spacebot::OutboundResponse::StreamEnd, + ), + Some(spacebot::api::ApiEvent::OutboundStreamEnd { + agent_id, + channel_id, + }) if agent_id == "agent-1" && channel_id == "channel-1" + )); + } + + #[test] + fn terminal_non_text_responses_require_stream_cleanup() { + assert!(!should_end_stream_before_terminal_response( + &spacebot::OutboundResponse::Text("done".to_string()), + )); + assert!(should_end_stream_before_terminal_response( + &spacebot::OutboundResponse::ThreadReply { + thread_name: "status".to_string(), + text: "done".to_string(), + }, + )); + assert!(should_end_stream_before_terminal_response( + &spacebot::OutboundResponse::RichMessage { + text: "done".to_string(), + blocks: Vec::new(), + cards: Vec::new(), + interactive_elements: Vec::new(), + poll: None, + }, + )); + } + + #[test] + fn outbound_terminal_reason_matches_terminal_variants() { + assert_eq!( + outbound_terminal_reason(&spacebot::OutboundResponse::Text("done".to_string())), + "text" + ); + assert_eq!( + outbound_terminal_reason(&spacebot::OutboundResponse::ThreadReply { + thread_name: "status".to_string(), + text: "done".to_string(), + }), + "thread_reply" + ); + assert_eq!( + outbound_terminal_reason(&spacebot::OutboundResponse::RichMessage { + text: "done".to_string(), + blocks: Vec::new(), + cards: Vec::new(), + interactive_elements: Vec::new(), + poll: None, + }), + "rich_message" + ); + } } diff --git a/src/messaging/discord.rs b/src/messaging/discord.rs index afd8c1c87..2a8f576fc 100644 --- a/src/messaging/discord.rs +++ b/src/messaging/discord.rs @@ -17,8 +17,38 @@ use serenity::all::{ }; use std::collections::HashMap; use std::sync::Arc; +use std::time::{Duration, Instant}; use tokio::sync::{RwLock, mpsc}; +#[derive(Clone)] +struct ActiveStream { + message_id: serenity::all::MessageId, + last_edit: Instant, + text: String, +} + +const STREAM_EDIT_INTERVAL: Duration = Duration::from_secs(1); + +fn should_throttle_stream_edit(last_edit: Instant, now: Instant) -> bool { + now.duration_since(last_edit) < STREAM_EDIT_INTERVAL +} + +async fn with_removed_active_stream( + active_messages: &RwLock>, + message_id: &str, + cleanup: F, +) -> Option +where + F: FnOnce(ActiveStream) -> Fut, + Fut: std::future::Future, +{ + let stream = { + let mut active_messages = active_messages.write().await; + active_messages.remove(message_id) + }?; + Some(cleanup(stream).await) +} + /// Discord adapter state. pub struct DiscordAdapter { runtime_key: String, @@ -26,8 +56,8 @@ pub struct DiscordAdapter { permissions: Arc>, http: Arc>>>, bot_user_id: Arc>>, - /// Maps InboundMessage.id to the Discord MessageId being edited during streaming. - active_messages: Arc>>, + /// Maps InboundMessage.id to the Discord streaming placeholder state. + active_messages: Arc>>, /// Typing handles per message. Typing stops when the handle is dropped. typing_tasks: Arc>>, shard_manager: Arc>>>, @@ -149,18 +179,69 @@ impl Messaging for DiscordAdapter { OutboundResponse::Text(text) => { self.stop_typing(message).await; let reply_to = Self::extract_reply_message_id(message); - - for (index, chunk) in split_message(&text, 2000).into_iter().enumerate() { - let mut builder = CreateMessage::new().content(chunk); - if index == 0 - && let Some(reply_message_id) = reply_to + let chunks = split_message(&text, 2000); + let active_placeholder = + self.active_messages.read().await.get(&message.id).cloned(); + + if let Some(message_id) = active_placeholder { + let first_chunk = chunks + .first() + .cloned() + .unwrap_or_else(|| "\u{200B}".to_string()); + let builder = EditMessage::new().content(first_chunk); + if let Err(error) = channel_id + .edit_message(&*http, message_id.message_id, builder) + .await { - builder = builder.reference_message((channel_id, reply_message_id)); + tracing::warn!( + %error, + "failed to finalize streaming message, falling back to normal send" + ); + if let Err(delete_error) = channel_id + .delete_message(&*http, message_id.message_id) + .await + { + tracing::warn!( + %delete_error, + "failed to delete discord stream placeholder after finalize fallback" + ); + } + self.active_messages.write().await.remove(&message.id); + for (index, chunk) in chunks.iter().enumerate() { + let mut builder = CreateMessage::new().content(chunk.clone()); + if index == 0 + && let Some(reply_message_id) = reply_to + { + builder = builder.reference_message((channel_id, reply_message_id)); + } + channel_id + .send_message(&*http, builder) + .await + .context("failed to send discord message")?; + } + } else { + self.active_messages.write().await.remove(&message.id); + for chunk in chunks.iter().skip(1) { + let builder = CreateMessage::new().content(chunk.clone()); + channel_id + .send_message(&*http, builder) + .await + .context("failed to send discord message")?; + } + } + } else { + for (index, chunk) in chunks.into_iter().enumerate() { + let mut builder = CreateMessage::new().content(chunk); + if index == 0 + && let Some(reply_message_id) = reply_to + { + builder = builder.reference_message((channel_id, reply_message_id)); + } + channel_id + .send_message(&*http, builder) + .await + .context("failed to send discord message")?; } - channel_id - .send_message(&*http, builder) - .await - .context("failed to send discord message")?; } } OutboundResponse::RichMessage { @@ -334,26 +415,49 @@ impl Messaging for DiscordAdapter { } OutboundResponse::StreamStart => { self.stop_typing(message).await; + let reply_to = Self::extract_reply_message_id(message); + let mut builder = CreateMessage::new().content("\u{200B}"); + if let Some(reply_message_id) = reply_to { + builder = builder.reference_message((channel_id, reply_message_id)); + } let placeholder = channel_id - .say(&*http, "\u{200B}") + .send_message(&*http, builder) .await .context("failed to send stream placeholder")?; - self.active_messages - .write() - .await - .insert(message.id.clone(), placeholder.id); + self.active_messages.write().await.insert( + message.id.clone(), + ActiveStream { + message_id: placeholder.id, + last_edit: Instant::now(), + text: String::new(), + }, + ); } OutboundResponse::StreamChunk(text) => { - let active = self.active_messages.read().await; - if let Some(&message_id) = active.get(&message.id) { - let display_text = if text.len() > 2000 { - let end = text.floor_char_boundary(1997); - format!("{}...", &text[..end]) + let now = Instant::now(); + let update = { + let mut active = self.active_messages.write().await; + if let Some(stream) = active.get_mut(&message.id) { + stream.text.push_str(&text); + if should_throttle_stream_edit(stream.last_edit, now) { + return Ok(()); + } + stream.last_edit = now; + let display_text = if stream.text.len() > 2000 { + let end = stream.text.floor_char_boundary(1997); + format!("{}...", &stream.text[..end]) + } else { + stream.text.clone() + }; + Some((stream.message_id, display_text)) } else { - text - }; + None + } + }; + + if let Some((message_id, display_text)) = update { let builder = EditMessage::new().content(display_text); if let Err(error) = channel_id.edit_message(&*http, message_id, builder).await { tracing::warn!(%error, "failed to edit streaming message"); @@ -361,7 +465,18 @@ impl Messaging for DiscordAdapter { } } OutboundResponse::StreamEnd => { - self.active_messages.write().await.remove(&message.id); + with_removed_active_stream( + &self.active_messages, + &message.id, + |stream| async move { + if let Err(error) = + channel_id.delete_message(&*http, stream.message_id).await + { + tracing::warn!(%error, "failed to delete discord stream placeholder"); + } + }, + ) + .await; } OutboundResponse::Status(status) => { self.send_status(message, status).await?; @@ -1179,6 +1294,8 @@ fn build_poll( mod tests { use super::*; use crate::{Button, ButtonStyle, Card, CardField, InteractiveElements, Poll}; + use std::sync::Arc; + use tokio::sync::oneshot; #[test] fn test_build_embed_limits() { @@ -1236,4 +1353,72 @@ mod tests { let _ = build_poll(&poll); // Again, can't easily inspect CreatePoll fields, but we verify it runs. } + + #[test] + fn adapter_streaming_throttle() { + let last_edit = Instant::now(); + let before_interval = last_edit + STREAM_EDIT_INTERVAL - Duration::from_millis(1); + let at_interval = last_edit + STREAM_EDIT_INTERVAL; + + assert!( + should_throttle_stream_edit(last_edit, before_interval), + "stream chunk edits should be skipped before the configured throttle interval" + ); + assert!( + !should_throttle_stream_edit(last_edit, at_interval), + "stream chunk edits should be allowed once the configured throttle interval elapses" + ); + } + + #[tokio::test] + async fn stream_end_releases_active_messages_lock_before_delete() { + let active_messages = Arc::new(RwLock::new(HashMap::from([( + "message-id".to_string(), + ActiveStream { + message_id: MessageId::new(1), + last_edit: Instant::now(), + text: "partial".to_string(), + }, + )]))); + let (delete_started_tx, delete_started_rx) = oneshot::channel(); + let (allow_delete_tx, allow_delete_rx) = oneshot::channel::<()>(); + + let cleanup_task = { + let active_messages = active_messages.clone(); + tokio::spawn(async move { + with_removed_active_stream( + active_messages.as_ref(), + "message-id", + |stream| async move { + delete_started_tx.send(stream.message_id).ok(); + allow_delete_rx.await.ok(); + }, + ) + .await + }) + }; + + delete_started_rx + .await + .expect("cleanup future should begin after the stream entry is removed"); + + let active_messages_guard = tokio::time::timeout( + Duration::from_millis(100), + active_messages.read(), + ) + .await + .expect("stream cleanup should not hold the active_messages lock while delete is pending"); + assert!( + !active_messages_guard.contains_key("message-id"), + "the stream entry should be removed before the delete future completes" + ); + drop(active_messages_guard); + + allow_delete_tx + .send(()) + .expect("test should release the simulated delete"); + cleanup_task + .await + .expect("cleanup task should complete cleanly"); + } } diff --git a/src/messaging/slack.rs b/src/messaging/slack.rs index 44e67997a..39e15eb9b 100644 --- a/src/messaging/slack.rs +++ b/src/messaging/slack.rs @@ -32,7 +32,7 @@ use slack_morphism::prelude::*; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::{RwLock, mpsc}; -use tokio::time::{Duration, timeout}; +use tokio::time::{Duration, Instant, timeout}; /// State shared with socket mode callbacks via `SlackClientEventsUserState`. struct SlackAdapterState { @@ -56,6 +56,39 @@ struct SlackUserIdentity { username: Option, } +#[derive(Debug, Clone)] +struct ActiveStream { + ts: String, + last_edit: Instant, + text: String, +} + +const STREAM_EDIT_INTERVAL: Duration = Duration::from_secs(1); +const STREAM_TEXT_BUFFER_LIMIT: usize = 12_000; + +fn initial_stream_last_edit() -> Instant { + Instant::now() + .checked_sub(STREAM_EDIT_INTERVAL) + .unwrap_or_else(Instant::now) +} + +fn should_throttle_stream_edit(last_edit: Instant, now: Instant) -> bool { + now.duration_since(last_edit) < STREAM_EDIT_INTERVAL +} + +fn cap_stream_text_buffer(text: &mut String) -> bool { + if text.len() <= STREAM_TEXT_BUFFER_LIMIT { + return false; + } + + let mut end = STREAM_TEXT_BUFFER_LIMIT; + while end > 0 && !text.is_char_boundary(end) { + end -= 1; + } + text.truncate(end); + true +} + /// Slack adapter. pub struct SlackAdapter { runtime_key: String, @@ -68,8 +101,8 @@ pub struct SlackAdapter { client: Arc, /// Pre-built API token wrapping `bot_token`. Created once alongside `client`. token: SlackApiToken, - /// Maps InboundMessage.id → Slack ts for streaming edits. - active_messages: Arc>>, + /// Maps InboundMessage.id → Slack streaming placeholder state. + active_messages: Arc>>, shutdown_tx: Arc>>>, /// Slash command routing: command string → agent_id. commands: Arc>, @@ -922,17 +955,74 @@ impl Messaging for SlackAdapter { match response { OutboundResponse::Text(text) => { let thread_ts = extract_thread_ts(message); - - for chunk in split_message(&text, 12_000) { - let mut req = SlackApiChatPostMessageRequest::new( + let chunks = split_message(&text, 12_000); + let active_stream = self.active_messages.read().await.get(&message.id).cloned(); + + if let Some(active_stream) = active_stream { + let first_chunk = chunks + .first() + .cloned() + .unwrap_or_else(|| "\u{200B}".to_string()); + let update = SlackApiChatUpdateRequest::new( channel_id.clone(), - markdown_content(chunk), + markdown_content(first_chunk), + SlackTs(active_stream.ts.clone()), ); - req = req.opt_thread_ts(thread_ts.clone()); - session - .chat_post_message(&req) - .await - .context("failed to send slack message")?; + + if let Err(error) = session.chat_update(&update).await { + tracing::warn!( + %error, + "failed to finalize streaming message, falling back to normal send" + ); + let delete = SlackApiChatDeleteRequest { + channel: channel_id.clone(), + ts: SlackTs(active_stream.ts.clone()), + as_user: None, + }; + if let Err(delete_error) = session.chat_delete(&delete).await { + tracing::warn!( + %delete_error, + "failed to delete slack stream placeholder after finalize fallback" + ); + } + self.active_messages.write().await.remove(&message.id); + for chunk in &chunks { + let mut req = SlackApiChatPostMessageRequest::new( + channel_id.clone(), + markdown_content(chunk.clone()), + ); + req = req.opt_thread_ts(thread_ts.clone()); + session + .chat_post_message(&req) + .await + .context("failed to send slack message")?; + } + } else { + self.active_messages.write().await.remove(&message.id); + for chunk in chunks.iter().skip(1) { + let mut req = SlackApiChatPostMessageRequest::new( + channel_id.clone(), + markdown_content(chunk.clone()), + ); + req = req.opt_thread_ts(thread_ts.clone()); + session + .chat_post_message(&req) + .await + .context("failed to send slack message")?; + } + } + } else { + for chunk in chunks { + let mut req = SlackApiChatPostMessageRequest::new( + channel_id.clone(), + markdown_content(chunk), + ); + req = req.opt_thread_ts(thread_ts.clone()); + session + .chat_post_message(&req) + .await + .context("failed to send slack message")?; + } } } OutboundResponse::ThreadReply { @@ -1087,42 +1177,73 @@ impl Messaging for SlackAdapter { } OutboundResponse::StreamStart => { + let thread_ts = extract_thread_ts(message); let req = SlackApiChatPostMessageRequest::new( channel_id.clone(), SlackMessageContent::new().with_text("\u{200B}".into()), - ); + ) + .opt_thread_ts(thread_ts); let resp = session .chat_post_message(&req) .await .context("failed to send stream placeholder")?; - self.active_messages - .write() - .await - .insert(message.id.clone(), resp.ts.0); + self.active_messages.write().await.insert( + message.id.clone(), + ActiveStream { + ts: resp.ts.0, + last_edit: initial_stream_last_edit(), + text: String::new(), + }, + ); } OutboundResponse::StreamChunk(text) => { - let active = self.active_messages.read().await; - if let Some(ts) = active.get(&message.id) { - let display_text = if text.len() > 12_000 { - let end = text.floor_char_boundary(11_997); - format!("{}...", &text[..end]) + let now = Instant::now(); + let update = { + let mut active = self.active_messages.write().await; + if let Some(stream) = active.get_mut(&message.id) { + stream.text.push_str(&text); + let truncated = cap_stream_text_buffer(&mut stream.text); + if should_throttle_stream_edit(stream.last_edit, now) { + return Ok(()); + } + stream.last_edit = now; + let display_text = if truncated { + format!("{}...", stream.text) + } else { + stream.text.clone() + }; + Some((stream.ts.clone(), display_text)) } else { - text - }; + None + } + }; + + if let Some((stream_ts, display_text)) = update { let req = SlackApiChatUpdateRequest::new( channel_id.clone(), markdown_content(display_text), - SlackTs(ts.clone()), + SlackTs(stream_ts), ); if let Err(error) = session.chat_update(&req).await { tracing::warn!(%error, "failed to edit streaming message"); } + } else { + return Ok(()); } } OutboundResponse::StreamEnd => { - self.active_messages.write().await.remove(&message.id); + if let Some(stream) = self.active_messages.write().await.remove(&message.id) { + let req = SlackApiChatDeleteRequest { + channel: channel_id.clone(), + ts: SlackTs(stream.ts), + as_user: None, + }; + if let Err(error) = session.chat_delete(&req).await { + tracing::warn!(%error, "failed to delete slack stream placeholder"); + } + } } OutboundResponse::Status(_) => { @@ -1702,6 +1823,7 @@ fn resolve_slack_user_identity(user: &SlackUser, user_id: &str) -> SlackUserIden #[cfg(test)] mod tests { use super::*; + use std::time::Duration; #[test] fn sanitize_reaction_name_unicode_emoji_with_shortcode() { @@ -1774,4 +1896,36 @@ mod tests { let result = sanitize_reaction_name(":partyparrot:"); assert_eq!(result, "partyparrot"); } + + #[test] + fn adapter_streaming_throttle() { + let last_edit = Instant::now(); + let before_interval = last_edit + STREAM_EDIT_INTERVAL - Duration::from_millis(1); + let at_interval = last_edit + STREAM_EDIT_INTERVAL; + + assert!( + should_throttle_stream_edit(last_edit, before_interval), + "stream chunk edits should be skipped before the configured throttle interval" + ); + assert!( + !should_throttle_stream_edit(last_edit, at_interval), + "stream chunk edits should be allowed once the configured throttle interval elapses" + ); + assert!( + !should_throttle_stream_edit(initial_stream_last_edit(), Instant::now()), + "the first streamed edit should not be throttled" + ); + } + + #[test] + fn stream_text_buffer_is_capped() { + let mut text = "a".repeat(STREAM_TEXT_BUFFER_LIMIT + 128); + let truncated = cap_stream_text_buffer(&mut text); + assert!(truncated, "oversized buffers should report truncation"); + assert_eq!( + text.len(), + STREAM_TEXT_BUFFER_LIMIT, + "stream display buffers should stay bounded" + ); + } } diff --git a/src/messaging/telegram.rs b/src/messaging/telegram.rs index f806e7537..eeb781639 100644 --- a/src/messaging/telegram.rs +++ b/src/messaging/telegram.rs @@ -41,10 +41,12 @@ pub struct TelegramAdapter { } /// Tracks an in-progress streaming message edit. +#[derive(Clone)] struct ActiveStream { chat_id: ChatId, message_id: MessageId, last_edit: Instant, + text: String, } /// Telegram's per-message character limit. @@ -56,6 +58,22 @@ const FORMATTED_SPLIT_LENGTH: usize = MAX_MESSAGE_LENGTH / 2; /// Minimum interval between streaming edits to avoid rate limits. const STREAM_EDIT_INTERVAL: std::time::Duration = std::time::Duration::from_millis(1000); +async fn with_removed_active_stream( + active_messages: &RwLock>, + conversation_id: &str, + cleanup: F, +) -> Option +where + F: FnOnce(ActiveStream) -> Fut, + Fut: std::future::Future, +{ + let stream = { + let mut active_messages = active_messages.write().await; + active_messages.remove(conversation_id) + }?; + Some(cleanup(stream).await) +} + impl TelegramAdapter { pub fn new( runtime_key: impl Into, @@ -303,7 +321,71 @@ impl Messaging for TelegramAdapter { match response { OutboundResponse::Text(text) => { self.stop_typing(&message.conversation_id).await; - send_formatted(&self.bot, chat_id, &text, None).await?; + let chunks = split_message(&text, MAX_MESSAGE_LENGTH); + let active_stream = self + .active_messages + .read() + .await + .get(&message.conversation_id) + .cloned(); + let reply_to = self.extract_message_id(message).ok(); + + if let Some(stream) = active_stream { + let (first_chunk, remaining_chunks) = stream_finalization_chunks(&chunks); + let html = markdown_to_telegram_html(first_chunk); + let finalized: std::result::Result<(), RequestError> = if let Err(html_error) = + self.bot + .edit_message_text(stream.chat_id, stream.message_id, &html) + .parse_mode(ParseMode::Html) + .send() + .await + { + tracing::debug!( + %html_error, + "HTML finalize edit failed, retrying as plain text" + ); + self.bot + .edit_message_text(stream.chat_id, stream.message_id, first_chunk) + .send() + .await + .map(|_| ()) + } else { + Ok(()) + }; + + if let Err(error) = finalized { + tracing::warn!( + %error, + "failed to finalize streaming message, falling back to normal send" + ); + if let Err(delete_error) = self + .bot + .delete_message(stream.chat_id, stream.message_id) + .send() + .await + { + tracing::warn!( + %delete_error, + "failed to delete telegram stream placeholder after finalize fallback" + ); + } + self.active_messages + .write() + .await + .remove(&message.conversation_id); + send_formatted(&self.bot, chat_id, &text, reply_to).await?; + } else { + self.active_messages + .write() + .await + .remove(&message.conversation_id); + for chunk in remaining_chunks { + send_formatted(&self.bot, chat_id, chunk, None).await?; + } + } + } else { + send_formatted(&self.bot, chat_id, &text, None).await?; + } } OutboundResponse::RichMessage { text, poll, .. } => { self.stop_typing(&message.conversation_id).await; @@ -427,10 +509,13 @@ impl Messaging for TelegramAdapter { } OutboundResponse::StreamStart => { self.stop_typing(&message.conversation_id).await; + let reply_to = self.extract_message_id(message).ok(); - let placeholder = self - .bot - .send_message(chat_id, "...") + let mut request = self.bot.send_message(chat_id, "..."); + if let Some(reply_id) = reply_to { + request = request.reply_parameters(ReplyParameters::new(reply_id)); + } + let placeholder = request .send() .await .context("failed to send stream placeholder")?; @@ -441,49 +526,71 @@ impl Messaging for TelegramAdapter { chat_id, message_id: placeholder.id, last_edit: Instant::now(), + text: String::new(), }, ); } OutboundResponse::StreamChunk(text) => { - let mut active = self.active_messages.write().await; - if let Some(stream) = active.get_mut(&message.conversation_id) { - if stream.last_edit.elapsed() < STREAM_EDIT_INTERVAL { - return Ok(()); + let now = Instant::now(); + let update = { + let mut active = self.active_messages.write().await; + if let Some(stream) = active.get_mut(&message.conversation_id) { + stream.text.push_str(&text); + if should_throttle_stream_edit(stream.last_edit, now) { + return Ok(()); + } + stream.last_edit = now; + let display_text = if stream.text.len() > MAX_MESSAGE_LENGTH { + let end = stream.text.floor_char_boundary(MAX_MESSAGE_LENGTH - 3); + format!("{}...", &stream.text[..end]) + } else { + stream.text.clone() + }; + Some((stream.chat_id, stream.message_id, display_text)) + } else { + None } + }; - let display_text = if text.len() > MAX_MESSAGE_LENGTH { - let end = text.floor_char_boundary(MAX_MESSAGE_LENGTH - 3); - format!("{}...", &text[..end]) - } else { - text - }; + let Some((stream_chat_id, stream_message_id, display_text)) = update else { + return Ok(()); + }; - let html = markdown_to_telegram_html(&display_text); - if let Err(html_error) = self + let html = markdown_to_telegram_html(&display_text); + if let Err(html_error) = self + .bot + .edit_message_text(stream_chat_id, stream_message_id, &html) + .parse_mode(ParseMode::Html) + .send() + .await + { + tracing::debug!(%html_error, "HTML edit failed, retrying as plain text"); + if let Err(error) = self .bot - .edit_message_text(stream.chat_id, stream.message_id, &html) - .parse_mode(ParseMode::Html) + .edit_message_text(stream_chat_id, stream_message_id, &display_text) .send() .await { - tracing::debug!(%html_error, "HTML edit failed, retrying as plain text"); + tracing::debug!(%error, "failed to edit streaming message"); + } + } + } + OutboundResponse::StreamEnd => { + with_removed_active_stream( + &self.active_messages, + &message.conversation_id, + |stream| async move { if let Err(error) = self .bot - .edit_message_text(stream.chat_id, stream.message_id, &display_text) + .delete_message(stream.chat_id, stream.message_id) .send() .await { - tracing::debug!(%error, "failed to edit streaming message"); + tracing::warn!(%error, "failed to delete telegram stream placeholder"); } - } - stream.last_edit = Instant::now(); - } - } - OutboundResponse::StreamEnd => { - self.active_messages - .write() - .await - .remove(&message.conversation_id); + }, + ) + .await; } OutboundResponse::Status(status) => { self.send_status(message, status).await?; @@ -986,6 +1093,16 @@ fn split_message(text: &str, max_len: usize) -> Vec { chunks } +fn should_throttle_stream_edit(last_edit: Instant, now: Instant) -> bool { + now.duration_since(last_edit) < STREAM_EDIT_INTERVAL +} + +fn stream_finalization_chunks(chunks: &[String]) -> (&str, &[String]) { + let first_chunk = chunks.first().map(String::as_str).unwrap_or("\u{200B}"); + let remaining_chunks = chunks.get(1..).unwrap_or(&[]); + (first_chunk, remaining_chunks) +} + /// Return true when Telegram rejected rich text entities and a plain-caption retry is safe. fn should_retry_plain_caption(error: &RequestError) -> bool { matches!(error, RequestError::Api(ApiError::CantParseEntities(_))) @@ -1224,6 +1341,9 @@ async fn send_formatted( #[cfg(test)] mod tests { use super::*; + use std::sync::Arc; + use std::time::Duration; + use tokio::sync::oneshot; #[test] fn bold() { @@ -1392,4 +1512,90 @@ mod tests { assert!(should_retry_plain_caption(&parse_error)); assert!(!should_retry_plain_caption(&non_parse_error)); } + + #[test] + fn adapter_streaming_throttle() { + let last_edit = Instant::now(); + let before_interval = last_edit + STREAM_EDIT_INTERVAL - Duration::from_millis(1); + let at_interval = last_edit + STREAM_EDIT_INTERVAL; + + assert!( + should_throttle_stream_edit(last_edit, before_interval), + "stream chunk edits should be skipped before the configured throttle interval" + ); + assert!( + !should_throttle_stream_edit(last_edit, at_interval), + "stream chunk edits should be allowed once the throttle interval has elapsed" + ); + + let chunks = vec![ + "first chunk".to_string(), + "second chunk".to_string(), + "third chunk".to_string(), + ]; + let (finalized_chunk, trailing_chunks) = stream_finalization_chunks(&chunks); + + assert_eq!( + finalized_chunk, "first chunk", + "finalization should update the stream placeholder with the first chunk" + ); + assert_eq!( + trailing_chunks, + &["second chunk".to_string(), "third chunk".to_string()], + "finalization should emit only trailing chunks after placeholder update" + ); + } + + #[tokio::test] + async fn telegram_stream_end_releases_active_messages_lock_before_delete() { + let active_messages = Arc::new(RwLock::new(HashMap::from([( + "conversation-id".to_string(), + ActiveStream { + chat_id: ChatId(1), + message_id: MessageId(1), + last_edit: Instant::now(), + text: "partial".to_string(), + }, + )]))); + let (delete_started_tx, delete_started_rx) = oneshot::channel(); + let (allow_delete_tx, allow_delete_rx) = oneshot::channel::<()>(); + + let cleanup_task = { + let active_messages = active_messages.clone(); + tokio::spawn(async move { + with_removed_active_stream( + active_messages.as_ref(), + "conversation-id", + |stream| async move { + delete_started_tx.send(stream.message_id).ok(); + allow_delete_rx.await.ok(); + }, + ) + .await + }) + }; + + delete_started_rx + .await + .expect("cleanup future should begin after the stream entry is removed"); + + let active_messages_guard = tokio::time::timeout( + Duration::from_millis(100), + active_messages.read(), + ) + .await + .expect("stream cleanup should not hold the active_messages lock while delete is pending"); + assert!( + !active_messages_guard.contains_key("conversation-id"), + "the stream entry should be removed before the delete future completes" + ); + drop(active_messages_guard); + + allow_delete_tx + .send(()) + .expect("test should release the simulated delete"); + cleanup_task + .await + .expect("cleanup task should complete cleanly"); + } } diff --git a/src/telemetry/registry.rs b/src/telemetry/registry.rs index 55efad796..946693ac4 100644 --- a/src/telemetry/registry.rs +++ b/src/telemetry/registry.rs @@ -32,6 +32,26 @@ pub struct Metrics { /// Total memory save (write) operations. pub memory_writes_total: IntCounter, + /// Rig request semantics decisions. + /// Labels: process, provider, field, decision. + pub rig_request_semantics_total: IntCounterVec, + + /// Rig streaming sessions started. + /// Labels: process, provider, mode. + pub rig_stream_sessions_total: IntCounterVec, + + /// Time to first streamed text delta in milliseconds. + /// Labels: process, provider. + pub rig_time_to_first_delta_ms: HistogramVec, + + /// Worker request-level tool concurrency usage. + /// Labels: worker_type, concurrency. + pub rig_tool_concurrency_total: IntCounterVec, + + /// Detected Rig bridge drift surfaces. + /// Labels: surface. + pub rig_drift_detected_total: IntCounterVec, + // -- Histograms -- /// LLM request duration in seconds. pub llm_request_duration_seconds: HistogramVec, @@ -119,6 +139,55 @@ impl Metrics { ) .expect("hardcoded metric descriptor"); + let rig_request_semantics_total = IntCounterVec::new( + Opts::new( + "spacebot_rig_request_semantics_total", + "Rig request semantics decisions by field and outcome", + ), + &["process", "provider", "field", "decision"], + ) + .expect("hardcoded metric descriptor"); + + let rig_stream_sessions_total = IntCounterVec::new( + Opts::new( + "spacebot_rig_stream_sessions_total", + "Rig streaming sessions started by process, provider, and mode", + ), + &["process", "provider", "mode"], + ) + .expect("hardcoded metric descriptor"); + + let rig_time_to_first_delta_ms = HistogramVec::new( + HistogramOpts::new( + "spacebot_rig_time_to_first_delta_ms", + "Time to first streamed text delta in milliseconds", + ) + .buckets(vec![ + 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 1500.0, 2500.0, 3500.0, 5000.0, 10_000.0, + 30_000.0, + ]), + &["process", "provider"], + ) + .expect("hardcoded metric descriptor"); + + let rig_tool_concurrency_total = IntCounterVec::new( + Opts::new( + "spacebot_rig_tool_concurrency_total", + "Worker request-level tool concurrency usage", + ), + &["worker_type", "concurrency"], + ) + .expect("hardcoded metric descriptor"); + + let rig_drift_detected_total = IntCounterVec::new( + Opts::new( + "spacebot_rig_drift_detected_total", + "Detected Rig bridge drift surfaces", + ), + &["surface"], + ) + .expect("hardcoded metric descriptor"); + let llm_request_duration_seconds = HistogramVec::new( HistogramOpts::new( "spacebot_llm_request_duration_seconds", @@ -245,6 +314,21 @@ impl Metrics { registry .register(Box::new(memory_writes_total.clone())) .expect("hardcoded metric"); + registry + .register(Box::new(rig_request_semantics_total.clone())) + .expect("hardcoded metric"); + registry + .register(Box::new(rig_stream_sessions_total.clone())) + .expect("hardcoded metric"); + registry + .register(Box::new(rig_time_to_first_delta_ms.clone())) + .expect("hardcoded metric"); + registry + .register(Box::new(rig_tool_concurrency_total.clone())) + .expect("hardcoded metric"); + registry + .register(Box::new(rig_drift_detected_total.clone())) + .expect("hardcoded metric"); registry .register(Box::new(llm_request_duration_seconds.clone())) .expect("hardcoded metric"); @@ -291,6 +375,11 @@ impl Metrics { tool_calls_total, memory_reads_total, memory_writes_total, + rig_request_semantics_total, + rig_stream_sessions_total, + rig_time_to_first_delta_ms, + rig_tool_concurrency_total, + rig_drift_detected_total, llm_request_duration_seconds, tool_call_duration_seconds, active_workers, diff --git a/src/tools.rs b/src/tools.rs index e88d9e820..54ba858d4 100644 --- a/src/tools.rs +++ b/src/tools.rs @@ -509,6 +509,41 @@ pub fn create_worker_tool_server( server.run() } +/// Compute the concrete worker tool names exposed for a request. +/// +/// This is used to gate request-level tool concurrency to explicit allowlisted, +/// read-only surfaces. +pub fn list_worker_tool_names( + browser_enabled: bool, + brave_search_enabled: bool, + secrets_enabled: bool, + mcp_tools: &[McpToolAdapter], +) -> Vec { + let mut names = vec![ + ::NAME.to_string(), + ::NAME.to_string(), + ::NAME.to_string(), + ::NAME.to_string(), + ::NAME.to_string(), + ::NAME.to_string(), + ]; + + if secrets_enabled { + names.push(::NAME.to_string()); + } + if browser_enabled { + names.push(::NAME.to_string()); + } + if brave_search_enabled { + names.push(::NAME.to_string()); + } + for mcp_tool in mcp_tools { + names.push(mcp_tool.name()); + } + + names +} + /// Create a ToolServer for the cortex process. /// /// The cortex only needs memory_save for consolidation. Additional tools can be