diff --git a/CHANGELOG.md b/CHANGELOG.md
index 33b72912..9aada60e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,54 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
+### Added
+
+- Live SSE relay for `tools/call`: when an upstream MCP server answers a tool
+ call with a Server-Sent Events stream (Streamable HTTP transport, MCP spec
+ 2025-11-25), Atryum relays intermediate messages (progress, logging, other
+ notifications) to the connected agent as they arrive, instead of buffering
+ the whole response and returning only the terminal result. Activates only
+ when the agent's request sends `Accept: text/event-stream` and the
+ upstream responds with one; non-streaming tool calls are unaffected.
+ If an upstream closes a resumable SSE response before the terminal
+ JSON-RPC message, Atryum reconnects with `Last-Event-ID` and continues
+ from the last event acknowledged to the upstream.
+ Streamed events and terminal delivery are audited into `invocation_events`
+ (`invocation.stream_event`, `invocation.stream_completed`, and
+ `invocation.stream_delivery`). New
+ `[defaults]` config knobs: `stream_relay_enabled` (kill-switch, default
+ on), `stream_header_timeout_seconds`, `stream_idle_timeout_seconds`,
+ `stream_max_duration_seconds`, `stream_max_message_bytes`,
+ `stream_audit_max_events`,
+ `stream_audit_max_event_bytes`. See `docs/architecture.md` for the full
+ design.
+- The relay also listens on the Streamable HTTP standalone SSE stream (a GET
+ to the upstream endpoint, independent of any specific `tools/call`), since
+ some upstream SDKs (e.g. the reference MCP Python SDK's
+ `Context.report_progress`) send progress notifications there rather than on
+ the `tools/call` response itself. Atryum rewrites each call's
+ `_meta.progressToken` to a value unique to that call before forwarding it
+ upstream, so two unrelated concurrent callers who happen to choose the same
+ token can never have their progress notifications cross-delivered.
+ Standalone connections now follow session renewal and retry transient
+ failures with bounded backoff. Stream audit writes use a fixed shared worker
+ pool; standalone-buffer drops and downstream terminal-delivery outcomes are
+ recorded explicitly.
+- Relay behavior guarantees: server-initiated upstream requests (sampling,
+ elicitation, roots) are audited but never forwarded to the agent, whichever
+ connection carried them; SSE keepalive/comment lines count as upstream
+ activity for the idle timeout, so a busy-but-quiet tool heartbeating through
+ a long run is not cut off; and stream failure audit reasons distinguish a
+ proven downstream abort (`stream_aborted_downstream`, only when a write to
+ the agent failed) from a bare request-context cancellation
+ (`stream_canceled`) and upstream timeouts (`stream_timeout`).
+
+### Fixed
+
+- Buffered (non-streaming) tool calls whose downstream client disconnects
+ mid-call no longer leave the invocation row stuck `executing`: finalization
+ writes now run on their own bounded context, matching the streaming path.
+
## [0.2.0] - 2026-07-14
### Added
diff --git a/Justfile b/Justfile
index 7da94283..4bf1126e 100644
--- a/Justfile
+++ b/Justfile
@@ -289,6 +289,24 @@ judge-eval-check:
go test -tags judgeeval ./internal/invocation \
-run 'TestJudge(GarbageOutput|MarkdownFenced|Request|UnrecognizedVerdict)|TestConstantVerdictBaselines' -v
+# Real end-to-end test of the tools/call SSE relay against the official MCP
+# reference "everything" server (@modelcontextprotocol/server-everything),
+# spawned live via npx over its Streamable HTTP transport. Requires
+# Node/npm; network access on first run to fetch the package. Skips itself
+# if npx isn't on PATH. See internal/api/mcp_everything_test.go.
+mcp-everything-test:
+ go test -tags mcpeverything ./internal/api -run TestMCPToolsCallAgainstRealEverythingServer -v
+
+# Real end-to-end test of the tools/call SSE relay's standalone-stream path,
+# against a real MCP Python SDK (FastMCP) server spawned live via uv.
+# FastMCP's Context.report_progress sends progress on the standalone SSE
+# stream, never the tools/call response itself — the complement to
+# mcp-everything-test above. Requires uv (https://docs.astral.sh/uv/);
+# network access on first run to resolve the mcp package. Skips itself if uv
+# isn't on PATH. See internal/api/mcp_standalone_stream_test.go.
+mcp-standalone-stream-test:
+ go test -tags mcpstandalone ./internal/api -run TestMCPToolsCallAgainstRealStandaloneStreamServer -v
+
# List registered harnesses, auth protocols, and MCP targets
integration-list:
integrations/scripts/agent_harness_integration_tests.sh list
diff --git a/README.md b/README.md
index a65b528e..2a707beb 100644
--- a/README.md
+++ b/README.md
@@ -157,7 +157,7 @@ The embedded invocation view subscribes to the admin SSE stream, updates list/de
## Storage
-SQLite by default, PostgreSQL optional via `server.database_url`. Both are first-class — migrations live in `internal/store/migrations/` and apply at startup. Core tables:
+SQLite by default, PostgreSQL optional via `server.database_url`. Both are first-class. Built-in migration definitions live in `pkg/migrations/`; `internal/store` applies and tracks them at startup. Embedding programs can add separately tracked, namespaced migrations through `pkg/atryum.WithMigrations`. Core tables:
- `mcp_servers` — upstream connection settings and generic auth/connection status, including `connection_status`, `auth_status`, `reauth_needed`, `auth_type`, `last_checked_at`, `last_check_ok`, `last_error_summary`, and `action_required`.
- `oauth_credentials` and related OAuth client registration tables — tokens and client registrations held by Atryum on behalf of agents.
@@ -218,6 +218,21 @@ After first-run bootstrap, edit MCP servers through the UI/API; TOML `[[upstream
When `backend.base_url` is empty, the ValidMind backend connection check is skipped for local standalone runs. When it is set, startup fails if credentials are missing or `GET /api/atryum/unstable/connection` is rejected. Environment variables override TOML: `VM_BASE_URL`, `VM_MACHINE_KEY`, `VM_MACHINE_SECRET`, and `VM_CONNECTION_TIMEOUT_SECONDS`.
+## Embedding
+
+The stock executable in `cmd/atryum` is a thin wrapper around the public
+`pkg/atryum` package. Another Go program can call `atryum.Main` to run the same CLI and
+server with these options:
+
+- `WithRoutes` adds HTTP routes. Extension routes are outside Atryum's built-in
+ authentication middleware and must authenticate themselves.
+- `WithMigrations` adds namespaced migrations after the built-in sequence.
+- `WithDatabase` runs a database hook after all migrations and before the server starts.
+- `WithThirdPartyNotices` supplies the notices printed by the `licenses` command.
+
+Extensions run inside the Atryum process and share its HTTP server and database. See
+[`pkg/atryum`](pkg/atryum) and [`pkg/migrations`](pkg/migrations) for the public types.
+
## Running
Single-binary Go service.
diff --git a/atryum.example.toml b/atryum.example.toml
index d1e87dbd..7b708095 100644
--- a/atryum.example.toml
+++ b/atryum.example.toml
@@ -49,6 +49,32 @@ connection_timeout_seconds = 5
[defaults]
request_timeout_seconds = 30
+# Live relay for tools/call. When the agent accepts SSE, Atryum can forward
+# intermediate HTTP SSE or stdio JSON-RPC messages as they arrive, followed
+# by one terminal SSE response. A plain upstream JSON response remains plain
+# downstream. stream_relay_enabled disables the relay without a rollback.
+stream_relay_enabled = true
+# Bounds HTTP session initialization and tools/call response headers, or the
+# stdio initialize handshake. 0 falls back to request_timeout_seconds above.
+stream_header_timeout_seconds = 0
+# Bounds the gap in upstream activity while reading the response. Resets on
+# every message, and on SSE keepalive/comment lines, so a quiet-but-alive
+# upstream is not cut off. Does not bound the call's total duration.
+stream_idle_timeout_seconds = 60
+# Bounds the complete upstream response-reading phase. 0 = unlimited.
+stream_max_duration_seconds = 600
+# Rejects one upstream SSE event, stdio JSON-RPC line, or plain JSON response
+# body larger than this many bytes. This bounds memory used by one message.
+stream_max_message_bytes = 4194304
+# Caps how many invocation.stream_event audit rows are persisted per call.
+# Beyond the cap, events are still relayed live to the agent but only
+# counted, not stored individually. 0 disables this count cap. The bounded
+# shared audit queue reports storage backpressure drops in stream_completed.
+stream_audit_max_events = 100
+# Truncates each persisted stream_event row's data field beyond this many
+# bytes. 0 = no truncation.
+stream_audit_max_event_bytes = 4096
+
# Optional Claude Managed Agents events bridge. Declare one [[managed_agents]]
# table per Anthropic account/workspace you want to watch. When an entry's
# api_key is set, Atryum connects outbound to Anthropic's Managed Agents
diff --git a/docs/architecture.md b/docs/architecture.md
index f3ced24d..ddc086ac 100644
--- a/docs/architecture.md
+++ b/docs/architecture.md
@@ -8,24 +8,37 @@ the [README](../README.md).
Atryum is a Go service that mediates tool calls. It exposes runtime endpoints to agent
harnesses, an MCP-compatible proxy, and an admin API/UI. All ingress paths share the
-same invocation service, rule evaluation, and audit store.
+same invocation service, rule evaluation, and audit store. The stock binary and programs
+that embed Atryum both enter through the public `pkg/atryum` bootstrap package.
+
+Two terms recur throughout this document. **Downstream** is the agent or harness side —
+the caller waiting on Atryum's response. **Upstream** is the tool-server side — where
+Atryum forwards the call to do the actual work. Atryum sits between them: a request flows
+agent → Atryum → tool server, and the result (or, for a streamed call, progress updates
+along the way) flows back the other direction.
```mermaid
flowchart LR
+ Stock[cmd/atryum stock binary]
+ Embedder[Embedding Go program]
Hook[Agent harness or hook]
- MCP[MCP client]
+ MCP[Downstream MCP client]
Claude[Claude Managed Agents API]
Admin[Admin UI or API client]
subgraph Atryum[Atryum process]
+ Bootstrap[pkg/atryum bootstrap and extensions]
API[HTTP and MCP handlers]
Watcher[Managed Agents watcher]
Service[Invocation service]
Rules[Rule evaluation]
Store[(SQLite or PostgreSQL)]
- Upstream[MCP client]
+ MCPClient[Upstream MCP client]
end
+ Stock --> Bootstrap
+ Embedder -->|extension options| Bootstrap
+ Bootstrap -->|constructs| API
Hook -->|submit, poll, report outcome| API
MCP -->|MCP JSON-RPC| API
Admin -->|review and configuration| API
@@ -35,12 +48,13 @@ flowchart LR
Service --> Rules
Rules --> Store
Service --> Store
- Service --> Upstream
- Upstream -->|HTTP or stdio| Tools[Upstream MCP servers]
+ Service --> MCPClient
+ MCPClient -->|HTTP or stdio| Tools[Upstream MCP servers]
```
-The diagram corresponds to `internal/api`, `internal/managedagents`,
-`internal/invocation`, `internal/store`, and `internal/mcp`.
+The diagram corresponds to `cmd/atryum`, `pkg/atryum`, `pkg/migrations`,
+`internal/api`, `internal/managedagents`, `internal/invocation`, `internal/store`, and
+`internal/mcp`.
## Runtime ingress
@@ -61,9 +75,12 @@ or audit records are evaluated.
| Package | Responsibility | Must not own |
|---|---|---|
+| `cmd/atryum` | Thin stock executable that supplies bundled notices and calls `pkg/atryum` | CLI, server, or business logic |
+| `pkg/atryum` | Public CLI/server bootstrap and extension options for routes, migrations, database hooks, and notices | Invocation decisions or transport internals |
+| `pkg/migrations` | Portable definitions for Atryum's built-in schema migrations | Repository queries or migration execution state |
| `internal/api` | HTTP/MCP transport, authentication middleware, request/response mapping, embedded UI | Rule or invocation state transitions |
| `internal/invocation` | Invocation lifecycle, rule matching, AI evaluation dispatch, approval coordination | SQL or upstream transport details |
-| `internal/store` | SQLite/PostgreSQL repositories, schema migrations, durable query semantics | Policy decisions |
+| `internal/store` | SQLite/PostgreSQL repositories, migration execution and tracking, durable query semantics | Policy decisions |
| `internal/mcp` | Server resolution, MCP forwarding, upstream authentication and OAuth | Approval policy |
| `internal/auth` | Inbound OIDC/JWT validation and authenticated identity context | Upstream MCP credentials |
| `internal/managedagents` | Anthropic session discovery, event replay, confirmation delivery | Independent rule evaluation |
@@ -71,6 +88,23 @@ or audit records are evaluated.
The React application in `ui/` is compiled into `internal/api/web/` for the production
binary. During development it can run separately, but it still uses the same admin API.
+### Embedding and extension boundary
+
+`cmd/atryum` is only the stock executable. The reusable application lives in
+`pkg/atryum`, so another Go program can start the same server with additional options:
+
+- `WithRoutes` mounts extra HTTP routes after built-in routes. These routes are outside
+ Atryum's authentication middleware and must provide their own authentication.
+- `WithMigrations` registers extension-owned, namespaced schema migrations. They run
+ after all built-in migrations and are tracked separately by namespace and version.
+- `WithDatabase` runs a hook after built-in and extension migrations finish but before
+ the server starts accepting requests.
+- `WithThirdPartyNotices` replaces the notices text printed by the `licenses` command.
+
+Extensions execute inside the Atryum process and share its database and HTTP server.
+They are trusted application code, not isolated plugins. Route patterns must not collide
+with built-in routes, and a migration namespace must remain stable across releases.
+
## Decision pipeline
Rules are loaded in ascending `rule_order`. Matching uses server/source, tool, and the
@@ -111,7 +145,33 @@ rule matches; it is not evaluated after a matching rule defers.
An `ai_evaluation` rule selects one configured evaluator. Standalone deployments use
a local LLM configuration stored in `llm_configs`. Evaluation errors escalate to human
review; missing charter context denies; and an unknown verdict is treated as
-`next_rule`, eventually reaching human review if no later rule decides.
+`next_rule`, eventually reaching human review if no later rule decides. Drawn out,
+that rule's own outcomes are:
+
+```mermaid
+flowchart TD
+ Eval[ai_evaluation rule runs its configured evaluator]
+ Verdict{Evaluator outcome}
+ Approved[approved]
+ Denied[denied]
+ NextRule[next_rule]
+ EvalError[Evaluation call itself failed]
+ NoCharter[No charter context available for this agent or tool]
+ Unparseable[Verdict does not parse as approve, deny, or defer]
+
+ Eval --> Verdict
+ Verdict -->|approve| Approved
+ Verdict -->|deny| Denied
+ Verdict -->|defer| NextRule
+ Verdict -->|error| EvalError --> Human2[Escalate straight to human approval]
+ Verdict -->|missing charter| NoCharter --> Denied
+ Verdict -->|unknown verdict| Unparseable --> NextRule
+```
+
+Only the `defer` and `unknown verdict` branches feed back into "another matching
+rule?" in the pipeline diagram above; `error` skips that loop and goes straight to
+human review because a broken evaluator is not expected to fix itself on the next
+matching rule.
## Atryum-executed calls
@@ -136,6 +196,31 @@ process. Concretely:
remain `pending_approval`. Nothing resumes pending invocations at startup, so a
later approval updates that row but does not execute the tool.
+The first bullet end to end, across two replicas:
+
+```mermaid
+sequenceDiagram
+ autonumber
+ participant Client
+ participant ProcA as Atryum process A - handling this request
+ participant Store as PostgreSQL
+ participant Admin
+ participant ProcB as Atryum process B - handling the admin decision
+
+ Client->>ProcA: Invoke tool
+ ProcA->>Store: Persist pending_approval
+ ProcA->>ProcA: Block the request goroutine on an in-memory channel
+ Admin->>ProcB: Approve invocation
+ ProcB->>ProcB: Look up the in-memory channel for this invocation - not found here
+ Note over ProcB: The waiting channel exists only in process A's memory
+ ProcB->>Store: Persist approved directly, since no local waiter exists
+ Note over ProcA: Process A's request is still blocked; nothing woke it
+ Client-->>ProcA: Client disconnects, or the request context is cancelled
+ ProcA->>ProcA: Context done fires while still waiting
+ ProcA->>Store: Persist failed, error text "cancelled"
+ Note over Store: Row now reads failed, overwriting the approved status
process B just wrote. The tool was never executed.
+```
+
Multi-replica or crash-resumable execution requires durable coordination — for
example a work queue, or an outbox table that a worker drains — which is not
implemented.
@@ -186,6 +271,662 @@ The admin invocation stream is a polling SSE view: the handler queries the durab
invocation list every two seconds and emits when its signature changes. Database writes
do not directly publish to the stream.
+### Live SSE relay for tools/call
+
+#### Purpose and scope
+
+Atryum normally returns one response after an upstream tool finishes. The live relay
+also lets Atryum send progress updates while the tool is still running.
+
+This behavior applies only to `tools/call` requests sent through the MCP proxy at
+`/mcp/{server}`. It does not apply to the direct REST endpoint
+`POST /api/v1/invocations`.
+
+Streaming is opt-in and keeps the old behavior as its fallback:
+
+1. The downstream MCP client includes `Accept: text/event-stream` to say it can read
+ Server-Sent Events.
+2. Atryum calls the upstream MCP server.
+3. If the upstream returns a normal JSON response, Atryum returns one normal JSON
+ response.
+4. If the upstream starts a stream, Atryum relays each progress update and then the
+ final result.
+
+#### Terms used in this section
+
+| Term | Meaning here |
+|---|---|
+| **MCP** | Model Context Protocol, the protocol used to call tools. MCP messages use JSON-RPC. |
+| **Downstream MCP client** | The agent or agent harness calling Atryum. “Downstream” means the side receiving Atryum's response. |
+| **Atryum** | The relay between the downstream client and the upstream server. |
+| **Upstream MCP server** | The tool server that Atryum calls. “Upstream” means the side doing the tool work. |
+| **JSON-RPC request** | A message asking for work. It has an `id`, and the response must carry the same `id`. |
+| **JSON-RPC notification** | A one-way message such as a progress update. It has a `method` but no `id`, so no reply is expected. |
+| **Terminal response** | The final JSON-RPC success or error for the tool call. It is the one result a non-streaming call would return. |
+| **Terminal delivery** | The final write from Atryum to the downstream client. It can fail even after the upstream result was saved successfully. |
+| **SSE** | Server-Sent Events, a one-way HTTP response format that lets a server send multiple events over one open response. |
+| **Progress token** | A string or number used to match progress to a call. The client requests it with `_meta.progressToken`; notifications return it as `params.progressToken`. |
+| **MCP session** | A group of requests recognized by the upstream server through one session ID. Atryum shares that upstream session across active calls. |
+| **Call-response stream** | The response body of the upstream `tools/call` POST. It belongs to one call. |
+| **Standalone stream** | A separate SSE GET connection shared by active calls in one upstream MCP session. Some MCP servers send progress here instead of on the call-response stream. |
+| **Stdio upstream** | An MCP server run as a local process. Atryum exchanges messages through the process's standard input and output instead of HTTP. |
+
+An SSE event contains one or more `data:` lines and ends with a blank line. For example:
+
+```
+data: {"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1,"total":3}}
+
+data: {"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}
+```
+
+The first event is a progress notification. The second event is the terminal response.
+
+#### The two upstream paths
+
+Every HTTP tool call has its own POST response. Some upstream servers put both progress
+and the terminal response on that response. Other servers put progress on a separate,
+shared GET stream while still returning the terminal response on the POST response.
+Atryum listens to both paths.
+
+```mermaid
+flowchart TB
+ Downstream[Downstream MCP client]
+ Relay[Atryum]
+ Upstream[Upstream MCP server]
+ CallPath[Path A: one POST per call
Progress and terminal response]
+ StandalonePath[Path B: one shared SSE GET
Progress only]
+
+ Downstream <-->|tools/call request
JSON or SSE response| Relay
+ Relay --- CallPath
+ CallPath --- Upstream
+ Relay --- StandalonePath
+ StandalonePath --- Upstream
+```
+
+| | Path A: call-response stream | Path B: standalone stream |
+|---|---|---|
+| HTTP connection | The response to one `tools/call` POST | One SSE GET shared by active calls in an upstream session |
+| Carries progress | Yes | Yes |
+| Carries the terminal response | Yes | No; the terminal response still arrives on the POST response |
+| How an update is matched to a call | The response already belongs to that call | Atryum matches `params.progressToken` |
+
+For a stdio upstream there is no HTTP or SSE on the upstream side. The process sends one
+JSON-RPC message per line instead. Atryum applies the same message classification,
+timeout, audit, and downstream-relay rules to those messages.
+
+#### Normal call flow
+
+```mermaid
+sequenceDiagram
+ autonumber
+ participant Client as Downstream MCP client
+ participant Atryum
+ participant Upstream as Upstream MCP server
+
+ Client->>Atryum: tools/call, accepts SSE
+ Atryum->>Atryum: Evaluate rules and approval policy
+ Atryum->>Upstream: Execute the tool
+ alt Upstream returns one normal response
+ Upstream-->>Atryum: Terminal JSON-RPC response
+ Atryum-->>Client: One normal JSON response
+ else Upstream streams
+ loop For each progress notification, if any
+ Upstream-->>Atryum: Progress notification on Path A or B
+ Atryum->>Atryum: Queue audit write
+ Atryum-->>Client: Progress notification as SSE
+ end
+ Upstream-->>Atryum: Terminal response on Path A
+ Atryum->>Atryum: Save final invocation state
+ Atryum-->>Client: Terminal response as SSE, then close
+ Atryum->>Atryum: Record whether terminal delivery succeeded
+ end
+```
+
+Approval happens before Atryum contacts the upstream server. While waiting for a human
+decision, the downstream request remains open but Atryum has not started an SSE
+response. A denial stays a normal JSON response. After approval, execution follows the
+flow above.
+
+#### Error and timeout paths
+
+The flow above shows the two happy branches: a plain response, and a stream that runs
+to completion. The same call has four more ways to end, driven by the timers described
+in [Resource limits](#resource-limits) below and by the downstream connection itself:
+
+```mermaid
+sequenceDiagram
+ autonumber
+ participant Client as Downstream MCP client
+ participant Atryum
+ participant Upstream as Upstream MCP server
+
+ Client->>Atryum: tools/call, accepts SSE
+ Atryum->>Atryum: Evaluate rules and approval policy
+ Atryum->>Upstream: Execute the tool
+ Atryum->>Atryum: Arm the setup timer (stream_header_timeout_seconds)
+
+ alt Setup timer fires before headers or a stream start arrive
+ Atryum->>Atryum: Trip: stream setup timeout exceeded
+ Atryum-->>Client: Terminal SSE error, reason stream_timeout
+ else Upstream opens a stream in time
+ Atryum->>Atryum: Disarm the setup timer; arm the idle and max-duration timers
+ loop while the upstream keeps sending progress
+ Upstream-->>Atryum: Progress notification
+ Atryum->>Atryum: Reset the idle timer
+ Atryum-->>Client: Progress notification as SSE
+ end
+ alt Idle timer fires and the gap genuinely exceeds stream_idle_timeout_seconds
+ Atryum->>Atryum: Trip: idle timeout waiting for the next stream event
+ Atryum-->>Client: Terminal SSE error, reason stream_timeout
+ else Max-duration timer fires regardless of activity
+ Atryum->>Atryum: Trip: max stream duration exceeded
+ Atryum-->>Client: Terminal SSE error, reason stream_timeout
+ else Client disconnects or stops reading mid-stream
+ Atryum->>Atryum: The next write to the client misses its per-write deadline
+ Note over Atryum: Recorded as stream_aborted_downstream;
no further write is attempted
+ else Upstream sends the terminal response normally
+ Upstream-->>Atryum: Terminal response
+ Atryum->>Atryum: Save final invocation state
+ Atryum-->>Client: Terminal response as SSE, then close
+ end
+ end
+```
+
+The idle and max-duration branches both end in the same reported reason,
+`stream_timeout`, because both are the guard's own bound firing rather than a
+transport failure; see the classification flowchart below for the full priority order.
+
+#### Matching shared progress to the correct call
+
+The standalone stream is shared, so receiving an event does not by itself identify the
+call that owns it. Atryum uses this process:
+
+1. The downstream client supplies `_meta.progressToken`.
+2. Before sending the tool call upstream, Atryum replaces that token with a value unique
+ to this call.
+3. On the standalone path, Atryum uses its unique value to select the correct call. On
+ the call-response path, the HTTP response already identifies the call.
+4. Before relaying progress from either path, Atryum restores the client's original token.
+
+Rewriting is necessary because two unrelated clients can choose the same token. Without
+it, one client could receive another client's progress.
+
+A standalone notification without a progress token cannot always be matched safely. If
+exactly one call is currently using the shared stream, Atryum sends the notification to
+that call. If several calls are active, Atryum drops it rather than guess and risk
+cross-delivery.
+
+#### Code responsibilities
+
+| Layer | Responsibility for this feature |
+|---|---|
+| `pkg/atryum` (startup only) | Load stream configuration, inject timeout and audit limits into the invocation service, and set the relay kill switch on the HTTP handler. |
+| `internal/api` | Detect downstream SSE support, write and flush SSE frames, send heartbeats, enforce downstream write deadlines, and finish the stream. |
+| `internal/invocation` | Preserve rule and approval behavior, update invocation state, and audit stream events through a bounded shared dispatcher. |
+| `internal/mcp` | Call the upstream, parse and classify JSON-RPC messages, merge the two upstream paths, correlate progress tokens, reconnect resumable streams, and enforce upstream timeouts. |
+
+After startup wiring, each call crosses the runtime packages in this order:
+`internal/api` → `internal/invocation` → `internal/mcp`.
+
+#### Control flow in pseudocode
+
+The table above says which package owns what; this section walks the same call in
+the order it actually runs, in pseudocode, with the reason each step exists. Nothing
+here is transport-specific until the HTTP/stdio split.
+
+**1. Deciding whether to stream at all (`internal/api`)**
+
+```
+handle tools/call request:
+ if client sent "Accept: text/event-stream" and the relay is not killed-switched:
+ sink = new SSE sink wrapping this response
+ # Constructing the sink does not commit the response to SSE. Only the
+ # upstream actually starting a stream (below) does that. This is what
+ # keeps a plain-JSON upstream response possible right up until the
+ # last moment.
+ else:
+ sink = nil
+
+ result = invocationService.InvokeStreaming(request, sink)
+
+ if sink was actually started:
+ # Headers are already committed as SSE. From here the only spec-legal
+ # way to end the exchange is one more SSE frame carrying the result
+ # (or a JSON-RPC error) — never a second status line, never a bare
+ # close.
+ write result as a final SSE event, stop the heartbeat, close
+ else:
+ # Either there was no sink, or one existed but the upstream never
+ # opened a stream (a plain JSON tool). Nothing has reached the wire
+ # yet, so the ordinary response is still available.
+ write result as one normal JSON response
+```
+
+**2. Rule and approval evaluation stay sink-agnostic (`internal/invocation`)**
+
+```
+InvokeStreaming(request, sink):
+ resolve the upstream server
+ create the durable invocation row (status = received)
+ evaluate rules, or the default policy if none match # identical to the buffered call
+ record the decision as an audit event
+
+ switch decision:
+ auto_denied -> persist denied, return # sink never touched
+ auto_approved -> finishExecution(invocation, sink)
+ human_approval ->
+ persist pending_approval
+ block on an in-memory channel until an admin decides
+ # sink stays untouched for as long as this blocks: nothing streams
+ # while a call is waiting on a human, matching the plain
+ # Atryum-executed-call flow above.
+ if approved: finishExecution(invocation, sink)
+ else: persist denied, return
+
+finishExecution(invocation, sink):
+ if sink == nil:
+ finishExecutionBuffered(invocation) # today's non-streaming path, unchanged
+ else:
+ finishExecutionStreaming(invocation, sink)
+```
+
+Threading `sink` this deep exists so the decision pipeline never has to know streaming
+exists at all — it is the same code, gated only by whether `sink` happens to be `nil`.
+
+**3. Wrapping the sink for audit before it reaches the transport (`internal/invocation`)**
+
+```
+finishExecutionStreaming(invocation, sink):
+ audited = decorate(sink) so every relayed event is durably recorded as it passes through
+ # A decorator, not a fork: it never skips forwarding to the real
+ # sink to make its own recording succeed, and it records the
+ # event even if forwarding then fails downstream — the audit
+ # trail must reflect what the upstream actually sent, not just
+ # what the agent actually received.
+ result, err = mcpClient.InvokeStream(upstream, tool, input, audited, streamOptions)
+
+ if err:
+ reason = classify(err):
+ the audited sink saw its own forwarding call fail -> "the agent connection died"
+ our own guard's context was the one that cancelled -> "we gave up (timeout)"
+ anything else -> "the transport itself failed"
+ persist invocation as failed, with that reason on the audit row
+ else:
+ persist invocation as succeeded or failed, from the upstream's own result
+
+ audited.finish(outcome) # one summary audit row: events seen vs. persisted vs. dropped
+ return response
+```
+
+`classify(err)` is checked in a fixed priority order, because more than one condition
+can be true at once (for example, the context can be canceled *and* the guard can have
+tripped) and only the first match is recorded:
+
+```mermaid
+flowchart TD
+ Err[InvokeStream returned an error]
+ Q1{Did the audited sink's own
write to the agent fail first?}
+ Q2{Is the error, or the call's context,
context.Canceled?}
+ Q3{Is the error mcp.ErrStreamTimeout?}
+ Q4{Is the error mcp.ErrStreamSessionRetryRefused?}
+ Q5{Is the error mcp.ErrStreamMessageTooLarge?}
+ R1["stream_aborted_downstream:
the agent connection died"]
+ R2["stream_canceled:
no proof of who went away -
quiet disconnect and server shutdown look the same"]
+ R3["stream_timeout:
the guard's own setup, idle, or max-duration bound fired"]
+ R4["stream_session_retry_refused:
a relayed event already reached the agent,
so retrying would duplicate it"]
+ R5["stream_message_too_large:
an event exceeded stream_max_message_bytes"]
+ R6["transport_error:
the upstream connection failed on its own"]
+
+ Err --> Q1
+ Q1 -->|yes| R1
+ Q1 -->|no| Q2
+ Q2 -->|yes| R2
+ Q2 -->|no| Q3
+ Q3 -->|yes| R3
+ Q3 -->|no| Q4
+ Q4 -->|yes| R4
+ Q4 -->|no| Q5
+ Q5 -->|yes| R5
+ Q5 -->|no| R6
+```
+
+Classifying the error exists so an operator reading the audit trail can tell "the agent
+hung up" apart from "our own bound fired" apart from "the upstream broke" — otherwise
+every one of those looks like the same generic transport error.
+
+**4. Choosing a transport (`internal/mcp`)**
+
+```
+InvokeStream(upstream, tool, input, sink, opts):
+ if sink == nil: return Invoke(...) # buffered call, unchanged
+ if upstream is stdio: return invokeStdioStream(...)
+ else: return invokeHTTPStream(...)
+```
+
+**5. HTTP: the two-path merge — the core of the feature**
+
+```
+invokeHTTPStream(upstream, tool, input, sink, opts):
+ ensure the upstream HTTP session is initialized
+
+ if the caller supplied a progressToken:
+ wireToken = mint a value unique to this call; remember the caller's original token
+ # Atryum multiplexes every concurrent caller of one upstream onto a
+ # single shared session. If two unrelated callers happened to choose
+ # the same progressToken, forwarding it unchanged would let one
+ # caller's progress leak into another's stream. A per-call wire token
+ # makes routing on the standalone path (below) unambiguous, and the
+ # original token is restored before anything is relayed downstream —
+ # the caller never sees Atryum's internal value.
+
+ standalone = acquire the shared standalone-stream connection for this upstream session
+ # Lazily created: the first concurrent caller opens it;
+ # later ones just add themselves as listeners on the one
+ # connection already running.
+ register this call as a listener for wireToken
+ wait for the standalone connection's first connect attempt to resolve
+ # Acquiring only starts a background connection attempt —
+ # it does not wait for it. If the request below reached
+ # the upstream and produced its first progress
+ # notification before this GET had actually subscribed,
+ # that notification would have nowhere to land and would
+ # be silently lost. This wait closes that race. It is
+ # bounded by the header timeout, so a hung upstream
+ # degrades to "no standalone progress for this call"
+ # rather than blocking it forever.
+
+ send the tools/call request upstream, carrying wireToken in place of the caller's token
+
+ if the response is plain JSON:
+ return it directly # nothing below this line runs
+ else: # the response is SSE
+ start a background reader for this call's own POST response (Path A, below)
+ loop: # the merge loop
+ select whichever arrives first:
+ a message from the standalone stream, already routed to this call (Path B)
+ a message from this call's own POST-response reader (Path A)
+
+ on a progress/notification, from either path:
+ restore the caller's original progressToken
+ hand it to the audited sink (recorded, then relayed to the agent)
+
+ on the terminal response (only ever arrives via Path A):
+ briefly keep draining any standalone progress already in flight
+ # Progress and the terminal response travel on
+ # independent connections, so one can race past the
+ # other. A short settle window lets a near-simultaneous
+ # progress notification still arrive before the call is
+ # considered finished, instead of losing it by a few
+ # milliseconds.
+ unregister this call's standalone listener
+ return the terminal result
+```
+
+Path A — this call's own reader, the only one of the two paths that can resume:
+
+```
+Path A reader (its own goroutine, one per call):
+ loop:
+ read the next SSE event from the POST response
+ if the connection drops:
+ if no event carrying an id was ever seen: give up # nothing to resume from
+ else:
+ wait a bounded, jittered backoff
+ reconnect with "Last-Event-ID: "
+ # The Streamable HTTP transport lets a server end this
+ # response early and expects the client to resume from where
+ # it left off, rather than losing everything already sent or
+ # replaying the call from scratch.
+ if the server inclusively replays that same id, skip the one duplicate
+ keep reading from the resumed connection
+ else:
+ forward the event into the merge loop above
+```
+
+As a sequence, the resumable case looks like this:
+
+```mermaid
+sequenceDiagram
+ autonumber
+ participant Reader as Path A reader - one goroutine per call
+ participant Upstream
+ participant Merge as merge loop in invokeHTTPStream
+
+ Reader->>Upstream: Open this call's own POST response
+ loop while connected
+ Upstream-->>Reader: SSE event, some carrying an id
+ Reader->>Merge: Forward the event
+ end
+ Upstream-->>Reader: Connection drops
+ alt No event with an id was ever seen
+ Reader->>Merge: Report transport failure - nothing to resume from
+ else At least one id was seen
+ Reader->>Reader: Wait a bounded, jittered backoff
+ Reader->>Upstream: Reconnect with header Last-Event-ID: last id seen
+ alt Upstream inclusively replays that same id
+ Upstream-->>Reader: Event id N, a duplicate
+ Reader->>Reader: Skip the duplicate
+ end
+ Upstream-->>Reader: Subsequent events resume normally
+ Reader->>Merge: Forward events as before
+ end
+```
+
+Path B — the standalone stream, shared by every concurrent call on the session:
+
+```
+Standalone reader (one shared goroutine per upstream session):
+ loop:
+ open a GET SSE connection to the upstream (no Last-Event-ID — this
+ isn't resuming a specific call, just listening)
+ signal "first connect attempt resolved" # what invokeHTTPStream waits on above
+ loop:
+ read the next SSE event
+ if it carries a progressToken: hand it to whichever registered call owns that token
+ else if exactly one call is currently listening: hand it to that call
+ # With only one candidate, attributing an untokened message is
+ # safe. With several in flight, it would be a guess — and
+ # guessing wrong leaks one caller's message to another, which
+ # is worse than dropping it.
+ else: drop it # ambiguous; safety over completeness
+ on disconnect: reconnect with backoff # no replay cursor on this path —
+ # messages sent while disconnected can be missed
+```
+
+**6. stdio: a simpler single-reader path, no shared connections**
+
+```
+invokeStdioStream(upstream, tool, input, sink, opts):
+ start the upstream as a child process, wired to its stdin/stdout
+ run the initialize handshake, bounded by the header timeout
+ send the tools/call request as one JSON-RPC line
+
+ loop reading newline-delimited JSON-RPC messages:
+ on a notification or a server-to-client request: hand it to the sink
+ on the terminal response for this call: return it
+
+ # Always kill the whole process group on the way out, even after a
+ # timeout — killing only the direct child can leave a spawned
+ # grandchild process running forever.
+```
+
+**7. Enforcing timeouts across both transports (`internal/mcp`)**
+
+```
+callTimeoutGuard:
+ one cancellable context shared by the whole call
+
+ arm a setup timer covering "waiting on response headers" / "the initialize handshake"
+ once headers or the handshake are in: disarm the setup timer, arm two more:
+ idle timer — reset on every relayed message, including bare SSE
+ keepalive/comment lines
+ # A tool that is slow but alive must not be punished
+ # for going quiet between progress updates, as long as
+ # something keeps proving the connection is alive.
+ max-duration timer — never reset, a hard ceiling regardless of activity
+
+ whichever timer fires first cancels the shared context and records why,
+ so the caller can tell "we gave up" apart from "the transport failed on its own"
+```
+
+The idle timer specifically re-derives elapsed time from a recorded timestamp rather
+than trusting that its callback firing means the call is genuinely idle: resetting a
+timer concurrently with its own callback is a documented race, so the callback instead
+checks real elapsed time and re-arms itself if it turns out to have fired early.
+
+**8. Auditing without blocking delivery (`internal/invocation`)**
+
+```
+recording one relayed event:
+ hand it to a fixed pool of shared background workers, sharded by call
+ # Not one goroutine per call: that design lets a stalled audit store leak
+ # one goroutine per concurrent streaming call. A fixed shared pool bounds
+ # that no matter how many calls are in flight — at the cost of the queue
+ # filling under sustained overload, which is reported, not hidden.
+ if the shard's queue is full: count it as dropped, do not block the relay
+ # Audit persistence must never slow down or stall live delivery to the
+ # agent. A full queue means "keep serving live traffic; this one row will
+ # not be written," and that trade is made visible, not silent.
+
+when the call ends:
+ wait briefly for this call's in-flight audit writes to finish
+ write one summary row: events seen / persisted / dropped / whether the wait itself timed out
+ # Per-event rows can be capped or dropped under load or by configuration;
+ # the summary row is the one place an operator can trust to see the true
+ # totals even when individual event rows are missing.
+```
+
+**9. Delivering to the agent, and auditing that delivery separately (`internal/api`)**
+
+```
+sink.StreamStarted():
+ commit to SSE: write status 200 and SSE headers, flush
+ start a background heartbeat loop
+ # Intermediaries (proxies, load balancers) commonly kill connections
+ # after ~60s of silence. A tool that is busy but has nothing to report
+ # yet would otherwise have its downstream leg cut for reasons that have
+ # nothing to do with the tool call itself.
+
+sink.Event(evt):
+ if it is a server-to-client request (sampling, elicitation, roots): drop it
+ # Atryum doesn't advertise those capabilities in initialize, and the
+ # agent has no channel to answer a request arriving on what it
+ # expects to be a tools/call response stream. Audited, never relayed.
+ else: write one SSE frame, one "data:" line per line of the payload
+ # A raw newline embedded in a single "data:" line breaks SSE framing
+ # for any compliant parser.
+
+sink.finishStream(terminal):
+ stop the heartbeat first
+ # Otherwise the heartbeat goroutine could still be writing to the
+ # response after the handler has already returned.
+ write the terminal frame as the last SSE event, using the agent's own
+ request id — never the fixed id Atryum sends upstream on the wire
+
+after the handler returns:
+ record whether that terminal write actually succeeded, as its own audit row,
+ separate from the invocation's succeeded/failed outcome
+ # The upstream tool call can succeed even though the agent never received
+ # the terminal frame (e.g. it disconnected moments earlier). Conflating
+ # the two would misreport a delivery failure as a tool failure, or the
+ # reverse.
+```
+
+#### Resource limits
+
+A single timeout is not enough for a stream. A long-running tool can be healthy as long
+as it continues to send progress. The relay therefore separates setup, inactivity,
+total-duration, and message-size limits:
+
+| Limit | Configuration | What it measures |
+|---|---|---|
+| Header timeout | `stream_header_timeout_seconds` | How long Atryum waits for HTTP session initialization and response headers, or stdio session initialization. |
+| Idle timeout | `stream_idle_timeout_seconds` | The longest allowed gap in upstream activity. It resets after every message, and after SSE keepalive/comment lines, so a busy-but-quiet upstream that heartbeats is not cut off. |
+| Maximum duration | `stream_max_duration_seconds` | A hard limit for the upstream execution phase, even if progress continues. |
+| Message size | `stream_max_message_bytes` | The largest accepted SSE event, stdio JSON-RPC line, or plain JSON response body. The default is 4 MiB. |
+
+The downstream connection has a separate per-write deadline. If the downstream client
+disconnects or stops reading, Atryum aborts that call instead of leaving a goroutine
+blocked forever. While the upstream is quiet, Atryum sends SSE comment heartbeats so
+proxies and load balancers do not mistake the downstream connection for an abandoned
+one. The audit trail distinguishes an upstream timeout (`stream_timeout`), a proven
+downstream disconnect (`stream_aborted_downstream`, set only when a write to the agent
+actually failed), and other transport failures. A bare request-context cancellation is
+recorded as `stream_canceled`: with no failed downstream write, a quiet agent
+disconnect and a server shutdown are indistinguishable, so the audit does not guess.
+
+#### Reliability guarantees and limits
+
+- **Plain JSON remains the fallback.** Atryum does not start the downstream SSE response
+ unless the client accepts SSE and the upstream sends an SSE response or a stdio
+ intermediate message.
+- **Approval still comes first.** No upstream call or downstream stream starts while a
+ tool call is waiting for approval.
+- **Retry is safe only before delivery.** Atryum may retry session setup before it has
+ relayed anything. After the first relayed event, retrying could duplicate visible
+ progress, so Atryum returns a terminal stream error instead.
+- **A started stream gets a terminal frame.** If execution fails after streaming has
+ begun, Atryum sends a final JSON-RPC error event instead of silently closing the
+ connection or trying to change the HTTP status.
+- **The per-call upstream stream can resume.** If it closes after providing an SSE event
+ ID but before the terminal response, Atryum reconnects with a `Last-Event-ID` header
+ naming the last processed event. If the upstream inclusively replays that event,
+ Atryum skips the duplicate. Reconnects use bounded exponential backoff with jitter;
+ the upstream cannot cause a zero-delay retry loop.
+- **The standalone stream reconnects but does not replay.** A transient connection
+ failure is retried with backoff. Session renewal opens a new standalone GET carrying
+ the new session ID. Because this path has no replay cursor, messages sent while it was
+ disconnected may still be missed.
+- **The downstream stream does not resume.** Atryum intentionally sends no downstream
+ SSE event IDs because it does not store each client's last processed position across
+ restarts.
+- **Audit storage cannot stall delivery.** Each intermediate event is offered to a
+ process-wide bounded, sharded dispatcher served by a fixed worker pool. A call does
+ not create its own audit goroutine, and one call's events remain ordered. Slow or
+ failed writes are counted in the completion audit row; they do not delay the live
+ event.
+- **Audit volume is bounded.** `stream_audit_max_events` limits rows per call, and
+ `stream_audit_max_event_bytes` limits the retained bytes in each row. Events beyond
+ those storage limits are still relayed.
+- **A slow call cannot block the shared standalone stream.** Each call has a bounded
+ progress buffer. If it fills, Atryum drops that call's standalone progress event so
+ other calls can continue. `standalone_events_dropped` in
+ `invocation.stream_completed` makes that loss visible.
+- **Unsupported standalone GET is not fatal.** If an upstream does not provide this
+ optional path, Atryum stops trying while the current group of calls remains active.
+ The tool call and its POST response continue normally.
+- **Input memory is bounded.** Atryum rejects an upstream message above
+ `stream_max_message_bytes` instead of accumulating an arbitrarily large SSE event,
+ stdio line, or plain JSON body in memory.
+- **Execution and delivery are audited separately.** A durable succeeded/failed
+ invocation describes the upstream outcome. `invocation.stream_delivery` separately
+ records whether the terminal SSE frame reached the downstream connection. A
+ persistence failure is recorded as `persistence_failed`, never as a successful stream.
+- **Multi-line SSE data stays valid.** Atryum writes one `data:` field for every payload
+ line and terminates the event with a blank line.
+- **Upstream requests are not forwarded as notifications.** Server-to-client requests
+ such as sampling or elicitation require a response channel Atryum does not broker.
+ Atryum audits and drops them instead.
+- **The relay has a kill switch.** Setting `stream_relay_enabled = false` restores
+ buffered, single-response behavior for every `tools/call`.
+
+#### Advanced implementation notes
+
+These details matter when changing the implementation but are not needed to understand
+the normal flow:
+
+- An HTTP response becomes committed when Atryum writes its headers or first bytes.
+ Before that point Atryum can still return plain JSON. After that point every success or
+ error must be a final SSE frame; the handler cannot switch response formats.
+- Heartbeats and tool events can be produced by different Go lightweight threads
+ (goroutines). A lock serializes downstream writes so two frames can never be
+ interleaved and corrupted.
+- Resetting an idle timer races with the timer callback if implemented naively. The
+ timeout guard checks the actual elapsed idle time before ending the call.
+- A timed-out stdio tool can leave child processes behind. On operating systems that
+ support process groups, Atryum terminates the entire group rather than only the direct
+ child.
+
## Decision-only calls
`Submit` persists and returns a decision without contacting an MCP server. An external
@@ -272,8 +1013,10 @@ The core tables are:
| `managed_agent_bindings`, `managed_agent_sessions` | Anthropic agent/session ownership and replay state |
| `external_sessions` | Atryum-minted harness sessions linking external invocations for cross-call evaluation context |
-Schema changes are ordered migrations under `internal/store/migrations/` and are
-applied at startup for both SQLite and PostgreSQL.
+Built-in schema definitions are ordered migrations under `pkg/migrations/`.
+`internal/store` applies them at startup and records their versions for both SQLite and
+PostgreSQL. Embedding programs can register separately tracked, namespaced migrations
+through `pkg/atryum.WithMigrations`; these run after all built-in migrations.
An invocation's row is the authoritative record of current state;
`invocation_events` is best-effort event history. The two writes are separate
@@ -309,5 +1052,8 @@ Inbound and upstream authentication are separate trust boundaries:
`admin_enabled = true` and the configured admin claim is present.
- Upstream MCP authentication is owned by `internal/mcp/auth_provider`; credentials and
OAuth tokens are never returned to the agent caller.
+- Routes registered through `pkg/atryum.WithRoutes` are deliberately outside Atryum's
+ built-in authentication middleware. The embedding program must authenticate and
+ authorize those routes itself.
- No-auth mode is a local deployment option. Identity supplied by a caller in this mode
is attribution, not a cryptographic ownership guarantee.
diff --git a/integrations/config/mcp-targets.yaml b/integrations/config/mcp-targets.yaml
index 00046f8b..71ecdcd1 100644
--- a/integrations/config/mcp-targets.yaml
+++ b/integrations/config/mcp-targets.yaml
@@ -43,4 +43,38 @@ mcp_targets:
expect_substrings: ["40"]
prompt: |
Use the calc_via_atryum MCP server's math tool to evaluate 17 + 23.
- Reply with only the numeric result, nothing else.
\ No newline at end of file
+ Reply with only the numeric result, nothing else.
+
+ - id: everything-streaming
+ display_name: "@modelcontextprotocol/server-everything (stdio)"
+ description: >
+ Official MCP reference "everything" server, run over stdio. Its
+ trigger-long-running-operation tool sends live progress notifications
+ before the terminal result, so this target exercises Atryum's SSE
+ relay end-to-end over a real HTTP connection to a real running atryum
+ process — the "real process boundary" fake_agent.py's other targets
+ don't cover (see docs/architecture.md, "Live SSE relay for tools/call").
+ Shares its upstream dependency with the Go mcp-everything-test — an
+ accepted, minor redundancy in exchange for not maintaining a second
+ hand-rolled MCP fixture. verify.streaming tells verify_upstream_direct
+ to run fake_agent.py with --stream instead of the plain one-shot check
+ the other targets use.
+ upstream:
+ name: everything-streaming
+ mode: stdio
+ command: npx
+ args: ["-y", "@modelcontextprotocol/server-everything"]
+ timeout_seconds: 30
+ enabled: true
+ verify:
+ tool: trigger-long-running-operation
+ arguments:
+ duration: 3
+ steps: 3
+ expect_substrings: ["Long running operation completed"]
+ streaming: true
+ min_progress_events: 3
+ prompt: |
+ Use the everything_streaming_via_atryum MCP server's
+ trigger-long-running-operation tool with duration 3 and steps 3.
+ Reply with only the final message, nothing else.
\ No newline at end of file
diff --git a/integrations/lib/atryum.sh b/integrations/lib/atryum.sh
index 978a481b..e4c89d78 100644
--- a/integrations/lib/atryum.sh
+++ b/integrations/lib/atryum.sh
@@ -46,12 +46,16 @@ render_atryum_config() {
local auth_id="$1" target_id="$2" template_path="$3" out_path="$4"
local upstreams
upstreams="$(render_upstream_block "$target_id")"
+ local py="${INTEGRATIONS_PYTHON:-python3}"
sed \
-e "s|__ATRYUM_PORT__|${ATRYUM_PORT}|g" \
-e "s|__RUN_DIR__|${RUN_DIR}|g" \
-e "s|__MOCK_OIDC_ISSUER__|${MOCK_OIDC_ISSUER:-http://127.0.0.1:${MOCK_OIDC_PORT}/realms/atryum}|g" \
"$template_path" \
- | awk -v block="$upstreams" '{gsub(/__UPSTREAMS_BLOCK__/, block); print}' \
+ | UPSTREAMS_BLOCK="$upstreams" "$py" -c '
+import os, sys
+sys.stdout.write(sys.stdin.read().replace("__UPSTREAMS_BLOCK__", os.environ["UPSTREAMS_BLOCK"]))
+' \
>"$out_path"
}
@@ -105,7 +109,7 @@ seed_auto_approve_rules() {
verify_upstream_direct() {
local target_id="$1"
local auth_id="${2:-no-auth}"
- local parsed server_name tool_name args_json expect_joined
+ local parsed server_name tool_name args_json expect_joined streaming min_progress_events
local py="${INTEGRATIONS_PYTHON:-python3}"
parsed="$("$py" - "$INTEGRATIONS_ROOT" "$target_id" "$py" <<'PY'
import json, subprocess, sys
@@ -120,14 +124,21 @@ print(u["name"])
print(v["tool"])
print(json.dumps(v["arguments"]))
print("|".join(v["expect_substrings"]))
+# Normalized to a plain "true"/"false" literal regardless of whether the
+# registry loader gave a real bool (PyYAML) or a string (the no-PyYAML
+# fallback parser in lib/registry.py) — see verify.streaming below.
+print("true" if v.get("streaming") else "false")
+print(v.get("min_progress_events") or 1)
PY
)"
server_name="$(echo "$parsed" | sed -n '1p')"
tool_name="$(echo "$parsed" | sed -n '2p')"
args_json="$(echo "$parsed" | sed -n '3p')"
expect_joined="$(echo "$parsed" | sed -n '4p')"
+ streaming="$(echo "$parsed" | sed -n '5p')"
+ min_progress_events="$(echo "$parsed" | sed -n '6p')"
- log "Direct MCP smoke via fake_agent.py (server=$server_name tool=$tool_name auth=$auth_id)"
+ log "Direct MCP smoke via fake_agent.py (server=$server_name tool=$tool_name auth=$auth_id streaming=$streaming)"
local bearer_args=()
case "$auth_id" in
oauth-client-credentials|oauth-dcr|static-bearer)
@@ -136,6 +147,15 @@ PY
bearer_args=(--bearer "$token")
;;
esac
+ # verify.streaming targets (e.g. everything-streaming) prove Atryum's SSE
+ # relay works over a real HTTP connection to a real running atryum
+ # process — fake_agent.py --stream fails the case if progress doesn't
+ # actually arrive live before the terminal result, not just on a wrong
+ # final answer.
+ local stream_args=()
+ if [[ "$streaming" == "true" ]]; then
+ stream_args=(--stream --min-progress-events "$min_progress_events")
+ fi
local output
output="$(
ATRYUM_URL="$ATRYUM_URL" \
@@ -143,7 +163,8 @@ PY
"$server_name" \
--tool "$tool_name" \
--arguments "$args_json" \
- "${bearer_args[@]}" 2>&1
+ "${bearer_args[@]+"${bearer_args[@]}"}" \
+ "${stream_args[@]+"${stream_args[@]}"}" 2>&1
)" || return 1
local part
diff --git a/integrations/lib/harness.sh b/integrations/lib/harness.sh
index beb880d4..f87bc228 100644
--- a/integrations/lib/harness.sh
+++ b/integrations/lib/harness.sh
@@ -227,7 +227,7 @@ configure_harness_mcp() {
amp_env+=(--env "MCP_REMOTE_HEADERS=${MCP_REMOTE_HEADERS}")
fi
if ! AMP_SETTINGS_FILE="$AMP_SETTINGS_FILE" amp mcp add "$mcp_alias" \
- "${amp_env[@]}" \
+ "${amp_env[@]+"${amp_env[@]}"}" \
-- npx -y mcp-remote "$mcp_url"; then
warn "amp mcp add failed for $AMP_SETTINGS_FILE"
return 1
@@ -256,7 +256,7 @@ configure_harness_mcp() {
grok_env+=(--env "MCP_REMOTE_HEADERS=${MCP_REMOTE_HEADERS}")
fi
HOME="$GROK_TEST_HOME" grok mcp add "$mcp_alias" \
- "${grok_env[@]}" \
+ "${grok_env[@]+"${grok_env[@]}"}" \
--command npx \
--args -y mcp-remote "$mcp_url" >/dev/null 2>&1 || {
warn "grok mcp add failed; writing config.toml fallback"
@@ -376,7 +376,7 @@ PY
python3 "$REPO_ROOT/scripts/fake_agent.py" mcp "$server_name" \
--tool "$tool_name" \
--arguments "$args_json" \
- "${bearer_args[@]}" 2>&1
+ "${bearer_args[@]+"${bearer_args[@]}"}" 2>&1
)"
rc=$?
else
diff --git a/internal/api/agent_server_helpers_test.go b/internal/api/agent_server_helpers_test.go
new file mode 100644
index 00000000..50dae11c
--- /dev/null
+++ b/internal/api/agent_server_helpers_test.go
@@ -0,0 +1,74 @@
+package api
+
+import (
+ "context"
+ "database/sql"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ "github.com/validmind/atryum/internal/config"
+ "github.com/validmind/atryum/internal/invocation"
+ "github.com/validmind/atryum/internal/invocation/policy"
+ "github.com/validmind/atryum/internal/mcp"
+ "github.com/validmind/atryum/internal/store"
+)
+
+// pollUntil calls cond repeatedly, sleeping interval between tries, until it
+// returns true or timeout elapses. Shared by every test in this package that
+// needs to wait for asynchronous state (a port accepting connections, a row
+// reaching a given status) instead of sleeping a fixed guess.
+func pollUntil(t *testing.T, timeout, interval time.Duration, cond func() bool) bool {
+ t.Helper()
+ deadline := time.Now().Add(timeout)
+ for time.Now().Before(deadline) {
+ if cond() {
+ return true
+ }
+ time.Sleep(interval)
+ }
+ return cond()
+}
+
+// newTestAgentServer builds the db -> resolver -> invocation.Service ->
+// Handler -> httptest.Server stack backing a single HTTP upstream, and
+// registers cleanup for the db and the server. timeoutSeconds drives both
+// the upstream's own request timeout and the service's default timeout, as
+// every existing call site already did identically. enableStreaming
+// installs the same stream options and audit limits every existing
+// streaming e2e test used; pass false for tests that don't exercise
+// streaming at all.
+func newTestAgentServer(t *testing.T, upstreamName, upstreamURL string, timeoutSeconds int, enableStreaming bool) (*httptest.Server, *invocation.Service) {
+ t.Helper()
+ db, err := sql.Open("sqlite", ":memory:")
+ if err != nil {
+ t.Fatalf("open db: %v", err)
+ }
+ t.Cleanup(func() { db.Close() })
+ if err := store.InitDB(db); err != nil {
+ t.Fatalf("InitDB: %v", err)
+ }
+ serverRepo := store.NewServerRepo(db)
+ resolver := mcp.NewResolver(serverRepo, config.Config{
+ Upstreams: []config.UpstreamConfig{{Name: upstreamName, Mode: "http", BaseURL: upstreamURL, Enabled: true, TimeoutSeconds: timeoutSeconds}},
+ })
+ if err := resolver.BootstrapIfEmpty(context.Background()); err != nil {
+ t.Fatal(err)
+ }
+ svc := invocation.NewService(
+ store.NewInvocationRepo(db), store.NewEventRepo(db), resolver, mcp.NewHTTPClient(),
+ policy.AlwaysApproveProvider{}, time.Duration(timeoutSeconds)*time.Second, nil, nil, nil, nil,
+ )
+ if enableStreaming {
+ svc.SetStreamOptions(
+ mcp.StreamOptions{HeaderTimeout: 10 * time.Second, IdleTimeout: 10 * time.Second, MaxDuration: 60 * time.Second},
+ invocation.StreamAuditLimits{MaxEvents: 100, MaxEventBytes: 4096},
+ )
+ }
+
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ agentServer := httptest.NewServer(h.Routes())
+ t.Cleanup(agentServer.Close)
+
+ return agentServer, svc
+}
diff --git a/internal/api/handlers.go b/internal/api/handlers.go
index 563f9d78..cd3b297e 100644
--- a/internal/api/handlers.go
+++ b/internal/api/handlers.go
@@ -42,6 +42,8 @@ const upstreamMCPOAuthCallbackPath = "/api/v1/mcp/oauth/callback"
type service interface {
Invoke(ctx context.Context, req invocation.CreateInvocationRequest) (invocation.InvocationResponse, error)
+ InvokeStreaming(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error)
+ RecordStreamDelivery(ctx context.Context, invocationID, status, message string) error
ListTools(ctx context.Context, server string) ([]mcp.Tool, error)
Get(ctx context.Context, id string) (invocation.InvocationResponse, error)
List(ctx context.Context, filter invocation.InvocationListFilter) (invocation.InvocationListResponse, error)
@@ -162,6 +164,13 @@ type Handler struct {
authValidator *auth.Validator
apiKeyAuth auth.APIKeyConfig
+ // streamRelayEnabled is the kill-switch for the tools/call SSE relay
+ // (see handleMCPProxy). The downstream relay only activates when the
+ // agent accepts SSE and the upstream client enters live mode, either
+ // from an HTTP SSE response or an intermediate stdio message. It can be
+ // disabled globally without a rollback.
+ streamRelayEnabled bool
+
// clientInfoCache remembers the most recent `initialize.clientInfo`
// per MCP session key so that subsequent tools/call requests on the
// same session can attach client_name / client_version. The key is the
@@ -813,7 +822,18 @@ func NewHandler(svc service, serverSvc serverService, policyRegistry *policy.Reg
if f, ok := svc.(mcpEnvelopeForwarder); ok {
forwarder = f
}
- return &Handler{svc: svc, serverSvc: serverSvc, policyRegistry: policyRegistry, rulesRepo: rules, agentsRepo: agents, agentSyncSettingsRepo: agentSyncSettings, llmConfigsRepo: llmConfigs, backendClient: bc, summarizeClient: bc, localSummarizer: localSummarizer, syncAgentsFn: syncAgents, forwarder: forwarder, staticHTTP: http.FileServer(http.FS(staticSub)), debug: debug, clientInfoCache: make(map[string]clientInfoSnapshot)}
+ return &Handler{svc: svc, serverSvc: serverSvc, policyRegistry: policyRegistry, rulesRepo: rules, agentsRepo: agents, agentSyncSettingsRepo: agentSyncSettings, llmConfigsRepo: llmConfigs, backendClient: bc, summarizeClient: bc, localSummarizer: localSummarizer, syncAgentsFn: syncAgents, forwarder: forwarder, staticHTTP: http.FileServer(http.FS(staticSub)), debug: debug, clientInfoCache: make(map[string]clientInfoSnapshot), streamRelayEnabled: true}
+}
+
+// SetStreamRelayEnabled toggles the tools/call SSE relay kill-switch (on by
+// default). Disabling it forces every tools/call back to the buffered
+// application/json path regardless of what the agent's Accept header or the
+// upstream's response content-type would otherwise allow. Call it during
+// startup wiring only: the flag is an unsynchronized bool read on every
+// request, so flipping it while serving is a data race, not a runtime
+// toggle.
+func (h *Handler) SetStreamRelayEnabled(enabled bool) {
+ h.streamRelayEnabled = enabled
}
// SetAuthValidator installs the inbound auth validator. When non-nil, the
@@ -1431,6 +1451,7 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server
var params struct {
Name string `json:"name"`
Arguments map[string]any `json:"arguments"`
+ Meta map[string]any `json:"_meta"`
}
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
h.writeRPCError(w, req.ID, -32602, "invalid params")
@@ -1448,7 +1469,7 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server
h.handleMCPPlanGet(w, r, req.ID, params.Arguments)
return
}
- toolReq := invocation.CreateInvocationRequest{Server: server, Tool: params.Name, Input: params.Arguments}
+ toolReq := invocation.CreateInvocationRequest{Server: server, Tool: params.Name, Input: params.Arguments, Meta: params.Meta}
if requestID != "" {
toolReq.RequestID = stringPtr(requestID)
}
@@ -1459,22 +1480,81 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server
toolReq.ClientName = snap.Name
toolReq.ClientVersion = snap.Version
}
- resp, err := h.svc.Invoke(r.Context(), toolReq)
+ // A stream-capable agent gets a relay sink. The downstream response
+ // switches to SSE only when the upstream client calls StreamStarted:
+ // for an HTTP SSE response or the first intermediate stdio message.
+ // Until then nothing is written, so a buffered upstream response
+ // still follows the ordinary JSON path below.
+ var sink *sseRelaySink
+ if flusher, ok := w.(http.Flusher); ok && h.streamRelayEnabled && acceptsEventStream(r) {
+ sink = newSSERelaySink(w, flusher)
+ // Structural guarantee that the heartbeat goroutine can never
+ // outlive this handler, even on a panic path. A no-op on normal
+ // paths, where finishStream already stopped it.
+ defer sink.stopHeartbeat()
+ }
+ var resp invocation.InvocationResponse
+ var err error
+ if sink != nil {
+ resp, err = h.svc.InvokeStreaming(r.Context(), toolReq, sink)
+ } else {
+ resp, err = h.svc.Invoke(r.Context(), toolReq)
+ }
if err != nil {
+ // This branch IS reachable with an already-started stream:
+ // besides pre-execution validation failures (sink untouched),
+ // InvokeStreaming also returns an error if persisting the
+ // result fails after the stream completed. Once headers are
+ // out, the only spec-correct way to end the exchange is a
+ // terminal JSON-RPC error as the final SSE frame — never
+ // writeRPCError (a second WriteHeader), and never a bare
+ // close (the agent would be left with a request that ended
+ // without any response).
+ if sink != nil && sink.started {
+ // The full error is logged server-side (unconditionally —
+ // not debugf: when persistence fails, this log line is the
+ // only place the real error text survives, since the audit
+ // row can only say persistence_failed); the wire gets a
+ // static message. The only errors reachable with a started
+ // stream are internal finalization failures (result
+ // persistence), whose text can carry SQL/driver detail the
+ // agent has no business seeing.
+ log.Printf("[mcpToolsCall] finalize after stream started failed server=%s tool=%s invocation=%s: %v", server, params.Name, resp.InvocationID, err)
+ // Trace this path like the success path does: a stream that
+ // relayed events and then failed finalization is exactly the
+ // call an operator will want to find.
+ _ = h.emitTraceEvent(r.Context(), server, "mcp.tools.call", map[string]any{"request_id": requestID, "status": resp.Status, "invocation_id": resp.InvocationID, "tool": params.Name, "streamed": true, "stream_events": sink.eventCount, "finalize_failed": true})
+ errBody, _ := json.Marshal(map[string]any{"code": -32000, "message": "failed to finalize invocation"})
+ terminal, _ := json.Marshal(jsonRPCResponse{JSONRPC: "2.0", ID: req.ID, Error: errBody})
+ deliveryErr := sink.finishStream(terminal)
+ h.recordStreamDelivery(r.Context(), resp.InvocationID, deliveryErr)
+ return
+ }
h.writeRPCError(w, req.ID, -32000, err.Error())
return
}
- tracePayload := map[string]any{"request_id": requestID, "status": resp.Status, "invocation_id": resp.InvocationID, "tool": params.Name}
+ streamed := sink != nil && sink.started
+ tracePayload := map[string]any{"request_id": requestID, "status": resp.Status, "invocation_id": resp.InvocationID, "tool": params.Name, "streamed": streamed}
+ if streamed {
+ tracePayload["stream_events"] = sink.eventCount
+ }
_ = h.emitTraceEvent(r.Context(), server, "mcp.tools.call", tracePayload)
- if len(resp.Error) > 0 {
- result := normalizeToolCallResult(resp.Error, true)
- if resp.Status == invocation.StatusDenied {
- result = h.appendRulesContextToToolResult(r.Context(), result, auth.AgentIDFromContext(r.Context()), server, params.Name)
- }
- h.writeRPCResult(w, req.ID, result)
+
+ result := h.toolCallResult(r.Context(), server, params.Name, resp)
+ if streamed {
+ // Headers were already sent when the stream started: the
+ // terminal response is the final SSE frame, rewritten to the
+ // agent's own request id (upstream always sees id "1" — see
+ // mcp.Client.invokeHTTP/invokeStream). Never writeRPCResult or
+ // WriteHeader here. finishStream also stops the heartbeat
+ // goroutine — mandatory before returning from the handler.
+ body, _ := json.Marshal(result)
+ terminal, _ := json.Marshal(jsonRPCResponse{JSONRPC: "2.0", ID: req.ID, Result: body})
+ deliveryErr := sink.finishStream(terminal)
+ h.recordStreamDelivery(r.Context(), resp.InvocationID, deliveryErr)
return
}
- h.writeRPCResult(w, req.ID, normalizeToolCallResult(resp.Result, false))
+ h.writeRPCResult(w, req.ID, result)
default:
forwarded, forwardedOK := h.forwardProxyEnvelope(r.Context(), server, req, protocolVersion)
if !forwardedOK {
@@ -1903,6 +1983,22 @@ func effectiveActionForTool(rules []store.Rule, server, tool, agentCUID string)
return invocation.RuleActionHumanApproval, ""
}
+// toolCallResult builds the MCP tool-result value for a completed
+// invocation, appending rules context to a denial exactly as the
+// non-streaming path always has. Shared by both the buffered and the SSE
+// relay branches of the tools/call handler so the result shape is
+// identical either way.
+func (h *Handler) toolCallResult(ctx context.Context, server, tool string, resp invocation.InvocationResponse) any {
+ if len(resp.Error) > 0 {
+ result := normalizeToolCallResult(resp.Error, true)
+ if resp.Status == invocation.StatusDenied {
+ result = h.appendRulesContextToToolResult(ctx, result, auth.AgentIDFromContext(ctx), server, tool)
+ }
+ return result
+ }
+ return normalizeToolCallResult(resp.Result, false)
+}
+
// appendRulesContextToToolResult adds an extra text content block to a denied
// tool call result describing the applicable rules and effective action.
func (h *Handler) appendRulesContextToToolResult(ctx context.Context, result any, agentID, server, tool string) any {
@@ -4795,3 +4891,21 @@ func (h *Handler) debugf(format string, args ...any) {
}
log.Printf("[mcp] "+format, args...)
}
+
+func (h *Handler) recordStreamDelivery(ctx context.Context, invocationID string, deliveryErr error) {
+ if invocationID == "" {
+ return
+ }
+ status := "succeeded"
+ message := ""
+ if deliveryErr != nil {
+ status = "failed"
+ message = deliveryErr.Error()
+ h.debugf("mcp terminal stream delivery failed invocation_id=%s err=%v", invocationID, deliveryErr)
+ }
+ persistCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Second)
+ defer cancel()
+ if err := h.svc.RecordStreamDelivery(persistCtx, invocationID, status, message); err != nil {
+ h.debugf("mcp stream delivery audit failed invocation_id=%s err=%v", invocationID, err)
+ }
+}
diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go
index d930b983..c5e28039 100644
--- a/internal/api/handlers_test.go
+++ b/internal/api/handlers_test.go
@@ -48,19 +48,29 @@ type stubService struct {
recordReq *invocation.ExternalExecutionUpdate
recordCtx context.Context
+ streamDeliveryInvocationID string
+ streamDeliveryStatus string
+ streamDeliveryMessage string
+
createSessionReq *invocation.CreateSessionRequest
createSessionAgentID string
- plan invocation.Plan
- planErr error
- planSubmitReq *invocation.PlanSubmitRequest
- planApproveID string
- planTTL int
- planDenyID string
- planDenyMsg string
- planReviseID string
- planFeedback string
- planExpireID string
- planCancelID string
+
+ // invokeStreamingFn, when set, overrides InvokeStreaming's behavior so a
+ // test can simulate the sink actually being used. Nil means
+ // InvokeStreaming behaves exactly like Invoke (sink untouched).
+ invokeStreamingFn func(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error)
+
+ plan invocation.Plan
+ planErr error
+ planSubmitReq *invocation.PlanSubmitRequest
+ planApproveID string
+ planTTL int
+ planDenyID string
+ planDenyMsg string
+ planReviseID string
+ planFeedback string
+ planExpireID string
+ planCancelID string
}
func (s *stubService) Invoke(ctx context.Context, req invocation.CreateInvocationRequest) (invocation.InvocationResponse, error) {
@@ -68,6 +78,27 @@ func (s *stubService) Invoke(ctx context.Context, req invocation.CreateInvocatio
s.invokedCtx = ctx
return s.invoke, s.invErr
}
+
+// InvokeStreaming records the call like Invoke and, unless invokeStreamingFn
+// is set, never touches sink — matching the real Service.InvokeStreaming's
+// nil-sink behavior for every test that doesn't care about streaming.
+// invokeStreamingFn lets a test simulate the sink actually being used (e.g.
+// calling StreamStarted/Event) to exercise handleMCPProxy's relay path.
+func (s *stubService) InvokeStreaming(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error) {
+ s.invokedReq = &req
+ s.invokedCtx = ctx
+ if s.invokeStreamingFn != nil {
+ return s.invokeStreamingFn(ctx, req, sink)
+ }
+ return s.invoke, s.invErr
+}
+
+func (s *stubService) RecordStreamDelivery(_ context.Context, invocationID, status, message string) error {
+ s.streamDeliveryInvocationID = invocationID
+ s.streamDeliveryStatus = status
+ s.streamDeliveryMessage = message
+ return nil
+}
func (s *stubService) ListTools(context.Context, string) ([]mcp.Tool, error) {
return s.tools, s.listErr
}
@@ -1263,6 +1294,23 @@ func TestMCPToolsCallInterceptsInvocation(t *testing.T) {
}
}
+func TestMCPToolsCallForwardsMetaToInvocation(t *testing.T) {
+ now := time.Now().UTC()
+ svc := &stubService{invoke: invocation.InvocationResponse{InvocationID: "inv_123", ServerName: "demo", ToolName: "demo_tool", Status: invocation.StatusSucceeded, SubmittedAt: now, CompletedAt: &now, Result: json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`)}}
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"demo_tool","arguments":{"a":1},"_meta":{"progressToken":"tok-7"}}}`))
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ if svc.invokedReq == nil {
+ t.Fatal("expected invocation request")
+ }
+ if got := svc.invokedReq.Meta["progressToken"]; got != "tok-7" {
+ t.Fatalf("expected _meta.progressToken to reach the invocation request, got %#v", svc.invokedReq.Meta)
+ }
+}
+
func TestMCPRulesToolReturnsAgentRulesWithoutInvocation(t *testing.T) {
rules := &stubRulesRepo{rules: []store.Rule{
{ID: "read-auto", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"demo"}, ToolPatterns: []string{"Read"}, Enabled: true, Order: 0},
diff --git a/internal/api/mcp_everything_test.go b/internal/api/mcp_everything_test.go
new file mode 100644
index 00000000..09a1b005
--- /dev/null
+++ b/internal/api/mcp_everything_test.go
@@ -0,0 +1,152 @@
+//go:build mcpeverything
+
+// Real end-to-end test against the official MCP reference "everything"
+// server (@modelcontextprotocol/server-everything), run over its Streamable
+// HTTP transport, through a real running Atryum handler stack. Requires
+// Node/npm (npx) and, on first run, network access to fetch the package —
+// excluded from `go test ./...` via this build tag so the default suite
+// stays fast and hermetic. Run explicitly:
+//
+// just mcp-everything-test
+//
+// or:
+//
+// go test -tags mcpeverything ./internal/api -run TestMCPToolsCallAgainstRealEverythingServer -v
+//
+// trigger-long-running-operation's progress notification carries
+// relatedRequestId, so the reference server routes it onto the same
+// connection as the tools/call response — this test proves that path. It
+// does not exercise the standalone-channel path (see
+// internal/mcp/client_test.go's TestInvokeStreamStandaloneStream* tests,
+// and mcp_standalone_stream_test.go's live fixture for that).
+package api
+
+import (
+ "bufio"
+ "fmt"
+ "net"
+ "net/http"
+ "os/exec"
+ "strings"
+ "testing"
+ "time"
+)
+
+// startEverythingServer launches the real @modelcontextprotocol/server-everything
+// package over its Streamable HTTP transport on a free local port, and waits
+// for it to accept connections before returning. Skips the test entirely if
+// npx isn't available, rather than failing.
+func startEverythingServer(t *testing.T) (baseURL string) {
+ t.Helper()
+ if _, err := exec.LookPath("npx"); err != nil {
+ t.Skip("npx not found in PATH; skipping real server-everything e2e test")
+ }
+
+ port := freePort(t)
+ cmd := exec.Command("npx", "-y", "@modelcontextprotocol/server-everything", "streamableHttp")
+ cmd.Env = append(cmd.Environ(), fmt.Sprintf("PORT=%d", port))
+ configureExternalProcessGroup(cmd)
+ if err := cmd.Start(); err != nil {
+ t.Fatalf("start server-everything: %v", err)
+ }
+ t.Cleanup(func() {
+ // npx execs the actual server binary as a child process; killing just
+ // the npx process it directly spawned would leak that child (the same
+ // process-group leak internal/mcp's stdio handling guards against).
+ killExternalProcessGroup(cmd.Process.Pid)
+ _ = cmd.Wait()
+ })
+
+ baseURL = fmt.Sprintf("http://127.0.0.1:%d/mcp", port)
+ if !pollUntil(t, 30*time.Second, 200*time.Millisecond, func() bool {
+ conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 200*time.Millisecond)
+ if err != nil {
+ return false
+ }
+ _ = conn.Close()
+ return true
+ }) {
+ t.Fatalf("server-everything did not start listening on port within 30s (npx may need network access to fetch the package on first run)")
+ }
+ return baseURL
+}
+
+// TestMCPToolsCallAgainstRealEverythingServer proves the relay is live, not
+// buffered, by pinning the wall-clock gaps between frames read off a real
+// subprocess:
+//
+// t=0 POST tools/call trigger-long-running-operation
+// duration=3, steps=3, progressToken=real-e2e-token
+// t=~1s notifications/progress -> progressTimes[0]
+// t=~2s notifications/progress -> progressTimes[1]
+// t=~3s notifications/progress -> progressTimes[2]
+// t=~3s+ terminal "Long running operation completed" -> terminalTime
+//
+// The assertions pin the gaps, not just the count: terminalTime minus
+// progressTimes[0] must be >= 1500ms (a buffered-then-replayed body would
+// show the terminal moments after the first update instead), and each
+// progressTimes[i] minus progressTimes[i-1] must be >= 500ms (the three
+// updates are really spread over real seconds, not emitted back-to-back
+// once the tool finished).
+func TestMCPToolsCallAgainstRealEverythingServer(t *testing.T) {
+ baseURL := startEverythingServer(t)
+
+ agentServer, _ := newTestAgentServer(t, "everything", baseURL, 30, true)
+
+ reqBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"trigger-long-running-operation","arguments":{"duration":3,"steps":3},"_meta":{"progressToken":"real-e2e-token"}}}`
+ req, err := http.NewRequest(http.MethodPost, agentServer.URL+"/mcp/everything", strings.NewReader(reqBody))
+ if err != nil {
+ t.Fatal(err)
+ }
+ req.Header.Set("Content-Type", "application/json")
+ req.Header.Set("Accept", "application/json, text/event-stream")
+
+ resp, err := http.DefaultClient.Do(req)
+ if err != nil {
+ t.Fatalf("agent request: %v", err)
+ }
+ defer resp.Body.Close()
+
+ if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
+ t.Fatalf("expected text/event-stream, got %q", ct)
+ }
+
+ reader := bufio.NewReader(resp.Body)
+ var progressTimes []time.Time
+ var terminal sseEventFrame
+ var terminalTime time.Time
+ for {
+ frame := readNextSSEFrame(t, reader)
+ if strings.Contains(frame.data, "notifications/progress") {
+ progressTimes = append(progressTimes, time.Now())
+ if !strings.Contains(frame.data, "real-e2e-token") {
+ t.Fatalf("expected the agent's own progressToken restored, got %q", frame.data)
+ }
+ continue
+ }
+ terminal = frame
+ terminalTime = time.Now()
+ break
+ }
+
+ if len(progressTimes) != 3 {
+ t.Fatalf("expected 3 live progress notifications from the real everything server, got %d", len(progressTimes))
+ }
+ if !strings.Contains(terminal.data, "Long running operation completed") {
+ t.Fatalf("expected the real terminal result, got %q", terminal.data)
+ }
+
+ // The whole point: the agent must see the first update well before the
+ // terminal response could exist (duration=3s/steps=3 means the server
+ // hasn't even finished sleeping when the first update arrives). If Atryum
+ // buffered the whole response and replayed it at the end, every frame
+ // would land within milliseconds of each other instead of spread ~1s apart.
+ if gap := terminalTime.Sub(progressTimes[0]); gap < 1500*time.Millisecond {
+ t.Fatalf("expected several seconds between the first live update and the terminal result (proving live delivery, not buffering), got %s", gap)
+ }
+ for i := 1; i < len(progressTimes); i++ {
+ if gap := progressTimes[i].Sub(progressTimes[i-1]); gap < 500*time.Millisecond {
+ t.Fatalf("expected a real ~1s gap between progress notification %d and %d, got %s", i, i+1, gap)
+ }
+ }
+}
diff --git a/internal/api/mcp_external_fixture_helpers_test.go b/internal/api/mcp_external_fixture_helpers_test.go
new file mode 100644
index 00000000..3e5a6d2f
--- /dev/null
+++ b/internal/api/mcp_external_fixture_helpers_test.go
@@ -0,0 +1,22 @@
+//go:build mcpeverything || mcpstandalone
+
+package api
+
+import (
+ "net"
+ "testing"
+)
+
+// freePort allocates an ephemeral TCP port for a spawned fixture server to
+// listen on, closing the probe listener immediately so the port is free
+// again by the time the caller passes it to the child process. Shared by
+// mcp_everything_test.go and mcp_standalone_stream_test.go.
+func freePort(t *testing.T) int {
+ t.Helper()
+ l, err := net.Listen("tcp", "127.0.0.1:0")
+ if err != nil {
+ t.Fatalf("allocate free port: %v", err)
+ }
+ defer l.Close()
+ return l.Addr().(*net.TCPAddr).Port
+}
diff --git a/internal/api/mcp_external_process_unix_test.go b/internal/api/mcp_external_process_unix_test.go
new file mode 100644
index 00000000..6c8f21f3
--- /dev/null
+++ b/internal/api/mcp_external_process_unix_test.go
@@ -0,0 +1,33 @@
+//go:build (mcpeverything || mcpstandalone) && !windows
+
+package api
+
+import (
+ "os/exec"
+ "syscall"
+)
+
+// configureExternalProcessGroup puts cmd in its own process group so
+// killExternalProcessGroup can terminate it and any descendants it spawns
+// (e.g. npx execs the actual server binary as a child process; uv run may
+// do likewise) in one signal. Without this, killing just the directly
+// spawned process can leave that child running — the same class of leak
+// internal/mcp's stdio upstream handling guards against (see
+// stdio_process_unix.go) for the same reason. Shared by both real-server
+// e2e fixtures (mcp_everything_test.go, mcp_standalone_stream_test.go).
+func configureExternalProcessGroup(cmd *exec.Cmd) {
+ cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+}
+
+// killExternalProcessGroup sends SIGKILL to the process group headed by
+// pid, verifying pid still heads that group (Setpgid made it the group
+// leader at spawn time) immediately before signaling — see
+// internal/mcp/stdio_process_unix.go's killStdioProcessGroup for the full
+// reasoning on why this check matters and what residual race it leaves.
+func killExternalProcessGroup(pid int) {
+ pgid, err := syscall.Getpgid(pid)
+ if err != nil || pgid != pid {
+ return
+ }
+ _ = syscall.Kill(-pid, syscall.SIGKILL)
+}
diff --git a/internal/api/mcp_external_process_windows_test.go b/internal/api/mcp_external_process_windows_test.go
new file mode 100644
index 00000000..78860eb0
--- /dev/null
+++ b/internal/api/mcp_external_process_windows_test.go
@@ -0,0 +1,15 @@
+//go:build (mcpeverything || mcpstandalone) && windows
+
+package api
+
+import "os/exec"
+
+// configureExternalProcessGroup is a no-op on Windows: process-group
+// signaling works differently there, and these test fixtures aren't
+// exercised on Windows CI. killExternalProcessGroup falls back to killing
+// just the directly-spawned process, which may leak a grandchild on this
+// platform — see mcp_external_process_unix_test.go for the Unix behavior
+// this stands in for.
+func configureExternalProcessGroup(cmd *exec.Cmd) {}
+
+func killExternalProcessGroup(pid int) {}
diff --git a/internal/api/mcp_standalone_stream_test.go b/internal/api/mcp_standalone_stream_test.go
new file mode 100644
index 00000000..b6e8ed31
--- /dev/null
+++ b/internal/api/mcp_standalone_stream_test.go
@@ -0,0 +1,155 @@
+//go:build mcpstandalone
+
+// Real end-to-end test against a real MCP Python SDK (FastMCP) server run
+// via uv, through a real running Atryum handler stack. Requires uv
+// (https://docs.astral.sh/uv/) on PATH; uv resolves the `mcp` package into
+// an ephemeral environment on first run, which needs network access.
+// Excluded from `go test ./...` via this build tag so the default suite
+// stays fast and hermetic. Run explicitly:
+//
+// just mcp-standalone-stream-test
+//
+// or:
+//
+// go test -tags mcpstandalone ./internal/api -run TestMCPToolsCallAgainstRealStandaloneStreamServer -v
+//
+// FastMCP's Context.report_progress doesn't attribute its notification to
+// the request that triggered it (no related_request_id), so the server
+// routes every progress update to the standalone SSE stream — never to the
+// tools/call POST response body. This is the complement to
+// mcp_everything_test.go, which proves the opposite case (progress on the
+// same connection as the call). See internal/api/testdata/mcp_standalone_fixture/server.py
+// for the server, and internal/mcp/client_test.go's
+// TestInvokeStreamStandaloneStream* tests for the hermetic, hand-fixtured
+// version of this same code path.
+package api
+
+import (
+ "bufio"
+ "fmt"
+ "net"
+ "net/http"
+ "os/exec"
+ "strings"
+ "testing"
+ "time"
+)
+
+// startStandaloneFixtureServer launches the real FastMCP-based fixture
+// server via `uv run --with mcp python3 ...` on a free local port, and waits
+// for it to accept connections before returning. Skips the test entirely if
+// uv isn't available, rather than failing.
+func startStandaloneFixtureServer(t *testing.T) (baseURL string) {
+ t.Helper()
+ if _, err := exec.LookPath("uv"); err != nil {
+ t.Skip("uv not found in PATH; skipping real standalone-stream e2e test")
+ }
+
+ port := freePort(t)
+ cmd := exec.Command("uv", "run", "--with", "mcp", "python3", "testdata/mcp_standalone_fixture/server.py")
+ cmd.Env = append(cmd.Environ(), fmt.Sprintf("PORT=%d", port))
+ configureExternalProcessGroup(cmd)
+ if err := cmd.Start(); err != nil {
+ t.Fatalf("start standalone fixture server: %v", err)
+ }
+ t.Cleanup(func() {
+ // uv run may spawn python as a child process rather than exec'ing into
+ // it directly; killing just the uv process it directly spawned could
+ // leak that child (the same process-group leak internal/mcp's stdio
+ // handling guards against).
+ killExternalProcessGroup(cmd.Process.Pid)
+ _ = cmd.Wait()
+ })
+
+ baseURL = fmt.Sprintf("http://127.0.0.1:%d/mcp", port)
+ if !pollUntil(t, 30*time.Second, 200*time.Millisecond, func() bool {
+ conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 200*time.Millisecond)
+ if err != nil {
+ return false
+ }
+ _ = conn.Close()
+ return true
+ }) {
+ t.Fatalf("standalone fixture server did not start listening on port within 30s (uv may need network access to resolve the mcp package on first run)")
+ }
+ return baseURL
+}
+
+// TestMCPToolsCallAgainstRealStandaloneStreamServer is the standalone-stream
+// counterpart to mcp_everything_test.go's TestMCPToolsCallAgainstRealEverythingServer:
+// the same wall-clock-gap proof of live delivery, but every progress
+// notification here arrives on the standalone GET stream rather than the
+// tools/call POST response, since FastMCP's report_progress never attributes
+// progress to a related_request_id:
+//
+// t=0 POST tools/call slow_streaming_task
+// steps=3, delay_seconds=1, progressToken=standalone-e2e-token
+// t=~1s notifications/progress (via standalone GET) -> progressTimes[0]
+// t=~2s notifications/progress (via standalone GET) -> progressTimes[1]
+// t=~3s notifications/progress (via standalone GET) -> progressTimes[2]
+// t=~3s+ terminal "done after 3 real progress notifications" -> terminalTime
+//
+// Same two assertions as the everything-server test: terminalTime minus
+// progressTimes[0] >= 1500ms, and each consecutive gap >= 500ms. Reading only
+// the tools/call POST response (the pre-fix behavior) would show 0 progress
+// notifications here — see the comment above the read loop below.
+func TestMCPToolsCallAgainstRealStandaloneStreamServer(t *testing.T) {
+ baseURL := startStandaloneFixtureServer(t)
+
+ agentServer, _ := newTestAgentServer(t, "standalone-fixture", baseURL, 30, true)
+
+ reqBody := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"slow_streaming_task","arguments":{"steps":3,"delay_seconds":1},"_meta":{"progressToken":"standalone-e2e-token"}}}`
+ req, err := http.NewRequest(http.MethodPost, agentServer.URL+"/mcp/standalone-fixture", strings.NewReader(reqBody))
+ if err != nil {
+ t.Fatal(err)
+ }
+ req.Header.Set("Content-Type", "application/json")
+ req.Header.Set("Accept", "application/json, text/event-stream")
+
+ resp, err := http.DefaultClient.Do(req)
+ if err != nil {
+ t.Fatalf("agent request: %v", err)
+ }
+ defer resp.Body.Close()
+
+ if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
+ t.Fatalf("expected text/event-stream, got %q", ct)
+ }
+
+ reader := bufio.NewReader(resp.Body)
+ var progressTimes []time.Time
+ var terminal sseEventFrame
+ var terminalTime time.Time
+ for {
+ frame := readNextSSEFrame(t, reader)
+ if strings.Contains(frame.data, "notifications/progress") {
+ progressTimes = append(progressTimes, time.Now())
+ if !strings.Contains(frame.data, "standalone-e2e-token") {
+ t.Fatalf("expected the agent's own progressToken restored, got %q", frame.data)
+ }
+ continue
+ }
+ terminal = frame
+ terminalTime = time.Now()
+ break
+ }
+
+ // If Atryum only read the tools/call POST response (the pre-fix
+ // behavior), this would be 0: FastMCP's report_progress sends every
+ // update exclusively on the standalone stream, never here.
+ if len(progressTimes) != 3 {
+ t.Fatalf("expected 3 live progress notifications relayed from the real standalone stream, got %d", len(progressTimes))
+ }
+ if !strings.Contains(terminal.data, "done after 3 real progress notifications") {
+ t.Fatalf("expected the real terminal result, got %q", terminal.data)
+ }
+
+ if gap := terminalTime.Sub(progressTimes[0]); gap < 1500*time.Millisecond {
+ t.Fatalf("expected several seconds between the first live update and the terminal result (proving live delivery, not buffering), got %s", gap)
+ }
+ for i := 1; i < len(progressTimes); i++ {
+ if gap := progressTimes[i].Sub(progressTimes[i-1]); gap < 500*time.Millisecond {
+ t.Fatalf("expected a real ~1s gap between progress notification %d and %d, got %s", i, i+1, gap)
+ }
+ }
+}
diff --git a/internal/api/sse_relay.go b/internal/api/sse_relay.go
new file mode 100644
index 00000000..5ff29e07
--- /dev/null
+++ b/internal/api/sse_relay.go
@@ -0,0 +1,206 @@
+package api
+
+import (
+ "bytes"
+ "fmt"
+ "io"
+ "net/http"
+ "sync"
+ "time"
+
+ "github.com/validmind/atryum/internal/mcp"
+)
+
+// splitSSELines splits payload data on every line terminator the SSE spec
+// recognizes — CRLF, LF, or bare CR — so a raw CR never reaches the wire
+// inside a data field, where a compliant parser treats it as a frame break.
+// This is reachable, not theoretical: JSON legally allows CR as inter-token
+// whitespace, so a stdio upstream can emit one inside an otherwise valid
+// message.
+func splitSSELines(data []byte) [][]byte {
+ normalized := bytes.ReplaceAll(data, []byte("\r\n"), []byte("\n"))
+ normalized = bytes.ReplaceAll(normalized, []byte("\r"), []byte("\n"))
+ return bytes.Split(normalized, []byte("\n"))
+}
+
+// writeSSEEvent writes and flushes one SSE frame. Each payload line needs its
+// own data field; otherwise a compliant parser drops continuation lines. It
+// deliberately omits event IDs because the downstream relay is not resumable.
+func writeSSEEvent(w io.Writer, flusher http.Flusher, event string, data []byte) error {
+ if event != "" {
+ if _, err := fmt.Fprintf(w, "event: %s\n", event); err != nil {
+ return err
+ }
+ }
+ for _, line := range splitSSELines(data) {
+ if _, err := fmt.Fprintf(w, "data: %s\n", line); err != nil {
+ return err
+ }
+ }
+ if _, err := fmt.Fprint(w, "\n"); err != nil {
+ return err
+ }
+ flusher.Flush()
+ return nil
+}
+
+const (
+ // sseRelayHeartbeatInterval paces `: ping` comment frames on an open
+ // tools/call relay stream. Intermediary proxies and load balancers
+ // commonly kill connections with no traffic for ~60s (e.g. the ALB
+ // default); an upstream tool that is busy but silent would otherwise
+ // have its downstream leg severed mid-call. Matches the cadence the
+ // GET keepalive endpoint (handleMCPSSE) already uses.
+ sseRelayHeartbeatInterval = 15 * time.Second
+ // sseRelayWriteTimeout bounds each individual write to the agent. The
+ // server deliberately sets no global WriteTimeout (it would kill every
+ // long-lived stream), so without a per-write deadline an agent that
+ // stops reading would block the handler goroutine in Write forever —
+ // and, because the relay is synchronous, wedge the upstream read loop
+ // with it. A deadline turns the stalled agent into a write error, which
+ // aborts the relay as stream_aborted_downstream.
+ sseRelayWriteTimeout = 30 * time.Second
+)
+
+// sseRelaySink turns one agent-facing tools/call response into a live SSE
+// response. StreamStarted commits the headers; after that, finishStream must
+// write the terminal response and stop the heartbeat before the handler exits.
+// The mutex prevents heartbeat and event frames from interleaving.
+type sseRelaySink struct {
+ w http.ResponseWriter
+ flusher http.Flusher
+ rc *http.ResponseController
+
+ // heartbeatInterval defaults to sseRelayHeartbeatInterval; a test seam.
+ heartbeatInterval time.Duration
+
+ mu sync.Mutex // serializes writes to w: events, heartbeats, terminal frame
+ // writeErr is the first write failure, sticky. A heartbeat that fails
+ // (agent gone) surfaces here so the next Event aborts the relay
+ // promptly instead of waiting for its own write to fail.
+ writeErr error
+
+ started bool
+ eventCount int
+ heartbeatStop chan struct{}
+ heartbeatDone chan struct{}
+ stopOnce sync.Once
+}
+
+func newSSERelaySink(w http.ResponseWriter, flusher http.Flusher) *sseRelaySink {
+ return &sseRelaySink{
+ w: w,
+ flusher: flusher,
+ rc: http.NewResponseController(w),
+ heartbeatInterval: sseRelayHeartbeatInterval,
+ }
+}
+
+func (s *sseRelaySink) StreamStarted() {
+ s.started = true
+ s.w.Header().Set("Content-Type", "text/event-stream")
+ s.w.Header().Set("Cache-Control", "no-cache")
+ s.w.Header().Set("Connection", "keep-alive")
+ s.w.Header().Set("X-Accel-Buffering", "no")
+ s.w.WriteHeader(http.StatusOK)
+ s.flusher.Flush()
+
+ s.heartbeatStop = make(chan struct{})
+ s.heartbeatDone = make(chan struct{})
+ go s.heartbeatLoop()
+}
+
+func (s *sseRelaySink) heartbeatLoop() {
+ defer close(s.heartbeatDone)
+ ticker := time.NewTicker(s.heartbeatInterval)
+ defer ticker.Stop()
+ for {
+ select {
+ case <-s.heartbeatStop:
+ return
+ case <-ticker.C:
+ s.mu.Lock()
+ if s.writeErr != nil {
+ s.mu.Unlock()
+ return
+ }
+ s.setWriteDeadlineLocked()
+ if _, err := fmt.Fprint(s.w, ": ping\n\n"); err != nil {
+ s.writeErr = err
+ s.mu.Unlock()
+ return
+ }
+ s.flusher.Flush()
+ s.mu.Unlock()
+ }
+ }
+}
+
+// setWriteDeadlineLocked arms the per-write deadline, best-effort: not every
+// ResponseWriter supports it (httptest recorders don't), and an unsupported
+// deadline must not break the relay — it just loses the stalled-agent bound.
+func (s *sseRelaySink) setWriteDeadlineLocked() {
+ _ = s.rc.SetWriteDeadline(time.Now().Add(sseRelayWriteTimeout))
+}
+
+func (s *sseRelaySink) Event(evt mcp.StreamEvent) error {
+ if evt.ServerRequest {
+ // Atryum does not broker server-initiated requests (sampling,
+ // elicitation, roots) — it doesn't advertise those capabilities in
+ // initialize, and the agent has no channel to answer a request
+ // arriving on what it expects to be a tools/call response stream.
+ // Audited by the service-layer auditingSink already (server_request
+ // flag on the invocation.stream_event row); never written to the
+ // agent.
+ return nil
+ }
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ if s.writeErr != nil {
+ // A heartbeat already found the connection dead: abort the relay
+ // now rather than waiting for this write to discover it again.
+ return s.writeErr
+ }
+ s.setWriteDeadlineLocked()
+ if err := writeSSEEvent(s.w, s.flusher, "", evt.Data); err != nil {
+ s.writeErr = err
+ return err
+ }
+ // Incremented only after a successful write: eventCount reports frames
+ // delivered to the agent, not attempts.
+ s.eventCount++
+ return nil
+}
+
+// stopHeartbeat stops the heartbeat goroutine (if one was started) without
+// writing anything, and waits for it to exit. finishStream calls it before
+// the terminal frame; handlers additionally defer it right after building
+// the sink so a panic between StreamStarted and finishStream cannot leave
+// the goroutine writing to a ResponseWriter whose handler has returned.
+// Idempotent and safe to call on a never-started sink.
+func (s *sseRelaySink) stopHeartbeat() {
+ if s.heartbeatStop == nil {
+ return
+ }
+ s.stopOnce.Do(func() { close(s.heartbeatStop) })
+ <-s.heartbeatDone
+}
+
+// finishStream stops the heartbeat and writes the terminal frame as the
+// stream's final write. Must be called on every handler path once started
+// is true — it is what guarantees the heartbeat goroutine cannot write to
+// (or race on) the ResponseWriter after the handler returns.
+func (s *sseRelaySink) finishStream(terminal []byte) error {
+ s.stopHeartbeat()
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ if s.writeErr != nil {
+ return s.writeErr
+ }
+ s.setWriteDeadlineLocked()
+ if err := writeSSEEvent(s.w, s.flusher, "", terminal); err != nil {
+ s.writeErr = err
+ return err
+ }
+ return nil
+}
diff --git a/internal/api/sse_relay_test.go b/internal/api/sse_relay_test.go
new file mode 100644
index 00000000..5aafe660
--- /dev/null
+++ b/internal/api/sse_relay_test.go
@@ -0,0 +1,626 @@
+package api
+
+import (
+ "bufio"
+ "context"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "github.com/validmind/atryum/internal/invocation"
+ "github.com/validmind/atryum/internal/mcp"
+ "github.com/validmind/atryum/internal/store"
+)
+
+func TestMCPToolsCallRelaysStreamedEventsAndRewritesTerminalID(t *testing.T) {
+ now := time.Now().UTC()
+ svc := &stubService{invoke: invocation.InvocationResponse{
+ InvocationID: "inv_123", ServerName: "demo", ToolName: "demo_tool",
+ Status: invocation.StatusSucceeded, SubmittedAt: now, CompletedAt: &now,
+ Result: json.RawMessage(`{"content":[{"type":"text","text":"done"}]}`),
+ }}
+ svc.invokeStreamingFn = func(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error) {
+ sink.StreamStarted()
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)}); err != nil {
+ t.Fatalf("sink.Event: %v", err)
+ }
+ return svc.invoke, nil
+ }
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":99,"method":"tools/call","params":{"name":"demo_tool","arguments":{}}}`))
+ req.Header.Set("Accept", "application/json, text/event-stream")
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ if w.Code != http.StatusOK {
+ t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String())
+ }
+ if ct := w.Header().Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
+ t.Fatalf("expected text/event-stream, got %q", ct)
+ }
+ body := w.Body.String()
+ if !strings.Contains(body, "notifications/progress") {
+ t.Fatalf("expected the relayed progress notification in the body, got %q", body)
+ }
+ // The terminal frame must carry the agent's own id (99), not the fixed
+ // upstream envelope id ("1") mcp.Client always sends on the wire.
+ if !strings.Contains(body, `"id":99`) {
+ t.Fatalf("expected terminal frame rewritten to the agent's id 99, got %q", body)
+ }
+ if !strings.Contains(body, `"text":"done"`) {
+ t.Fatalf("expected terminal result body, got %q", body)
+ }
+}
+
+// TestMCPToolsCallRelaysMultiLineEventDataAsMultipleDataLines is a
+// regression test: writeSSEEvent must emit one "data:" line per line of a
+// multi-line payload rather than embedding raw newlines inside a single
+// "data:" line, which breaks SSE framing (a continuation line with no
+// field prefix is dropped by any compliant parser). Upstream SSE data can
+// genuinely be multi-line — see mcp.TestListToolsDecodesMultilineSSEData
+// for the receiving side of this same scenario.
+func TestMCPToolsCallRelaysMultiLineEventDataAsMultipleDataLines(t *testing.T) {
+ now := time.Now().UTC()
+ svc := &stubService{invoke: invocation.InvocationResponse{
+ InvocationID: "inv_multiline", ServerName: "demo", ToolName: "demo_tool",
+ Status: invocation.StatusSucceeded, SubmittedAt: now, CompletedAt: &now,
+ Result: json.RawMessage(`{"content":[{"type":"text","text":"done"}]}`),
+ }}
+ multiLineNotification := []byte("{\"jsonrpc\":\"2.0\",\n\"method\":\"notifications/progress\",\n\"params\":{\"progress\":1}}")
+ svc.invokeStreamingFn = func(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error) {
+ sink.StreamStarted()
+ if err := sink.Event(mcp.StreamEvent{Data: multiLineNotification}); err != nil {
+ t.Fatalf("sink.Event: %v", err)
+ }
+ return svc.invoke, nil
+ }
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"demo_tool","arguments":{}}}`))
+ req.Header.Set("Accept", "application/json, text/event-stream")
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ reader := bufio.NewReader(w.Body)
+ frame := readNextSSEFrame(t, reader)
+ if frame.data != string(multiLineNotification) {
+ t.Fatalf("expected the multi-line payload reconstructed exactly from multiple \"data:\" lines, got %q", frame.data)
+ }
+}
+
+// TestMCPToolsCallErrorAfterStreamStartedWritesTerminalErrorFrame is a
+// regression test: InvokeStreaming CAN return an error after the stream
+// has started (e.g. persisting the result fails after the relay
+// completed). A started stream must then end with a terminal JSON-RPC
+// error frame carrying the agent's request id — not writeRPCError (a
+// second WriteHeader) and not a bare close that would leave the request
+// unanswered.
+// TestSSERelaySinkHeartbeatsKeepStreamAliveWithoutCorruptingFrames covers
+// the intermediary-keepalive requirement: proxies/LBs kill connections
+// that carry no traffic (commonly ~60s idle), so an open relay must emit
+// `: ping` comments while the upstream is silent — and, because the
+// heartbeat runs on its own goroutine, its writes must never interleave
+// with (corrupt) an event or terminal frame. Run under -race this also
+// proves the mutex discipline.
+func TestSSERelaySinkHeartbeatsKeepStreamAliveWithoutCorruptingFrames(t *testing.T) {
+ w := httptest.NewRecorder()
+ sink := newSSERelaySink(w, w)
+ sink.heartbeatInterval = 2 * time.Millisecond
+
+ sink.StreamStarted()
+ deadline := time.Now().Add(60 * time.Millisecond)
+ for time.Now().Before(deadline) {
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)}); err != nil {
+ t.Fatalf("Event: %v", err)
+ }
+ }
+ if err := sink.finishStream([]byte(`{"jsonrpc":"2.0","id":1,"result":{}}`)); err != nil {
+ t.Fatalf("finishStream: %v", err)
+ }
+
+ body := w.Body.String()
+ if !strings.Contains(body, ": ping\n\n") {
+ t.Fatalf("expected heartbeat comments in the stream, got none in %d bytes", len(body))
+ }
+ // Frame integrity: every line must be a well-formed SSE line — a
+ // heartbeat interleaved mid-frame would produce a line that is neither.
+ for _, line := range strings.Split(body, "\n") {
+ if line == "" || strings.HasPrefix(line, "data: ") || strings.HasPrefix(line, ": ping") {
+ continue
+ }
+ t.Fatalf("malformed SSE line (heartbeat interleaved mid-frame?): %q", line)
+ }
+ if !strings.HasSuffix(strings.TrimRight(body, "\n"), `{"jsonrpc":"2.0","id":1,"result":{}}`) {
+ t.Fatalf("expected the terminal frame to be the stream's final write, got tail %q", body[max(0, len(body)-120):])
+ }
+}
+
+// switchableFailingWriter is a ResponseWriter whose writes succeed until
+// the test flips broken — simulating an agent whose connection died
+// mid-stream. broken is atomic because the heartbeat goroutine writes
+// concurrently with the test goroutine.
+type switchableFailingWriter struct {
+ *httptest.ResponseRecorder
+ broken atomic.Bool
+}
+
+func (f *switchableFailingWriter) Write(p []byte) (int, error) {
+ if f.broken.Load() {
+ return 0, fmt.Errorf("connection reset by peer")
+ }
+ return f.ResponseRecorder.Write(p)
+}
+
+// TestSSERelaySinkHeartbeatFailureSurfacesOnNextEvent covers stalled-agent
+// detection: when a heartbeat write discovers the connection is dead, the
+// failure must stick and abort the relay on the next Event — the
+// synchronous relay loop is otherwise blind to the downstream connection
+// between events.
+//
+// The dead connection is discovered on a goroutine the test never calls
+// directly, and the test's own polling loop only ever observes the result
+// through the mutex-guarded sink.writeErr:
+//
+// test (main goroutine) heartbeatLoop goroutine
+// ---------------------- -----------------------
+// StreamStarted() --------------> go heartbeatLoop()
+// Event(progress=1) -> ok
+// fw.broken.Store(true)
+// (connection now dead)
+// poll sink.writeErr
+// (mu.Lock/Unlock in a loop) ticker fires (2ms)
+// mu.Lock(); writeErr == nil
+// Write(": ping\n\n") -> fails
+// writeErr = err; mu.Unlock()
+// sees writeErr != nil, stops polling
+// Event(progress=2)
+// -> writeErr already set; returns
+// it without attempting a write
+// finishStream([]byte("{}"))
+// -> same sticky writeErr; returns
+// it without writing a terminal frame
+//
+// Event and finishStream never touch the dead connection themselves — both
+// only ever observe the heartbeat goroutine's writeErr.
+func TestSSERelaySinkHeartbeatFailureSurfacesOnNextEvent(t *testing.T) {
+ fw := &switchableFailingWriter{ResponseRecorder: httptest.NewRecorder()}
+ sink := newSSERelaySink(fw, fw.ResponseRecorder)
+ sink.heartbeatInterval = 2 * time.Millisecond
+
+ sink.StreamStarted()
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)}); err != nil {
+ t.Fatalf("first Event should succeed, got %v", err)
+ }
+ fw.broken.Store(true) // the agent's connection dies between events
+
+ // Wait until a heartbeat has hit the dead connection.
+ deadlineExceeded := time.Now().Add(2 * time.Second)
+ for {
+ sink.mu.Lock()
+ failed := sink.writeErr != nil
+ sink.mu.Unlock()
+ if failed {
+ break
+ }
+ if time.Now().After(deadlineExceeded) {
+ t.Fatal("heartbeat never observed the write failure")
+ }
+ time.Sleep(2 * time.Millisecond)
+ }
+
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":2}}`)}); err == nil {
+ t.Fatal("expected the heartbeat's sticky write error to abort the next Event")
+ }
+ if err := sink.finishStream([]byte(`{}`)); err == nil {
+ t.Fatal("expected finishStream to report the dead connection rather than pretend the terminal frame was written")
+ }
+}
+
+func TestMCPToolsCallAuditsTerminalDeliveryFailure(t *testing.T) {
+ now := time.Now().UTC()
+ svc := &stubService{invoke: invocation.InvocationResponse{
+ InvocationID: "inv_delivery", ServerName: "demo", ToolName: "demo_tool",
+ Status: invocation.StatusSucceeded, SubmittedAt: now, CompletedAt: &now,
+ Result: json.RawMessage(`{"content":[{"type":"text","text":"done"}]}`),
+ }}
+ writer := &switchableFailingWriter{ResponseRecorder: httptest.NewRecorder()}
+ svc.invokeStreamingFn = func(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error) {
+ sink.StreamStarted()
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)}); err != nil {
+ t.Fatalf("sink.Event: %v", err)
+ }
+ writer.broken.Store(true)
+ return svc.invoke, nil
+ }
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"demo_tool","arguments":{}}}`))
+ req.Header.Set("Accept", "text/event-stream")
+
+ h.Routes().ServeHTTP(writer, req)
+
+ if svc.streamDeliveryInvocationID != "inv_delivery" {
+ t.Fatalf("delivery audit invocation = %q, want inv_delivery", svc.streamDeliveryInvocationID)
+ }
+ if svc.streamDeliveryStatus != "failed" {
+ t.Fatalf("delivery audit status = %q, want failed", svc.streamDeliveryStatus)
+ }
+ if !strings.Contains(svc.streamDeliveryMessage, "connection reset") {
+ t.Fatalf("delivery audit message = %q, want terminal write failure", svc.streamDeliveryMessage)
+ }
+}
+
+func TestMCPToolsCallErrorAfterStreamStartedWritesTerminalErrorFrame(t *testing.T) {
+ svc := &stubService{}
+ svc.invokeStreamingFn = func(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error) {
+ sink.StreamStarted()
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)}); err != nil {
+ t.Fatalf("sink.Event: %v", err)
+ }
+ return invocation.InvocationResponse{}, fmt.Errorf("persisting result failed")
+ }
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":77,"method":"tools/call","params":{"name":"demo_tool","arguments":{}}}`))
+ req.Header.Set("Accept", "application/json, text/event-stream")
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ if ct := w.Header().Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
+ t.Fatalf("expected text/event-stream (headers were already sent), got %q", ct)
+ }
+ body := w.Body.String()
+ if !strings.Contains(body, `"id":77`) {
+ t.Fatalf("expected the terminal error frame rewritten to the agent's id 77, got %q", body)
+ }
+ if !strings.Contains(body, `"error"`) || !strings.Contains(body, "failed to finalize invocation") {
+ t.Fatalf("expected a terminal JSON-RPC error frame with the static message, got %q", body)
+ }
+ // The internal error text must never reach the wire: it can carry
+ // SQL/driver detail. It is logged server-side instead.
+ if strings.Contains(body, "persisting result failed") {
+ t.Fatalf("internal error detail leaked into the terminal frame: %q", body)
+ }
+}
+
+// TestMCPToolsCallNeverForwardsServerToClientRequestToAgent is a
+// regression test: Atryum does not broker server-initiated requests
+// (sampling, elicitation, roots) — the agent has no channel to answer one
+// arriving on a tools/call response stream, so it must never be forwarded,
+// even though it is still relayed to the sink's Event method for the
+// service layer to audit.
+func TestMCPToolsCallNeverForwardsServerToClientRequestToAgent(t *testing.T) {
+ now := time.Now().UTC()
+ svc := &stubService{invoke: invocation.InvocationResponse{
+ InvocationID: "inv_srvreq", ServerName: "demo", ToolName: "demo_tool",
+ Status: invocation.StatusSucceeded, SubmittedAt: now, CompletedAt: &now,
+ Result: json.RawMessage(`{"content":[{"type":"text","text":"done"}]}`),
+ }}
+ svc.invokeStreamingFn = func(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error) {
+ sink.StreamStarted()
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","id":"srv-1","method":"sampling/createMessage","params":{}}`), ServerRequest: true}); err != nil {
+ t.Fatalf("sink.Event(server request): %v", err)
+ }
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)}); err != nil {
+ t.Fatalf("sink.Event(notification): %v", err)
+ }
+ return svc.invoke, nil
+ }
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"demo_tool","arguments":{}}}`))
+ req.Header.Set("Accept", "application/json, text/event-stream")
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ body := w.Body.String()
+ if strings.Contains(body, "sampling/createMessage") {
+ t.Fatalf("expected the server-to-client request never to reach the agent, got %q", body)
+ }
+ if !strings.Contains(body, "notifications/progress") {
+ t.Fatalf("expected the notification to still be relayed, got %q", body)
+ }
+}
+
+func TestMCPToolsCallStreamedFailureWritesTerminalErrorFrame(t *testing.T) {
+ now := time.Now().UTC()
+ svc := &stubService{invoke: invocation.InvocationResponse{
+ InvocationID: "inv_failed", ServerName: "demo", ToolName: "demo_tool",
+ Status: invocation.StatusFailed, SubmittedAt: now, CompletedAt: &now,
+ Error: json.RawMessage(`{"content":[{"type":"text","text":"upstream exploded"}],"isError":true}`),
+ }}
+ svc.invokeStreamingFn = func(ctx context.Context, req invocation.CreateInvocationRequest, sink mcp.StreamSink) (invocation.InvocationResponse, error) {
+ sink.StreamStarted()
+ if err := sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)}); err != nil {
+ t.Fatalf("sink.Event: %v", err)
+ }
+ return svc.invoke, nil
+ }
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":55,"method":"tools/call","params":{"name":"demo_tool","arguments":{}}}`))
+ req.Header.Set("Accept", "application/json, text/event-stream")
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ if ct := w.Header().Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
+ t.Fatalf("expected text/event-stream, got %q", ct)
+ }
+ body := w.Body.String()
+ if !strings.Contains(body, "upstream exploded") {
+ t.Fatalf("expected the terminal error content, got %q", body)
+ }
+ if !strings.Contains(body, `"id":55`) {
+ t.Fatalf("expected terminal frame rewritten to the agent's id 55, got %q", body)
+ }
+}
+
+func TestMCPToolsCallWithoutStreamAcceptGetsPlainJSONEvenIfUpstreamWouldStream(t *testing.T) {
+ now := time.Now().UTC()
+ svc := &stubService{invoke: invocation.InvocationResponse{
+ InvocationID: "inv_1", ServerName: "demo", ToolName: "demo_tool",
+ Status: invocation.StatusSucceeded, SubmittedAt: now, CompletedAt: &now,
+ Result: json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`),
+ }}
+ svc.invokeStreamingFn = func(context.Context, invocation.CreateInvocationRequest, mcp.StreamSink) (invocation.InvocationResponse, error) {
+ t.Fatal("InvokeStreaming should never be called when the agent did not send Accept: text/event-stream")
+ return invocation.InvocationResponse{}, nil
+ }
+ h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"demo_tool","arguments":{}}}`))
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ if ct := w.Header().Get("Content-Type"); !strings.Contains(ct, "application/json") {
+ t.Fatalf("expected application/json, got %q", ct)
+ }
+ if !strings.Contains(w.Body.String(), `"text":"ok"`) {
+ t.Fatalf("expected buffered JSON result, got %s", w.Body.String())
+ }
+}
+
+func TestMCPToolsCallDenialWithStreamCapableAgentStaysJSON(t *testing.T) {
+ now := time.Now().UTC()
+ rules := &stubRulesRepo{rules: []store.Rule{
+ {ID: "bash-deny", Action: invocation.RuleActionAutoDeny, ServerPatterns: []string{"demo"}, ToolPatterns: []string{"Bash"}, Enabled: true, Order: 0},
+ }}
+ svc := &stubService{invoke: invocation.InvocationResponse{
+ InvocationID: "inv_denied", ServerName: "demo", ToolName: "Bash",
+ Status: invocation.StatusDenied,
+ SubmittedAt: now, CompletedAt: &now,
+ Error: json.RawMessage(`{"content":[{"type":"text","text":"Tool call denied by approval rule (auto_deny)."}],"isError":true}`),
+ }}
+ svc.invokeStreamingFn = func(context.Context, invocation.CreateInvocationRequest, mcp.StreamSink) (invocation.InvocationResponse, error) {
+ // A denial is decided before any upstream execution — the sink must
+ // never be touched, so this returns the same stubbed response
+ // Invoke would, without ever calling StreamStarted/Event.
+ return svc.invoke, nil
+ }
+ h := NewHandler(svc, stubServerService{}, nil, rules, nil, nil, nil, nil, nil, nil)
+ req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"Bash","arguments":{"cmd":"ls"}}}`))
+ req.Header.Set("Accept", "application/json, text/event-stream")
+ w := httptest.NewRecorder()
+
+ h.Routes().ServeHTTP(w, req)
+
+ if ct := w.Header().Get("Content-Type"); !strings.Contains(ct, "application/json") {
+ t.Fatalf("expected application/json (stream never starts on denial), got %q", ct)
+ }
+ var rpcResp struct {
+ Result struct {
+ Content []struct {
+ Type string `json:"type"`
+ Text string `json:"text"`
+ } `json:"content"`
+ IsError bool `json:"isError"`
+ } `json:"result"`
+ }
+ if err := json.Unmarshal(w.Body.Bytes(), &rpcResp); err != nil {
+ t.Fatal(err)
+ }
+ if !rpcResp.Result.IsError {
+ t.Fatalf("expected isError=true, got %#v", rpcResp.Result)
+ }
+ if len(rpcResp.Result.Content) < 2 {
+ t.Fatalf("expected denial text plus rules context, got %#v", rpcResp.Result.Content)
+ }
+}
+
+// sseEventFrame is one decoded "event:"/"data:" SSE frame read off a live
+// HTTP response body.
+type sseEventFrame struct {
+ event string
+ data string
+}
+
+// readNextSSEFrame blocks on reader until one full SSE frame (up to the
+// blank line that terminates it) has arrived, then returns it. Used to
+// prove live, incremental delivery: unlike parsing a fully-buffered body,
+// this only returns once that specific frame has actually been read off
+// the wire.
+func readNextSSEFrame(t *testing.T, reader *bufio.Reader) sseEventFrame {
+ t.Helper()
+ var frame sseEventFrame
+ for {
+ line, err := reader.ReadString('\n')
+ if err != nil {
+ t.Fatalf("read SSE frame: %v (partial line=%q)", err, line)
+ }
+ line = strings.TrimRight(line, "\n")
+ if line == "" {
+ return frame
+ }
+ if strings.HasPrefix(line, ":") {
+ continue
+ }
+ field, value, ok := strings.Cut(line, ": ")
+ if !ok {
+ field, value = line, ""
+ }
+ switch field {
+ case "event":
+ frame.event = value
+ case "data":
+ if frame.data != "" {
+ frame.data += "\n"
+ }
+ frame.data += value
+ }
+ }
+}
+
+// TestMCPToolsCallEndToEndRelaysLiveBeforeTerminalResponseExists is the
+// primary acceptance test for the streaming relay: a real Handler, real
+// invocation.Service, and real mcp.Client are wired to a fake upstream MCP
+// server that flushes one progress notification and then blocks — on a
+// channel this test controls — before it is even able to write its
+// terminal response. The agent (a real HTTP client reading the response
+// incrementally) must observe the notification before that channel is
+// released, which proves the intermediate event was relayed live rather
+// than after the fact from a buffered body: the terminal response cannot
+// exist yet at the point the test asserts the notification arrived.
+//
+// test (as the agent) agentServer (Handler+mcp.Client) fake upstream (httptest)
+// -------------------- --------------------------------- ------------------------
+// POST /mcp/demo tools/call
+// progressToken=tok-live -------------------------------------------> proxies tools/call upstream
+// writes progress frame (tok-live)
+// flusher.Flush()
+// relays the frame live <--------
+// first := readNextSSEFrame() <---
+// (blocks until this exact
+// frame is read off the wire)
+// assert first has tok-live
+// <-releaseTerminal
+// (blocked: the terminal response
+// cannot exist yet)
+// close(releaseTerminal) --------------------------------------------> unblocks
+// writes terminal frame (id "1")
+// rewrites id "1" -> 42 <-------
+// terminal := readNextSSEFrame() <-
+// (id:42, "all done")
+//
+// Because first is read before releaseTerminal is closed, the assertion on
+// first.data cannot be satisfied by a buffered replay: the terminal frame
+// literally does not exist in the upstream handler yet at that point.
+func TestMCPToolsCallEndToEndRelaysLiveBeforeTerminalResponseExists(t *testing.T) {
+ releaseTerminal := make(chan struct{})
+ upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ // The standalone SSE stream Atryum opens alongside a
+ // progressToken-bearing tools/call. This fake upstream doesn't
+ // support it (a legitimate, spec-allowed response); the agent's
+ // progress notification arrives via the tools/call POST
+ // response itself below, exercised independently of this.
+ http.Error(w, "not found", http.StatusNotFound)
+ return
+ }
+ var body map[string]any
+ if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+ t.Errorf("decode upstream request: %v", err)
+ return
+ }
+ switch body["method"] {
+ case "initialize":
+ _ = json.NewEncoder(w).Encode(map[string]any{
+ "jsonrpc": "2.0", "id": body["id"],
+ "result": map[string]any{"serverInfo": map[string]any{"name": "fake", "version": "0.1.0"}, "capabilities": map[string]any{}},
+ })
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+
+ // Echo the agent's own progressToken back — this only works if
+ // _meta forwarded all the way from the agent's request through
+ // to the upstream tools/call envelope (Phase 0).
+ params, _ := body["params"].(map[string]any)
+ meta, _ := params["_meta"].(map[string]any)
+ token, _ := meta["progressToken"].(string)
+ progress, _ := json.Marshal(map[string]any{
+ "jsonrpc": "2.0", "method": "notifications/progress",
+ "params": map[string]any{"progressToken": token, "progress": 1},
+ })
+ fmt.Fprintf(w, "event: message\ndata: %s\n\n", progress)
+ flusher.Flush()
+
+ <-releaseTerminal // the terminal response cannot exist until the test releases this
+
+ result, _ := json.Marshal(map[string]any{
+ "jsonrpc": "2.0", "id": "1",
+ "result": map[string]any{"content": []map[string]any{{"type": "text", "text": "all done"}}},
+ })
+ fmt.Fprintf(w, "event: message\ndata: %s\n\n", result)
+ flusher.Flush()
+ default:
+ t.Errorf("unexpected upstream method %v", body["method"])
+ }
+ }))
+ defer upstream.Close()
+
+ agentServer, _ := newTestAgentServer(t, "demo", upstream.URL, 5, false)
+
+ reqBody := `{"jsonrpc":"2.0","id":42,"method":"tools/call","params":{"name":"demo_tool","arguments":{},"_meta":{"progressToken":"tok-live"}}}`
+ req, err := http.NewRequest(http.MethodPost, agentServer.URL+"/mcp/demo", strings.NewReader(reqBody))
+ if err != nil {
+ t.Fatal(err)
+ }
+ req.Header.Set("Content-Type", "application/json")
+ req.Header.Set("Accept", "application/json, text/event-stream")
+
+ resp, err := http.DefaultClient.Do(req)
+ if err != nil {
+ t.Fatalf("agent request: %v", err)
+ }
+ defer resp.Body.Close()
+
+ if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
+ t.Fatalf("expected text/event-stream, got %q", ct)
+ }
+
+ reader := bufio.NewReader(resp.Body)
+ first := readNextSSEFrame(t, reader)
+ if !strings.Contains(first.data, "notifications/progress") || !strings.Contains(first.data, "tok-live") {
+ t.Fatalf("expected the progress notification (echoing the agent's progressToken) first, got %q", first.data)
+ }
+
+ // Only now — after the agent has actually received the intermediate
+ // event over the wire — does the fake upstream get to write its
+ // terminal response.
+ close(releaseTerminal)
+
+ terminal := readNextSSEFrame(t, reader)
+ if !strings.Contains(terminal.data, `"id":42`) {
+ t.Fatalf("expected the terminal frame rewritten to the agent's id 42, got %q", terminal.data)
+ }
+ if !strings.Contains(terminal.data, "all done") {
+ t.Fatalf("expected the terminal result body, got %q", terminal.data)
+ }
+}
+
+// TestWriteSSEEventNormalizesCarriageReturns pins that no raw CR ever
+// reaches the wire inside a data field: CR is an SSE line terminator, so an
+// upstream that embeds one (legal JSON inter-token whitespace) would
+// otherwise have its frame split mid-message by a compliant agent parser.
+func TestWriteSSEEventNormalizesCarriageReturns(t *testing.T) {
+ rec := httptest.NewRecorder()
+ payload := []byte("{\"a\":\r\r1,\r\n\"b\":2}")
+ if err := writeSSEEvent(rec, rec, "", payload); err != nil {
+ t.Fatalf("writeSSEEvent returned error: %v", err)
+ }
+ body := rec.Body.String()
+ if strings.Contains(body, "\r") {
+ t.Fatalf("raw CR reached the wire: %q", body)
+ }
+ want := "data: {\"a\":\ndata: \ndata: 1,\ndata: \"b\":2}\n\n"
+ if body != want {
+ t.Fatalf("framed body = %q, want %q", body, want)
+ }
+}
diff --git a/internal/api/testdata/mcp_standalone_fixture/server.py b/internal/api/testdata/mcp_standalone_fixture/server.py
new file mode 100644
index 00000000..9cf2b9f0
--- /dev/null
+++ b/internal/api/testdata/mcp_standalone_fixture/server.py
@@ -0,0 +1,37 @@
+"""Real MCP server (official Python SDK, FastMCP, Streamable HTTP transport)
+used by mcp_standalone_stream_test.go to exercise Atryum's standalone-SSE-
+stream relay path against genuine SDK behavior, not a hand-rolled fixture.
+
+FastMCP's Context.report_progress() calls send_progress_notification()
+without related_request_id, so the server's message router sends every
+progress update to the standalone GET stream — never to the tools/call POST
+response body. That's exactly the case internal/mcp/client.go's
+standaloneStream machinery exists for (see docs/architecture.md's "Live SSE
+relay for tools/call" section). This fixture is not a workaround for that
+behavior; it demonstrates it, because it's the real SDK's actual behavior.
+
+Reads PORT from the environment (default 8642) so the Go test harness can
+pick a free port per run.
+"""
+
+import asyncio
+import os
+
+from mcp.server.fastmcp import Context, FastMCP
+
+PORT = int(os.environ.get("PORT", "8642"))
+
+mcp = FastMCP("atryum-standalone-fixture", host="127.0.0.1", port=PORT)
+
+
+@mcp.tool()
+async def slow_streaming_task(ctx: Context, steps: int = 3, delay_seconds: float = 1.0) -> str:
+ """Report progress steps times with real delays, then return a result."""
+ for i in range(1, steps + 1):
+ await ctx.report_progress(progress=i, total=steps, message=f"step {i}/{steps}")
+ await asyncio.sleep(delay_seconds)
+ return f"done after {steps} real progress notifications"
+
+
+if __name__ == "__main__":
+ mcp.run(transport="streamable-http")
diff --git a/internal/config/config.go b/internal/config/config.go
index fccb6cf2..54f2199b 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -148,6 +148,50 @@ type ServerConfig struct {
type DefaultsConfig struct {
RequestTimeoutSeconds int `toml:"request_timeout_seconds"`
+
+ // StreamRelayEnabled is the kill-switch for the tools/call SSE relay
+ // (see api.Handler.SetStreamRelayEnabled / docs/architecture.md). The
+ // downstream relay only activates when the agent sends Accept:
+ // text/event-stream and the upstream client enters live mode. HTTP
+ // upstreams do that by returning SSE; stdio upstreams do it when they
+ // emit an intermediate message.
+ StreamRelayEnabled bool `toml:"stream_relay_enabled"`
+ // StreamHeaderTimeoutSeconds bounds setup before tool response reading:
+ // HTTP session initialization and response headers, or the stdio
+ // initialize handshake. Zero falls back to RequestTimeoutSeconds.
+ StreamHeaderTimeoutSeconds int `toml:"stream_header_timeout_seconds"`
+ // StreamIdleTimeoutSeconds bounds response-reading inactivity. Streaming
+ // events reset it; for a plain HTTP JSON response it bounds the complete
+ // body read. It does not limit the call's total duration.
+ StreamIdleTimeoutSeconds int `toml:"stream_idle_timeout_seconds"`
+ // StreamMaxDurationSeconds bounds response reading after setup completes.
+ // Zero disables the bound (unlimited).
+ StreamMaxDurationSeconds int `toml:"stream_max_duration_seconds"`
+ // StreamMaxMessageBytes bounds one upstream JSON-RPC message or plain
+ // response body. Zero at the mcp package boundary uses its 4 MiB default.
+ StreamMaxMessageBytes int `toml:"stream_max_message_bytes"`
+ // StreamAuditMaxEvents caps how many invocation.stream_event audit
+ // rows get persisted per call; beyond the cap, events are still
+ // relayed live to the agent but only counted, not stored
+ // individually. Zero disables this count cap; the bounded audit queue
+ // may still drop events under storage backpressure, which is reported
+ // in the stream completion audit row.
+ StreamAuditMaxEvents int `toml:"stream_audit_max_events"`
+ // StreamAuditMaxEventBytes truncates each persisted stream_event
+ // row's data field beyond this size. Zero disables truncation.
+ StreamAuditMaxEventBytes int `toml:"stream_audit_max_event_bytes"`
+}
+
+// EffectiveStreamHeaderTimeoutSeconds resolves the documented zero-value
+// fallback for stream_header_timeout_seconds: zero (or negative) falls back
+// to request_timeout_seconds. Lives here so every consumer of DefaultsConfig
+// resolves the fallback identically instead of re-implementing it at each
+// wiring site.
+func (d DefaultsConfig) EffectiveStreamHeaderTimeoutSeconds() int {
+ if d.StreamHeaderTimeoutSeconds > 0 {
+ return d.StreamHeaderTimeoutSeconds
+ }
+ return d.RequestTimeoutSeconds
}
type UpstreamConfig struct {
@@ -183,7 +227,13 @@ func Load(path string) (Config, error) {
ConnectionTimeoutSecs: 5,
},
Defaults: DefaultsConfig{
- RequestTimeoutSeconds: 30,
+ RequestTimeoutSeconds: 30,
+ StreamRelayEnabled: true,
+ StreamIdleTimeoutSeconds: 60,
+ StreamMaxDurationSeconds: 600,
+ StreamMaxMessageBytes: 4 * 1024 * 1024,
+ StreamAuditMaxEvents: 100,
+ StreamAuditMaxEventBytes: 4096,
},
}
_, err := toml.DecodeFile(path, &cfg)
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
index 28321fb2..852ee92e 100644
--- a/internal/config/config_test.go
+++ b/internal/config/config_test.go
@@ -72,6 +72,54 @@ func TestLoadMissingConfigUsesDefaultsAndEnv(t *testing.T) {
if cfg.Backend.APISecret != "env-api-secret" {
t.Fatalf("Backend.APISecret = %q", cfg.Backend.APISecret)
}
+ if !cfg.Defaults.StreamRelayEnabled {
+ t.Fatal("Defaults.StreamRelayEnabled = false, want true (safe default: doubly gated on agent Accept + upstream content-type)")
+ }
+ if cfg.Defaults.StreamIdleTimeoutSeconds != 60 {
+ t.Fatalf("Defaults.StreamIdleTimeoutSeconds = %d, want 60", cfg.Defaults.StreamIdleTimeoutSeconds)
+ }
+ if cfg.Defaults.StreamMaxDurationSeconds != 600 {
+ t.Fatalf("Defaults.StreamMaxDurationSeconds = %d, want 600", cfg.Defaults.StreamMaxDurationSeconds)
+ }
+ if cfg.Defaults.StreamAuditMaxEvents != 100 {
+ t.Fatalf("Defaults.StreamAuditMaxEvents = %d, want 100", cfg.Defaults.StreamAuditMaxEvents)
+ }
+ if cfg.Defaults.StreamAuditMaxEventBytes != 4096 {
+ t.Fatalf("Defaults.StreamAuditMaxEventBytes = %d, want 4096", cfg.Defaults.StreamAuditMaxEventBytes)
+ }
+ if cfg.Defaults.StreamMaxMessageBytes != 4*1024*1024 {
+ t.Fatalf("Defaults.StreamMaxMessageBytes = %d, want 4194304", cfg.Defaults.StreamMaxMessageBytes)
+ }
+ if cfg.Defaults.StreamHeaderTimeoutSeconds != 0 {
+ t.Fatalf("Defaults.StreamHeaderTimeoutSeconds = %d, want 0 (falls back to RequestTimeoutSeconds at the call site)", cfg.Defaults.StreamHeaderTimeoutSeconds)
+ }
+}
+
+func TestLoadStreamRelayCanBeDisabledViaTOML(t *testing.T) {
+ dir := t.TempDir()
+ path := filepath.Join(dir, "atryum.toml")
+ if err := os.WriteFile(path, []byte("[defaults]\nstream_relay_enabled = false\nstream_idle_timeout_seconds = 30\n"), 0o644); err != nil {
+ t.Fatal(err)
+ }
+
+ cfg, err := Load(path)
+ if err != nil {
+ t.Fatalf("Load: %v", err)
+ }
+ if cfg.Defaults.StreamRelayEnabled {
+ t.Fatal("Defaults.StreamRelayEnabled = true, want false (explicitly disabled in TOML)")
+ }
+ if cfg.Defaults.StreamIdleTimeoutSeconds != 30 {
+ t.Fatalf("Defaults.StreamIdleTimeoutSeconds = %d, want 30 (explicit TOML override)", cfg.Defaults.StreamIdleTimeoutSeconds)
+ }
+ // Fields the TOML fragment didn't mention keep their Go-level default,
+ // proving partial overrides don't blow away the rest of [defaults].
+ if cfg.Defaults.StreamMaxDurationSeconds != 600 {
+ t.Fatalf("Defaults.StreamMaxDurationSeconds = %d, want 600 (untouched default)", cfg.Defaults.StreamMaxDurationSeconds)
+ }
+ if cfg.Defaults.RequestTimeoutSeconds != 30 {
+ t.Fatalf("Defaults.RequestTimeoutSeconds = %d, want 30 (untouched default)", cfg.Defaults.RequestTimeoutSeconds)
+ }
}
func TestLoadAuthAdminClaimValueAcceptsBool(t *testing.T) {
@@ -122,6 +170,21 @@ recent_chat_messages_limit = 42
}
}
+func TestEffectiveStreamHeaderTimeoutSecondsFallsBackToRequestTimeout(t *testing.T) {
+ d := DefaultsConfig{RequestTimeoutSeconds: 30}
+ if got := d.EffectiveStreamHeaderTimeoutSeconds(); got != 30 {
+ t.Fatalf("zero stream header timeout resolved to %d, want the request timeout 30", got)
+ }
+ d.StreamHeaderTimeoutSeconds = 12
+ if got := d.EffectiveStreamHeaderTimeoutSeconds(); got != 12 {
+ t.Fatalf("explicit stream header timeout resolved to %d, want 12", got)
+ }
+ d.StreamHeaderTimeoutSeconds = -1
+ if got := d.EffectiveStreamHeaderTimeoutSeconds(); got != 30 {
+ t.Fatalf("negative stream header timeout resolved to %d, want the request timeout 30", got)
+ }
+}
+
func TestLoadPlansTTLBounds(t *testing.T) {
dir := t.TempDir()
diff --git a/internal/invocation/model.go b/internal/invocation/model.go
index 63292678..f4f35182 100644
--- a/internal/invocation/model.go
+++ b/internal/invocation/model.go
@@ -108,6 +108,11 @@ type CreateInvocationRequest struct {
// the per-row fallback for anonymous (no agent_id) invocations.
ClientName string `json:"-"`
ClientVersion string `json:"-"`
+ // Meta carries the agent's MCP tools/call params._meta (e.g.
+ // progressToken) through to the upstream call so progress/logging
+ // notifications the upstream emits can be correlated back to the
+ // agent's own request.
+ Meta map[string]any `json:"-"`
}
// ExternalSubmitRequest is used by callers that execute the tool themselves
diff --git a/internal/invocation/service.go b/internal/invocation/service.go
index ff3821d8..597a0a58 100644
--- a/internal/invocation/service.go
+++ b/internal/invocation/service.go
@@ -178,13 +178,31 @@ type resolver interface {
}
type upstreamClient interface {
- Invoke(ctx context.Context, upstream mcp.Upstream, tool string, input map[string]any, requestID *string) (mcp.InvokeResult, error)
+ Invoke(ctx context.Context, upstream mcp.Upstream, tool string, input map[string]any, requestID *string, meta map[string]any) (mcp.InvokeResult, error)
+ InvokeStream(ctx context.Context, upstream mcp.Upstream, tool string, input map[string]any, requestID *string, meta map[string]any, sink mcp.StreamSink, opts mcp.StreamOptions) (mcp.InvokeResult, error)
ListTools(ctx context.Context, upstream mcp.Upstream) ([]mcp.Tool, error)
ForwardEnvelope(ctx context.Context, upstream mcp.Upstream, envelope mcp.Envelope, protocolVersion string) (mcp.ForwardResult, error)
}
+// StreamAuditLimits bounds how many per-call invocation.stream_event audit
+// rows get persisted, and how large each one's data field is, for one
+// streaming tools/call execution. A zero value disables that particular
+// cap (no configured event-count limit / untruncated data). The shared audit
+// dispatcher's bounded queues remain a final backpressure guard even with
+// MaxEvents zero and report drops in invocation.stream_completed.
+type StreamAuditLimits struct {
+ MaxEvents int
+ MaxEventBytes int
+}
+
const toolCatalogTTL = 5 * time.Minute
+// terminalPersistenceTimeout bounds writes that finalize an invocation after
+// execution. Those writes must outlive the agent-facing request context: a
+// downstream disconnect cancels that context before the service can record the
+// resulting failure, which would otherwise leave the durable row executing.
+const terminalPersistenceTimeout = 5 * time.Second
+
type toolCatalogEntry struct {
tools map[string]mcp.Tool
fetchedAt time.Time
@@ -212,6 +230,15 @@ type Service struct {
mu sync.Mutex
pendingApprovals map[string]chan approvalDecision
+ // streamOptions and streamAuditLimits govern InvokeStreaming's execution
+ // once a sink is present (see finishExecutionStreaming). SetStreamOptions
+ // installs configured values; until then effectiveStreamOptions
+ // substitutes safe derived bounds so the zero value never means
+ // "unbounded relay".
+ streamOptions mcp.StreamOptions
+ streamOptionsSet bool
+ streamAuditLimits StreamAuditLimits
+
toolCatalogMu sync.Mutex
toolCatalog map[string]toolCatalogEntry
}
@@ -250,6 +277,81 @@ func (s *Service) SetInvocationSummarizer(client SummaryClient) {
s.summarizer = client
}
+// SetStreamOptions configures the header/idle/max-duration timeout scheme
+// and per-event audit caps used by InvokeStreaming once a caller supplies a
+// sink (see finishExecutionStreaming). Calls with a nil sink are unaffected.
+func (s *Service) SetStreamOptions(opts mcp.StreamOptions, auditLimits StreamAuditLimits) {
+ s.streamOptions = opts
+ s.streamOptionsSet = true
+ s.streamAuditLimits = auditLimits
+}
+
+// Fallback stream bounds for a Service whose owner never called
+// SetStreamOptions. They mirror the documented config defaults
+// (stream_idle_timeout_seconds, stream_max_duration_seconds,
+// stream_audit_max_events, stream_audit_max_event_bytes in
+// atryum.example.toml) so an embedder wiring NewService directly gets the
+// same protection the stock binary configures explicitly.
+const (
+ fallbackStreamIdleTimeout = 60 * time.Second
+ fallbackStreamMaxDuration = 10 * time.Minute
+ fallbackStreamAuditMaxEvents = 100
+ fallbackStreamAuditMaxEvtBytes = 4096
+)
+
+// effectiveStreamOptions returns the configured stream bounds, or — when
+// SetStreamOptions was never called — bounds derived from the service's own
+// default timeout plus the documented defaults. A zero mcp.StreamOptions
+// disables every bound, and the api-layer heartbeat keeps idle proxies from
+// killing the connection, so passing the zero value through would turn "the
+// embedder skipped optional wiring" into "relays run unbounded" — strictly
+// worse than the buffered path, which is always bounded by defaultTimeout.
+// The explicit flag (not a zero-value comparison) keeps an operator's
+// deliberate all-zeros configuration meaning what it says: unbounded.
+func (s *Service) effectiveStreamOptions() mcp.StreamOptions {
+ if s.streamOptionsSet {
+ return s.streamOptions
+ }
+ return mcp.StreamOptions{
+ HeaderTimeout: s.defaultTimeout,
+ IdleTimeout: fallbackStreamIdleTimeout,
+ MaxDuration: fallbackStreamMaxDuration,
+ }
+}
+
+// effectiveStreamAuditLimits is effectiveStreamOptions' counterpart for the
+// audit caps: the same never-wired embedder must not get unlimited
+// stream_event rows at up to 4 MiB of payload each, for the same reason it
+// must not get an unbounded relay.
+func (s *Service) effectiveStreamAuditLimits() StreamAuditLimits {
+ if s.streamOptionsSet {
+ return s.streamAuditLimits
+ }
+ return StreamAuditLimits{
+ MaxEvents: fallbackStreamAuditMaxEvents,
+ MaxEventBytes: fallbackStreamAuditMaxEvtBytes,
+ }
+}
+
+// RecordStreamDelivery records whether the handler delivered the terminal SSE
+// frame. Upstream execution and durable invocation state are separate from
+// this final agent-facing write, so delivery gets its own audit event.
+func (s *Service) RecordStreamDelivery(ctx context.Context, invocationID, status, message string) error {
+ if s.events == nil || invocationID == "" {
+ return nil
+ }
+ payload := map[string]any{"status": status}
+ if message != "" {
+ payload["message"] = message
+ }
+ return s.events.Create(ctx, Event{
+ InvocationID: invocationID,
+ EventType: "invocation.stream_delivery",
+ Payload: mustJSON(payload),
+ CreatedAt: time.Now().UTC(),
+ })
+}
+
// SetSessionStore installs the optional store backing the Invocations API
// session feature (POST /api/v1/external/sessions + session_id on Submit). When
// not installed, CreateSession returns an error and Submit ignores session_id.
@@ -304,7 +406,20 @@ func (s *Service) CreateSession(ctx context.Context, req CreateSessionRequest, a
}, nil
}
+// Invoke runs one tool call with no live relay: the upstream call is fully
+// buffered and returned as a single result, exactly as InvokeStreaming(ctx,
+// req, nil) would.
func (s *Service) Invoke(ctx context.Context, req CreateInvocationRequest) (InvocationResponse, error) {
+ return s.InvokeStreaming(ctx, req, nil)
+}
+
+// InvokeStreaming runs one tool call exactly like Invoke, except that a
+// non-nil sink can receive intermediate JSON-RPC messages from either HTTP
+// SSE or stdio as they arrive — see finishExecutionStreaming.
+// Rule evaluation, policy, and the human-approval gate below are entirely
+// unaware of sink: it is not touched until execution begins, so an
+// approval-gated call pauses with nothing relayed, the same as today.
+func (s *Service) InvokeStreaming(ctx context.Context, req CreateInvocationRequest, sink mcp.StreamSink) (InvocationResponse, error) {
if req.Server == "" {
return InvocationResponse{}, fmt.Errorf("server is required")
}
@@ -419,7 +534,7 @@ func (s *Service) Invoke(ctx context.Context, req CreateInvocationRequest) (Invo
}),
CreatedAt: time.Now().UTC(),
})
- return s.finishExecution(ctx, inv, upstream, req)
+ return s.finishExecution(ctx, inv, upstream, req, sink)
case planGateDeny:
planPayload["disposition"] = "plan_denied"
@@ -440,7 +555,7 @@ func (s *Service) Invoke(ctx context.Context, req CreateInvocationRequest) (Invo
if err := s.invocations.UpdateResult(ctx, inv); err != nil {
return InvocationResponse{}, err
}
- return s.waitForHumanApproval(ctx, inv, upstream, req)
+ return s.waitForHumanApproval(ctx, inv, upstream, req, sink)
}
}
@@ -526,12 +641,12 @@ func (s *Service) Invoke(ctx context.Context, req CreateInvocationRequest) (Invo
case policy.DispositionNever:
return s.denyByPolicy(ctx, inv, decision.Reason, aiConfidence)
case policy.DispositionAuto:
- return s.executeNow(ctx, inv, upstream, req, decision.Reason, aiConfidence)
+ return s.executeNow(ctx, inv, upstream, req, decision.Reason, aiConfidence, sink)
default:
// DispositionHuman, DispositionWorkflow, and dispositionAIEscalated all gate
// on a human decision. AI-escalated invocations are already tagged on inv.Approval
// above; waitForHumanApproval will persist the pending_approval status.
- return s.waitForHumanApproval(ctx, inv, upstream, req)
+ return s.waitForHumanApproval(ctx, inv, upstream, req, sink)
}
}
@@ -1073,7 +1188,7 @@ func (s *Service) denyByPolicy(ctx context.Context, inv Invocation, reason strin
}
// executeNow runs the tool call immediately without waiting for human approval.
-func (s *Service) executeNow(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest, reason string, confidence *float64) (InvocationResponse, error) {
+func (s *Service) executeNow(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest, reason string, confidence *float64, sink mcp.StreamSink) (InvocationResponse, error) {
inv.Status = StatusExecuting
inv.Approval = newApproval("auto_approved", reason, confidence)
if err := s.invocations.UpdateResult(ctx, inv); err != nil {
@@ -1089,11 +1204,40 @@ func (s *Service) executeNow(ctx context.Context, inv Invocation, upstream mcp.U
}),
CreatedAt: time.Now().UTC(),
})
- return s.finishExecution(ctx, inv, upstream, req)
+ return s.finishExecution(ctx, inv, upstream, req, sink)
}
// waitForHumanApproval blocks until an operator approves or denies, or the context is cancelled.
-func (s *Service) waitForHumanApproval(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest) (InvocationResponse, error) {
+//
+// pendingApprovals is an in-memory map local to this *Service* (one per OS
+// process). With multiple replicas, Approve/Deny landing on a different
+// process than the one blocked here cannot reach this goroutine's channel —
+// it can only update the durable row directly (see recordExternalDecision).
+// That produces a race whenever the original caller later disconnects:
+//
+// process A (this one) process B durable row
+// --------------------- --------- -----------
+// blocks here, waiting on
+// ch := pendingApprovals[id] -------------------------------> pending_approval
+// Approve(id) called
+// looks up its own
+// pendingApprovals[id]
+// -> not found locally
+// falls through to
+// recordExternalDecision -> approved
+// (still blocked; ch was
+// never written)
+// ctx.Done() fires
+// (client disconnected)
+// overwrites unconditionally -------------------------------> failed
+//
+// The tool is never executed even though the row briefly read "approved".
+// Durable, crash-resumable coordination (a work queue, an outbox table)
+// would close this gap; today a single active process is required for
+// Invoke-path human approval. TestMultiReplicaApprovalRaceOverwritesApprovedRowAsFailed
+// pins this exact sequence; docs/architecture.md "Atryum-executed calls" has
+// the full walkthrough.
+func (s *Service) waitForHumanApproval(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest, sink mcp.StreamSink) (InvocationResponse, error) {
ch := make(chan approvalDecision, 1)
s.mu.Lock()
s.pendingApprovals[inv.InvocationID] = ch
@@ -1196,27 +1340,45 @@ func (s *Service) waitForHumanApproval(ctx context.Context, inv Invocation, upst
Payload: mustJSON(executingPayload),
CreatedAt: approvedAt,
})
- return s.finishExecution(ctx, inv, upstream, req)
+ return s.finishExecution(ctx, inv, upstream, req, sink)
+}
+
+// finishExecution calls the upstream client and persists the outcome. When
+// sink is nil, the call is fully buffered exactly as before. When sink is
+// non-nil, the call is relayed live via InvokeStream: intermediate events
+// reach sink (wrapped in an audit decorator) as they arrive, and the fixed
+// s.defaultTimeout is replaced by s.streamOptions, since a long-lived relay
+// should not be judged against the same budget as an ordinary buffered call.
+func (s *Service) finishExecution(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest, sink mcp.StreamSink) (InvocationResponse, error) {
+ if sink == nil {
+ return s.finishExecutionBuffered(ctx, inv, upstream, req)
+ }
+ return s.finishExecutionStreaming(ctx, inv, upstream, req, sink)
}
-// finishExecution calls the upstream client and persists the outcome.
-func (s *Service) finishExecution(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest) (InvocationResponse, error) {
+func (s *Service) finishExecutionBuffered(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest) (InvocationResponse, error) {
execCtx, cancel := context.WithTimeout(ctx, s.defaultTimeout)
defer cancel()
- result, err := s.client.Invoke(execCtx, upstream, req.Tool, req.Input, req.RequestID)
+ result, err := s.client.Invoke(execCtx, upstream, req.Tool, req.Input, req.RequestID, req.Meta)
completed := time.Now().UTC()
inv.CompletedAt = &completed
+ // Finalization writes use their own bounded context, detached from the
+ // request: a downstream that disconnected mid-call (canceling ctx) must
+ // not leave the row stuck "executing" — the same guarantee the
+ // streaming path provides via terminalPersistenceTimeout.
+ persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout)
+ defer cancelPersist()
if err != nil {
inv.Status = StatusFailed
inv.Error = mustJSON(map[string]any{"message": err.Error()})
- _ = s.invocations.UpdateResult(ctx, inv)
- _ = s.events.Create(ctx, Event{InvocationID: inv.InvocationID, EventType: "invocation.failed", Payload: inv.Error, CreatedAt: completed})
+ _ = s.invocations.UpdateResult(persistCtx, inv)
+ _ = s.events.Create(persistCtx, Event{InvocationID: inv.InvocationID, EventType: "invocation.failed", Payload: inv.Error, CreatedAt: completed})
return s.toResponse(inv), nil
}
if result.Failed {
inv.Status = StatusFailed
inv.Error = result.Body
- _ = s.events.Create(ctx, Event{
+ _ = s.events.Create(persistCtx, Event{
InvocationID: inv.InvocationID, EventType: "invocation.failed",
Payload: mustJSON(map[string]any{"request_id": req.RequestID, "input": json.RawMessage(inv.Input), "arguments": json.RawMessage(inv.Input), "body": json.RawMessage(result.Body)}),
CreatedAt: completed,
@@ -1224,14 +1386,14 @@ func (s *Service) finishExecution(ctx context.Context, inv Invocation, upstream
} else {
inv.Status = StatusSucceeded
inv.Response = result.Body
- _ = s.events.Create(ctx, Event{
+ _ = s.events.Create(persistCtx, Event{
InvocationID: inv.InvocationID, EventType: "invocation.succeeded",
Payload: mustJSON(map[string]any{"request_id": req.RequestID, "input": json.RawMessage(inv.Input), "arguments": json.RawMessage(inv.Input), "body": json.RawMessage(result.Body)}),
CreatedAt: completed,
})
}
- if err := s.invocations.UpdateResult(ctx, inv); err != nil {
- return InvocationResponse{}, err
+ if err := s.invocations.UpdateResult(persistCtx, inv); err != nil {
+ return InvocationResponse{}, fmt.Errorf("persist buffered invocation result: %w", err)
}
if inv.Status == StatusSucceeded {
s.completePlanAfterSuccessfulFinalAction(ctx, inv)
diff --git a/internal/invocation/service_test.go b/internal/invocation/service_test.go
index c8edb432..1602fbac 100644
--- a/internal/invocation/service_test.go
+++ b/internal/invocation/service_test.go
@@ -70,12 +70,21 @@ while IFS= read -r line; do
if [[ -z "$line" ]]; then
continue
fi
+ # Echo back whatever id the request actually carried — Atryum's stdio
+ # client assigns ids from a shared per-client counter (notifications
+ # consume one too, even though it's stripped before sending), so a
+ # hardcoded id here would drift from what's actually sent. Only computed
+ # inside the branches that have one: the notification line below carries
+ # no "id" field at all, and grep finding nothing there would (under
+ # pipefail) abort the script if this ran unconditionally.
if [[ "$line" == *'"method":"initialize"'* ]]; then
- printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"serverInfo":{"name":"fake-shortcut","version":"0.1.0"},"capabilities":{}}}'
+ id=$(echo "$line" | grep -o '"id":[0-9]*' | head -1 | cut -d: -f2)
+ printf '%s\n' "{\"jsonrpc\":\"2.0\",\"id\":${id},\"result\":{\"serverInfo\":{\"name\":\"fake-shortcut\",\"version\":\"0.1.0\"},\"capabilities\":{}}}"
elif [[ "$line" == *'"method":"notifications/initialized"'* ]]; then
continue
elif [[ "$line" == *'"method":"tools/call"'* ]]; then
- printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{"content":[{"type":"text","text":"ok"}]}}'
+ id=$(echo "$line" | grep -o '"id":[0-9]*' | head -1 | cut -d: -f2)
+ printf '%s\n' "{\"jsonrpc\":\"2.0\",\"id\":${id},\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"ok\"}]}}"
exit 0
fi
done
@@ -188,6 +197,153 @@ func TestInvokeFailsClosedWhenRuleLoadFails(t *testing.T) {
}
}
+// TestMultiReplicaApprovalRaceOverwritesApprovedRowAsFailed pins the
+// documented multi-replica race for Atryum-executed (Invoke) human approval
+// (see docs/architecture.md "Atryum-executed calls"): pendingApprovals is an
+// in-memory map local to one *Service*, so an Approve call landing on a
+// different Service instance — standing in for a second replica — cannot
+// wake the goroutine blocked in waitForHumanApproval on the first instance.
+// It can only update the durable row directly. When the original caller's
+// context is later cancelled, that goroutine wakes on ctx.Done() and
+// unconditionally overwrites the row to failed, clobbering the approval a
+// different process just persisted. The tool must never execute.
+//
+// The production comment on waitForHumanApproval shows the race in terms of
+// process A / process B / the durable row; this test drives that exact
+// sequence with its own goroutine, polling loop, and assertions:
+//
+// background goroutine (processA) main goroutine (test)
+// -------------------------------- ----------------------
+// processA.Invoke(ctx, ...)
+// blocks in waitForHumanApproval
+// pollUntil polls processB.List
+// until StatusPendingApproval,
+// captures invocationID
+// processB.Approve(invocationID,
+// "operator")
+// processA.Get(invocationID)
+// -> Approved
+// select done: default
+// (still blocked — the approval
+// never woke processA.Invoke)
+// cancel()
+// ctx.Done() fires, overwrites the row
+// unconditionally -> Failed
+// done <- {Status: Failed, "...cancelled..."}
+// <-done (within 2s)
+// assert Status == Failed
+// processA.Get(invocationID)
+// -> Failed (processB's
+// approved write is gone)
+func TestMultiReplicaApprovalRaceOverwritesApprovedRowAsFailed(t *testing.T) {
+ upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var body map[string]any
+ _ = json.NewDecoder(r.Body).Decode(&body)
+ switch body["method"] {
+ case "initialize":
+ _ = json.NewEncoder(w).Encode(map[string]any{
+ "jsonrpc": "2.0", "id": body["id"],
+ "result": map[string]any{"serverInfo": map[string]any{"name": "fake", "version": "0.1.0"}, "capabilities": map[string]any{}},
+ })
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ // Should never be reached: the row is overwritten to failed before
+ // anything is ever approved from process A's own perspective.
+ t.Error("tool executed even though the invocation was ultimately marked failed")
+ default:
+ w.WriteHeader(http.StatusBadRequest)
+ }
+ }))
+ defer upstream.Close()
+
+ db := newSQLiteTestDB(t)
+ serverRepo := store.NewServerRepo(db)
+ resolver := mcp.NewResolver(serverRepo, config.Config{
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+ if err := resolver.BootstrapIfEmpty(context.Background()); err != nil {
+ t.Fatal(err)
+ }
+ newProcess := func() *invocation.Service {
+ return invocation.NewService(store.NewInvocationRepo(db), store.NewEventRepo(db), resolver, mcp.NewHTTPClient(), policy.ManualApprovalProvider{}, 5*time.Second, nil, nil, nil, nil)
+ }
+ processA := newProcess() // blocks the call below in its own in-memory waiter
+ processB := newProcess() // simulates a second replica handling the admin decision; shares the same db, not the same pendingApprovals map
+
+ ctx, cancel := context.WithCancel(context.Background())
+ defer cancel()
+
+ type invokeResult struct {
+ resp invocation.InvocationResponse
+ err error
+ }
+ done := make(chan invokeResult, 1)
+ go func() {
+ resp, err := processA.Invoke(ctx, invocation.CreateInvocationRequest{Server: "shortcut", Tool: "dangerous-tool", Input: map[string]any{}})
+ done <- invokeResult{resp, err}
+ }()
+
+ var invocationID string
+ pollUntil(t, 2*time.Second, 10*time.Millisecond, func() bool {
+ list, err := processB.List(context.Background(), invocation.InvocationListFilter{Limit: 10})
+ if err == nil && len(list.Items) == 1 && list.Items[0].Status == invocation.StatusPendingApproval {
+ invocationID = list.Items[0].InvocationID
+ return true
+ }
+ return false
+ })
+ if invocationID == "" {
+ t.Fatal("invocation never reached pending_approval")
+ }
+
+ if err := processB.Approve(context.Background(), invocationID, "operator"); err != nil {
+ t.Fatalf("approve from process B: %v", err)
+ }
+
+ approvedRead, err := processA.Get(context.Background(), invocationID)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if approvedRead.Status != invocation.StatusApproved {
+ t.Fatalf("expected process B's approval to be visible as %q, got %q", invocation.StatusApproved, approvedRead.Status)
+ }
+
+ select {
+ case <-done:
+ t.Fatal("process A's Invoke returned before its context was cancelled; the approval should not have woken it")
+ default:
+ }
+
+ // The caller now disconnects (or its request times out): cancel the
+ // context processA.Invoke is still blocked on.
+ cancel()
+
+ var result invokeResult
+ select {
+ case result = <-done:
+ case <-time.After(2 * time.Second):
+ t.Fatal("processA.Invoke did not return after its context was cancelled")
+ }
+ if result.err != nil {
+ t.Fatal(result.err)
+ }
+ if result.resp.Status != invocation.StatusFailed {
+ t.Fatalf("status = %q, want failed (context cancellation must overwrite even an already-approved row)", result.resp.Status)
+ }
+ if !strings.Contains(string(result.resp.Error), "cancelled") {
+ t.Fatalf("expected the stored error text to mention cancellation, got %s", result.resp.Error)
+ }
+
+ finalRead, err := processA.Get(context.Background(), invocationID)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if finalRead.Status != invocation.StatusFailed {
+ t.Fatalf("final row status = %q, want failed — process B's approved write must be overwritten, not preserved", finalRead.Status)
+ }
+}
+
// TestSubmitLogsAndAuditsRuleLoadFailure verifies that Submit records a
// rule-lookup failure in the invocation's audit trail instead of discarding
// it. Submit already fell back to pending_approval before this fix (unlike
@@ -367,17 +523,13 @@ func TestSubmitPendingApprovalAutomaticallySummarizesInvocation(t *testing.T) {
}
var got invocation.InvocationResponse
- deadline := time.Now().Add(2 * time.Second)
- for time.Now().Before(deadline) {
+ pollUntil(t, 2*time.Second, 10*time.Millisecond, func() bool {
got, err = service.Get(context.Background(), resp.InvocationID)
if err != nil {
t.Fatal(err)
}
- if got.Summary == "Run shell command ls." {
- break
- }
- time.Sleep(10 * time.Millisecond)
- }
+ return got.Summary == "Run shell command ls."
+ })
if got.Summary != "Run shell command ls." {
t.Fatalf("expected automatic summary, got %q", got.Summary)
}
@@ -1392,6 +1544,22 @@ func TestInvokeAIEvaluationToleratesMissingToolDescription(t *testing.T) {
}
}
+// pollUntil calls cond repeatedly, sleeping interval between tries, until it
+// returns true or timeout elapses. Shared by tests waiting on asynchronous
+// state (a row reaching a given status, a field getting populated) instead
+// of sleeping a fixed guess that's either too slow or too flaky.
+func pollUntil(t *testing.T, timeout, interval time.Duration, cond func() bool) bool {
+ t.Helper()
+ deadline := time.Now().Add(timeout)
+ for time.Now().Before(deadline) {
+ if cond() {
+ return true
+ }
+ time.Sleep(interval)
+ }
+ return cond()
+}
+
func newTestService(t *testing.T, cfg config.Config) *invocation.Service {
t.Helper()
db := newSQLiteTestDB(t)
diff --git a/internal/invocation/stream_execution.go b/internal/invocation/stream_execution.go
new file mode 100644
index 00000000..8384b61c
--- /dev/null
+++ b/internal/invocation/stream_execution.go
@@ -0,0 +1,108 @@
+package invocation
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "time"
+
+ "github.com/validmind/atryum/internal/mcp"
+)
+
+// finishExecutionStreaming is finishExecution's live-relay path. It audits
+// intermediate events and the durable upstream outcome; the API handler
+// records terminal-frame delivery separately because that write happens only
+// after this method returns.
+func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, upstream mcp.Upstream, req CreateInvocationRequest, sink mcp.StreamSink) (InvocationResponse, error) {
+ audited := newAuditingSink(sink, s.events, inv.InvocationID, req.RequestID, upstream.Name, s.effectiveStreamAuditLimits())
+ result, err := s.client.InvokeStream(ctx, upstream, req.Tool, req.Input, req.RequestID, req.Meta, audited, s.effectiveStreamOptions())
+ completed := time.Now().UTC()
+ inv.CompletedAt = &completed
+
+ if err != nil {
+ inv.Status = StatusFailed
+ reason, message := classifyStreamError(ctx, audited, err)
+ inv.Error = mustJSON(map[string]any{"message": message})
+ persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout)
+ defer cancelPersist()
+ if updateErr := s.invocations.UpdateResult(persistCtx, inv); updateErr != nil {
+ audited.finish(completed, "persistence_failed")
+ return s.toResponse(inv), fmt.Errorf("persist streaming invocation failure: %w", updateErr)
+ }
+ audited.finish(completed, "failed")
+ // stream_completed (the summary of what happened during the relay)
+ // is written before the invocation-level failed/succeeded event, so
+ // an audit trail read chronologically sees "here's what the stream
+ // did" before "here's how the invocation ended" — the natural
+ // narrative order, even though both share the same timestamp.
+ _ = s.events.Create(persistCtx, Event{
+ InvocationID: inv.InvocationID, EventType: "invocation.failed",
+ Payload: mustJSON(map[string]any{"reason": reason, "message": message, "events_total": audited.seq}),
+ CreatedAt: completed,
+ })
+ return s.toResponse(inv), nil
+ }
+ var terminalEvent Event
+ if result.Failed {
+ inv.Status = StatusFailed
+ inv.Error = result.Body
+ terminalEvent = Event{
+ InvocationID: inv.InvocationID, EventType: "invocation.failed",
+ Payload: mustJSON(map[string]any{"request_id": req.RequestID, "input": json.RawMessage(inv.Input), "arguments": json.RawMessage(inv.Input), "body": json.RawMessage(result.Body)}),
+ CreatedAt: completed,
+ }
+ } else {
+ inv.Status = StatusSucceeded
+ inv.Response = result.Body
+ terminalEvent = Event{
+ InvocationID: inv.InvocationID, EventType: "invocation.succeeded",
+ Payload: mustJSON(map[string]any{"request_id": req.RequestID, "input": json.RawMessage(inv.Input), "arguments": json.RawMessage(inv.Input), "body": json.RawMessage(result.Body)}),
+ CreatedAt: completed,
+ }
+ }
+ persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout)
+ defer cancelPersist()
+ if err := s.invocations.UpdateResult(persistCtx, inv); err != nil {
+ audited.finish(completed, "persistence_failed")
+ return s.toResponse(inv), fmt.Errorf("persist streaming invocation result: %w", err)
+ }
+ if result.Failed {
+ audited.finish(completed, "failed")
+ } else {
+ audited.finish(completed, "succeeded")
+ }
+ _ = s.events.Create(persistCtx, terminalEvent)
+ return s.toResponse(inv), nil
+}
+
+// classifyStreamError distinguishes why InvokeStream returned an error, in
+// priority order: the sink itself (i.e. the downstream agent connection)
+// failing first — that's the caller's own signal and always the most
+// specific one available — then the client's own header/idle/max-duration
+// bound (mcp.ErrStreamTimeout), then anything else as a generic transport
+// failure. The reason is persisted on the invocation.failed audit event so
+// it can be told apart from an ordinary transport error after the fact.
+func classifyStreamError(ctx context.Context, audited *auditingSink, err error) (reason string, message string) {
+ if audited.downstreamErr != nil {
+ return "stream_aborted_downstream", audited.downstreamErr.Error()
+ }
+ if errors.Is(err, context.Canceled) || errors.Is(ctx.Err(), context.Canceled) {
+ // The request context died with no failed downstream write to prove
+ // who went away: a quietly-disconnected agent and a server shutdown
+ // are indistinguishable from here, so the reason stays neutral
+ // rather than blaming the downstream for every in-flight call when
+ // the process stops.
+ return "stream_canceled", err.Error()
+ }
+ if errors.Is(err, mcp.ErrStreamTimeout) {
+ return "stream_timeout", err.Error()
+ }
+ if errors.Is(err, mcp.ErrStreamSessionRetryRefused) {
+ return "stream_session_retry_refused", err.Error()
+ }
+ if errors.Is(err, mcp.ErrStreamMessageTooLarge) {
+ return "stream_message_too_large", err.Error()
+ }
+ return "transport_error", err.Error()
+}
diff --git a/internal/invocation/stream_execution_test.go b/internal/invocation/stream_execution_test.go
new file mode 100644
index 00000000..a7b6bd31
--- /dev/null
+++ b/internal/invocation/stream_execution_test.go
@@ -0,0 +1,823 @@
+package invocation_test
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "net/http"
+ "net/http/httptest"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "github.com/validmind/atryum/internal/config"
+ "github.com/validmind/atryum/internal/invocation"
+ "github.com/validmind/atryum/internal/invocation/policy"
+ "github.com/validmind/atryum/internal/mcp"
+ "github.com/validmind/atryum/internal/store"
+)
+
+// recordingSink is a test mcp.StreamSink that records what it received.
+// onEvent, when set, lets a test hook into delivery (e.g. to synchronize
+// with a background approval goroutine). Safe for concurrent use: Event may
+// run on the InvokeStreaming call's goroutine while a test's assertions run
+// on another.
+type recordingSink struct {
+ mu sync.Mutex
+ started bool
+ events []mcp.StreamEvent
+ onEvent func(mcp.StreamEvent) error
+}
+
+type blockingStreamEventRepo struct {
+ inner *store.EventRepo
+ started chan struct{}
+ once sync.Once
+}
+
+type failTerminalUpdateRepo struct {
+ *store.InvocationRepo
+}
+
+func (r *failTerminalUpdateRepo) UpdateResult(ctx context.Context, inv invocation.Invocation) error {
+ if inv.Status == invocation.StatusSucceeded || inv.Status == invocation.StatusFailed {
+ return errors.New("terminal persistence unavailable")
+ }
+ return r.InvocationRepo.UpdateResult(ctx, inv)
+}
+
+func (r *blockingStreamEventRepo) Create(ctx context.Context, evt invocation.Event) error {
+ if evt.EventType == "invocation.stream_event" {
+ r.once.Do(func() { close(r.started) })
+ <-ctx.Done()
+ return ctx.Err()
+ }
+ return r.inner.Create(ctx, evt)
+}
+
+func (r *blockingStreamEventRepo) ListByInvocation(ctx context.Context, invocationID string, filter invocation.EventListFilter) ([]invocation.Event, int, error) {
+ return r.inner.ListByInvocation(ctx, invocationID, filter)
+}
+
+func (s *recordingSink) StreamStarted() {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.started = true
+}
+
+func (s *recordingSink) Event(evt mcp.StreamEvent) error {
+ s.mu.Lock()
+ s.events = append(s.events, evt)
+ s.mu.Unlock()
+ if s.onEvent != nil {
+ return s.onEvent(evt)
+ }
+ return nil
+}
+
+func (s *recordingSink) touched() bool {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ return s.started || len(s.events) > 0
+}
+
+func (s *recordingSink) eventCount() int {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ return len(s.events)
+}
+
+// sseToolCallUpstream builds an httptest.Server implementing the
+// initialize/notifications.initialized handshake, dispatching tools/call to
+// callHandler so a test controls exactly what SSE bytes are written.
+func sseToolCallUpstream(t *testing.T, callHandler func(w http.ResponseWriter, r *http.Request, body map[string]any)) *httptest.Server {
+ t.Helper()
+ return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var body map[string]any
+ if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+ // t.Errorf + an HTTP error, never t.Fatal: FailNow from a
+ // non-test goroutine only kills the handler, and the
+ // never-answered client would hang the test until the suite
+ // timeout instead of failing it cleanly.
+ t.Errorf("decode request: %v", err)
+ http.Error(w, "bad request body", http.StatusBadRequest)
+ return
+ }
+ switch body["method"] {
+ case "initialize":
+ _ = json.NewEncoder(w).Encode(map[string]any{
+ "jsonrpc": "2.0", "id": body["id"],
+ "result": map[string]any{"serverInfo": map[string]any{"name": "fake", "version": "0.1.0"}, "capabilities": map[string]any{}},
+ })
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ callHandler(w, r, body)
+ default:
+ t.Errorf("unexpected method %q", body["method"])
+ http.Error(w, "unexpected method", http.StatusInternalServerError)
+ }
+ }))
+}
+
+func writeSSEEvent(w http.ResponseWriter, flusher http.Flusher, data string) {
+ _, _ = w.Write([]byte("event: message\ndata: " + data + "\n\n"))
+ flusher.Flush()
+}
+
+func TestInvokeStreamingRelaysEventsAndAuditsThem(t *testing.T) {
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer upstream.Close()
+
+ service := newTestService(t, config.Config{
+ Defaults: config.DefaultsConfig{RequestTimeoutSeconds: 5},
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+
+ sink := &recordingSink{}
+ resp, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{"n": 1}}, sink)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusSucceeded {
+ t.Fatalf("expected succeeded, got %s: %s", resp.Status, resp.Error)
+ }
+ if !sink.started {
+ t.Fatal("expected StreamStarted to fire")
+ }
+ if sink.eventCount() != 1 {
+ t.Fatalf("expected exactly one relayed event, got %d", sink.eventCount())
+ }
+ if !jsonContains(resp.Result, "done") {
+ t.Fatalf("expected terminal result body, got %s", resp.Result)
+ }
+
+ events, err := service.Events(context.Background(), resp.InvocationID, invocation.EventListFilter{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ var sawStreamEvent, sawStreamCompleted bool
+ for _, evt := range events.Items {
+ switch evt.Type {
+ case "invocation.stream_event":
+ sawStreamEvent = true
+ var payload struct {
+ Seq int `json:"seq"`
+ UpstreamName string `json:"upstream_name"`
+ }
+ if err := json.Unmarshal(evt.Data, &payload); err != nil {
+ t.Fatalf("decode invocation.stream_event payload: %v", err)
+ }
+ if payload.Seq != 1 {
+ t.Fatalf("expected seq 1, got %d", payload.Seq)
+ }
+ if payload.UpstreamName != "shortcut" {
+ t.Fatalf("expected upstream_name shortcut, got %q", payload.UpstreamName)
+ }
+ case "invocation.stream_completed":
+ sawStreamCompleted = true
+ var payload struct {
+ EventsTotal int `json:"events_total"`
+ Terminal string `json:"terminal"`
+ }
+ if err := json.Unmarshal(evt.Data, &payload); err != nil {
+ t.Fatalf("decode invocation.stream_completed payload: %v", err)
+ }
+ if payload.EventsTotal != 1 {
+ t.Fatalf("expected events_total 1, got %d", payload.EventsTotal)
+ }
+ if payload.Terminal != "succeeded" {
+ t.Fatalf("expected terminal succeeded, got %q", payload.Terminal)
+ }
+ }
+ }
+ if !sawStreamEvent {
+ t.Fatal("expected an invocation.stream_event audit row")
+ }
+ if !sawStreamCompleted {
+ t.Fatal("expected an invocation.stream_completed audit row")
+ }
+}
+
+func TestInvokeStreamingAuditCapsEnforcedWithoutSuppressingRelay(t *testing.T) {
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ for i := 0; i < 3; i++ {
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ }
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer upstream.Close()
+
+ db := newSQLiteTestDB(t)
+ serverRepo := store.NewServerRepo(db)
+ cfg := config.Config{Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}}
+ resolver := mcp.NewResolver(serverRepo, cfg)
+ if err := resolver.BootstrapIfEmpty(context.Background()); err != nil {
+ t.Fatal(err)
+ }
+ service := invocation.NewService(
+ store.NewInvocationRepo(db), store.NewEventRepo(db), resolver, mcp.NewHTTPClient(),
+ policy.AlwaysApproveProvider{}, 5*time.Second, nil, nil, nil, nil,
+ )
+ service.SetStreamOptions(mcp.StreamOptions{}, invocation.StreamAuditLimits{MaxEvents: 1})
+
+ sink := &recordingSink{}
+ resp, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, sink)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusSucceeded {
+ t.Fatalf("expected succeeded, got %s: %s", resp.Status, resp.Error)
+ }
+ // The cap bounds what's persisted, not what's relayed: the agent-facing
+ // sink must still see every event even once the audit log stops
+ // recording them individually.
+ if sink.eventCount() != 3 {
+ t.Fatalf("expected all 3 events relayed to the sink despite the audit cap, got %d", sink.eventCount())
+ }
+
+ events, err := service.Events(context.Background(), resp.InvocationID, invocation.EventListFilter{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ streamEventRows := 0
+ for _, evt := range events.Items {
+ if evt.Type == "invocation.stream_event" {
+ streamEventRows++
+ }
+ if evt.Type == "invocation.stream_completed" {
+ var payload struct {
+ EventsTotal int `json:"events_total"`
+ }
+ if err := json.Unmarshal(evt.Data, &payload); err != nil {
+ t.Fatalf("decode invocation.stream_completed payload: %v", err)
+ }
+ if payload.EventsTotal != 3 {
+ t.Fatalf("expected events_total to reflect the true count (3) even though only 1 was persisted, got %d", payload.EventsTotal)
+ }
+ }
+ }
+ if streamEventRows != 1 {
+ t.Fatalf("expected exactly 1 persisted invocation.stream_event row (MaxEvents cap), got %d", streamEventRows)
+ }
+}
+
+// TestInvokeStreamingBlockedAuditWriteDoesNotDelayRelay proves a stalled
+// audit repository cannot delay live delivery: recordEvent's enqueue onto
+// the shared streamAuditDispatcher is non-blocking, so the worker goroutine
+// stuck inside Create is a separate actor from the relay path. Three
+// channels pin that interleaving regardless of scheduling:
+//
+// main goroutine InvokeStreaming goroutine audit worker (runWorker)
+// --------------- -------------------------- ------------------------
+// go func(){ InvokeStreaming }()
+// sink.Event -> recordEvent
+// enqueue(write) --non-blocking--> dequeues write
+// Create(): close(events.started),
+// blocks on <-writeCtx.Done()
+// <-events.started returns
+// inner.Event(evt) -> onEvent hook
+// close(delivered)
+// <-delivered returns (<200ms:
+// relay was never gated on
+// the still-blocked write)
+// writeCtx (streamAuditWriteTimeout)
+// expires -> Create returns ctx.Err()
+// finish() drains pending writes,
+// InvokeStreaming returns Succeeded
+// <-done returns (<3s)
+//
+// Had the audit write instead run synchronously before relaying, "delivered"
+// could never close until the blocked write's own timeout had passed.
+func TestInvokeStreamingBlockedAuditWriteDoesNotDelayRelay(t *testing.T) {
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer upstream.Close()
+
+ db := newSQLiteTestDB(t)
+ serverRepo := store.NewServerRepo(db)
+ resolver := mcp.NewResolver(serverRepo, config.Config{Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}})
+ if err := resolver.BootstrapIfEmpty(context.Background()); err != nil {
+ t.Fatal(err)
+ }
+ events := &blockingStreamEventRepo{inner: store.NewEventRepo(db), started: make(chan struct{})}
+ service := invocation.NewService(
+ store.NewInvocationRepo(db), events, resolver, mcp.NewHTTPClient(),
+ policy.AlwaysApproveProvider{}, 5*time.Second, nil, nil, nil, nil,
+ )
+
+ delivered := make(chan struct{})
+ var deliveredOnce sync.Once
+ sink := &recordingSink{onEvent: func(mcp.StreamEvent) error {
+ deliveredOnce.Do(func() { close(delivered) })
+ return nil
+ }}
+ done := make(chan error, 1)
+ go func() {
+ resp, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, sink)
+ if err == nil && resp.Status != invocation.StatusSucceeded {
+ err = fmt.Errorf("status = %s, want succeeded", resp.Status)
+ }
+ done <- err
+ }()
+
+ select {
+ case <-events.started:
+ case <-time.After(time.Second):
+ t.Fatal("audit write did not start")
+ }
+ select {
+ case <-delivered:
+ case <-time.After(200 * time.Millisecond):
+ t.Fatal("relay waited for blocked audit storage")
+ }
+ select {
+ case err := <-done:
+ if err != nil {
+ t.Fatal(err)
+ }
+ case <-time.After(3 * time.Second):
+ t.Fatal("stream did not finish after bounded audit write timed out")
+ }
+}
+
+func TestInvokeStreamingSinkAbortMarksFailedAsDownstreamAborted(t *testing.T) {
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":2}}`)
+ })
+ defer upstream.Close()
+
+ service := newTestService(t, config.Config{
+ Defaults: config.DefaultsConfig{RequestTimeoutSeconds: 5},
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+
+ sink := &recordingSink{onEvent: func(mcp.StreamEvent) error {
+ return errors.New("downstream connection closed")
+ }}
+ resp, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, sink)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusFailed {
+ t.Fatalf("expected failed status, got %s", resp.Status)
+ }
+
+ events, err := service.Events(context.Background(), resp.InvocationID, invocation.EventListFilter{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ found := false
+ for _, evt := range events.Items {
+ if evt.Type == "invocation.failed" && jsonContains(evt.Data, "stream_aborted_downstream") {
+ found = true
+ }
+ }
+ if !found {
+ t.Fatal("expected an invocation.failed event with reason stream_aborted_downstream")
+ }
+}
+
+func TestInvokeStreamingSinkAbortPersistsFailureAfterRequestContextCancellation(t *testing.T) {
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ })
+ defer upstream.Close()
+
+ service := newTestService(t, config.Config{
+ Defaults: config.DefaultsConfig{RequestTimeoutSeconds: 5},
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+
+ ctx, cancel := context.WithCancel(context.Background())
+ sink := &recordingSink{onEvent: func(mcp.StreamEvent) error {
+ cancel() // net/http cancels the request context when the agent disconnects.
+ return errors.New("downstream connection closed")
+ }}
+ resp, err := service.InvokeStreaming(ctx, invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, sink)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+
+ persisted, err := service.Get(context.Background(), resp.InvocationID)
+ if err != nil {
+ t.Fatalf("read persisted invocation: %v", err)
+ }
+ if persisted.Status != invocation.StatusFailed {
+ t.Fatalf("persisted status = %s, want failed after downstream disconnect", persisted.Status)
+ }
+
+ events, err := service.Events(context.Background(), resp.InvocationID, invocation.EventListFilter{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ for _, evt := range events.Items {
+ if evt.Type == "invocation.failed" && jsonContains(evt.Data, "stream_aborted_downstream") {
+ return
+ }
+ }
+ t.Fatal("expected persisted invocation.failed event with reason stream_aborted_downstream")
+}
+
+func TestInvokeStreamingQuietRequestCancellationIsAuditedAsCanceled(t *testing.T) {
+ callStarted := make(chan struct{})
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ w.(http.Flusher).Flush()
+ close(callStarted)
+ <-r.Context().Done()
+ })
+ defer upstream.Close()
+
+ service := newTestService(t, config.Config{
+ Defaults: config.DefaultsConfig{RequestTimeoutSeconds: 5},
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+
+ ctx, cancel := context.WithCancel(context.Background())
+ resultCh := make(chan invocation.InvocationResponse, 1)
+ errCh := make(chan error, 1)
+ go func() {
+ resp, err := service.InvokeStreaming(ctx, invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, &recordingSink{})
+ resultCh <- resp
+ errCh <- err
+ }()
+ <-callStarted
+ cancel()
+
+ resp := <-resultCh
+ if err := <-errCh; err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusFailed {
+ t.Fatalf("status = %s, want failed", resp.Status)
+ }
+ events, err := service.Events(context.Background(), resp.InvocationID, invocation.EventListFilter{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ for _, evt := range events.Items {
+ // stream_canceled, not stream_aborted_downstream: with no failed
+ // downstream write to prove who went away, a quiet disconnect and a
+ // server shutdown are indistinguishable, so the audit reason must
+ // not blame the agent.
+ if evt.Type == "invocation.failed" && jsonContains(evt.Data, "stream_canceled") {
+ return
+ }
+ }
+ t.Fatal("quiet request cancellation was not audited as stream_canceled")
+}
+
+func TestInvokeStreamingDoesNotAuditSuccessBeforeTerminalPersistence(t *testing.T) {
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer upstream.Close()
+
+ db := newSQLiteTestDB(t)
+ serverRepo := store.NewServerRepo(db)
+ resolver := mcp.NewResolver(serverRepo, config.Config{
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+ if err := resolver.BootstrapIfEmpty(context.Background()); err != nil {
+ t.Fatal(err)
+ }
+ eventRepo := store.NewEventRepo(db)
+ service := invocation.NewService(
+ &failTerminalUpdateRepo{InvocationRepo: store.NewInvocationRepo(db)},
+ eventRepo,
+ resolver,
+ mcp.NewHTTPClient(),
+ policy.AlwaysApproveProvider{},
+ 5*time.Second,
+ nil, nil, nil, nil,
+ )
+
+ resp, err := service.InvokeStreaming(
+ context.Background(),
+ invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}},
+ &recordingSink{},
+ )
+ if err == nil {
+ t.Fatal("expected terminal persistence error")
+ }
+ if resp.InvocationID == "" {
+ t.Fatal("expected the invocation ID to remain available for delivery audit")
+ }
+
+ events, _, listErr := eventRepo.ListByInvocation(context.Background(), resp.InvocationID, invocation.EventListFilter{Limit: 100})
+ if listErr != nil {
+ t.Fatal(listErr)
+ }
+ for _, evt := range events {
+ if evt.EventType == "invocation.succeeded" {
+ t.Fatal("found invocation.succeeded despite terminal persistence failure")
+ }
+ if evt.EventType == "invocation.stream_completed" && jsonContains(evt.Payload, `"terminal":"succeeded"`) {
+ t.Fatal("stream completion claimed success before terminal persistence")
+ }
+ }
+}
+
+func TestInvokeStreamingIdleTimeoutMarksFailedAsStreamTimeout(t *testing.T) {
+ blockUntilTestDone := make(chan struct{})
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ <-blockUntilTestDone // never send the terminal event
+ })
+ t.Cleanup(func() {
+ close(blockUntilTestDone)
+ upstream.Close()
+ })
+
+ db := newSQLiteTestDB(t)
+ serverRepo := store.NewServerRepo(db)
+ cfg := config.Config{Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}}
+ resolver := mcp.NewResolver(serverRepo, cfg)
+ if err := resolver.BootstrapIfEmpty(context.Background()); err != nil {
+ t.Fatal(err)
+ }
+ service := invocation.NewService(
+ store.NewInvocationRepo(db), store.NewEventRepo(db), resolver, mcp.NewHTTPClient(),
+ policy.AlwaysApproveProvider{}, 5*time.Second, nil, nil, nil, nil,
+ )
+ service.SetStreamOptions(mcp.StreamOptions{IdleTimeout: 50 * time.Millisecond}, invocation.StreamAuditLimits{})
+
+ sink := &recordingSink{}
+ resp, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, sink)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusFailed {
+ t.Fatalf("expected failed status, got %s", resp.Status)
+ }
+
+ events, err := service.Events(context.Background(), resp.InvocationID, invocation.EventListFilter{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ found := false
+ for _, evt := range events.Items {
+ if evt.Type == "invocation.failed" && jsonContains(evt.Data, "stream_timeout") {
+ found = true
+ }
+ }
+ if !found {
+ t.Fatal("expected an invocation.failed event with reason stream_timeout")
+ }
+}
+
+func TestInvokeStreamingMidStreamSessionRetryRefusalMarksFailedWithDistinctReason(t *testing.T) {
+ // Atomic: incremented on handler goroutines, read on the test goroutine.
+ var toolsCallCount atomic.Int32
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ toolsCallCount.Add(1)
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"No session ID provided for non-initialization request"}}`)
+ })
+ defer upstream.Close()
+
+ service := newTestService(t, config.Config{
+ Defaults: config.DefaultsConfig{RequestTimeoutSeconds: 5},
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+
+ sink := &recordingSink{}
+ resp, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, sink)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusFailed {
+ t.Fatalf("expected failed status, got %s", resp.Status)
+ }
+ if got := toolsCallCount.Load(); got != 1 {
+ t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed mid-stream)", got)
+ }
+
+ events, err := service.Events(context.Background(), resp.InvocationID, invocation.EventListFilter{})
+ if err != nil {
+ t.Fatal(err)
+ }
+ found := false
+ for _, evt := range events.Items {
+ if evt.Type == "invocation.failed" && jsonContains(evt.Data, "stream_session_retry_refused") {
+ found = true
+ }
+ }
+ if !found {
+ t.Fatal("expected an invocation.failed event with reason stream_session_retry_refused, distinguishable from a generic transport_error")
+ }
+}
+
+// TestInvokeStreamingApprovalGateDoesNotTouchSinkBeforeApproval races a
+// background approval against the InvokeStreaming call blocked in
+// waitForHumanApproval, proving the sink is never touched until approval
+// unblocks execution:
+//
+// main goroutine (InvokeStreaming) background goroutine
+// --------------------------------- ---------------------
+// InvokeStreaming(..., sink) called
+// ManualApprovalProvider -> pending
+// blocks in waitForHumanApproval
+// pollUntil polls service.List
+// until an item reaches
+// StatusPendingApproval
+// sink.touched() checked -> false
+// service.Approve(pendingID, "")
+// unblocks, executeNow runs the tool
+// sink.Event(...) fires -> touched=true
+// InvokeStreaming returns Succeeded
+//
+// Polling for pending_approval (instead of a fixed sleep) is what keeps this
+// deterministic: Approve only runs once the row is confirmed pending, so the
+// sink-touched check above always happens before it.
+func TestInvokeStreamingApprovalGateDoesNotTouchSinkBeforeApproval(t *testing.T) {
+ upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeSSEEvent(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer upstream.Close()
+
+ db := newSQLiteTestDB(t)
+ serverRepo := store.NewServerRepo(db)
+ cfg := config.Config{Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}}
+ resolver := mcp.NewResolver(serverRepo, cfg)
+ if err := resolver.BootstrapIfEmpty(context.Background()); err != nil {
+ t.Fatal(err)
+ }
+ service := invocation.NewService(
+ store.NewInvocationRepo(db), store.NewEventRepo(db), resolver, mcp.NewHTTPClient(),
+ policy.ManualApprovalProvider{}, 5*time.Second, nil, nil, nil, nil,
+ )
+
+ sink := &recordingSink{}
+ go func() {
+ // Poll for the invocation to actually reach pending_approval rather
+ // than sleeping a fixed interval: on a slow machine a fixed sleep can
+ // catch the row while still "received", making Approve fail and
+ // InvokeStreaming block on the approval channel until the suite
+ // timeout.
+ var pendingID string
+ pollUntil(t, 10*time.Second, 5*time.Millisecond, func() bool {
+ list, err := service.List(context.Background(), invocation.InvocationListFilter{Limit: 10})
+ if err != nil {
+ return false
+ }
+ for _, item := range list.Items {
+ if item.Status == invocation.StatusPendingApproval {
+ pendingID = item.InvocationID
+ return true
+ }
+ }
+ return false
+ })
+ if pendingID == "" {
+ t.Error("timed out waiting for a pending-approval invocation")
+ return
+ }
+ if sink.touched() {
+ t.Errorf("sink touched before approval — approval gating must precede any relay")
+ }
+ if err := service.Approve(context.Background(), pendingID, ""); err != nil {
+ t.Errorf("approve: %v", err)
+ }
+ }()
+
+ resp, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, sink)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusSucceeded {
+ t.Fatalf("expected succeeded, got %s: %s", resp.Status, resp.Error)
+ }
+ if !sink.touched() {
+ t.Fatal("expected the sink to have been touched after approval unblocked execution")
+ }
+}
+
+func TestInvokeStreamingNilSinkMatchesInvoke(t *testing.T) {
+ upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var body map[string]any
+ _ = json.NewDecoder(r.Body).Decode(&body)
+ switch body["method"] {
+ case "initialize":
+ _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": body["id"], "result": map[string]any{"serverInfo": map[string]any{"name": "fake", "version": "0.1.0"}, "capabilities": map[string]any{}}})
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": 1, "result": map[string]any{"content": []map[string]any{{"type": "text", "text": "ok"}}}})
+ default:
+ w.WriteHeader(http.StatusBadRequest)
+ }
+ }))
+ defer upstream.Close()
+
+ service := newTestService(t, config.Config{
+ Defaults: config.DefaultsConfig{RequestTimeoutSeconds: 5},
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+
+ viaInvoke, err := service.Invoke(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}})
+ if err != nil {
+ t.Fatalf("Invoke returned error: %v", err)
+ }
+ viaStreaming, err := service.InvokeStreaming(context.Background(), invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}}, nil)
+ if err != nil {
+ t.Fatalf("InvokeStreaming returned error: %v", err)
+ }
+ if viaInvoke.Status != viaStreaming.Status {
+ t.Fatalf("status mismatch: Invoke=%s InvokeStreaming(nil)=%s", viaInvoke.Status, viaStreaming.Status)
+ }
+ if string(viaInvoke.Result) != string(viaStreaming.Result) {
+ t.Fatalf("result mismatch: Invoke=%s InvokeStreaming(nil)=%s", viaInvoke.Result, viaStreaming.Result)
+ }
+}
+
+// TestInvokeBufferedPersistsFailureAfterRequestContextCancellation mirrors
+// the streaming regression test above for the buffered path: a downstream
+// that disconnects mid-call (canceling the request context) must still get
+// its invocation finalized as failed — not left stuck "executing" — because
+// finalization writes run on their own bounded context, detached from the
+// request.
+func TestInvokeBufferedPersistsFailureAfterRequestContextCancellation(t *testing.T) {
+ callStarted := make(chan struct{})
+ upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var body map[string]any
+ if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+ t.Errorf("decode request: %v", err)
+ http.Error(w, "bad request body", http.StatusBadRequest)
+ return
+ }
+ switch body["method"] {
+ case "initialize":
+ _ = json.NewEncoder(w).Encode(map[string]any{"jsonrpc": "2.0", "id": body["id"], "result": map[string]any{"serverInfo": map[string]any{"name": "fake", "version": "0.1.0"}, "capabilities": map[string]any{}}})
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ close(callStarted)
+ <-r.Context().Done() // hold the call open until the client vanishes
+ default:
+ t.Errorf("unexpected method %q", body["method"])
+ http.Error(w, "unexpected method", http.StatusInternalServerError)
+ }
+ }))
+ defer upstream.Close()
+
+ service := newTestService(t, config.Config{
+ Defaults: config.DefaultsConfig{RequestTimeoutSeconds: 5},
+ Upstreams: []config.UpstreamConfig{{Name: "shortcut", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}},
+ })
+
+ ctx, cancel := context.WithCancel(context.Background())
+ respCh := make(chan invocation.InvocationResponse, 1)
+ errCh := make(chan error, 1)
+ go func() {
+ resp, err := service.Invoke(ctx, invocation.CreateInvocationRequest{Server: "shortcut", Tool: "demo", Input: map[string]any{}})
+ respCh <- resp
+ errCh <- err
+ }()
+ <-callStarted
+ cancel()
+
+ resp := <-respCh
+ if err := <-errCh; err != nil {
+ t.Fatalf("Invoke returned error: %v", err)
+ }
+ if resp.Status != invocation.StatusFailed {
+ t.Fatalf("status = %s, want failed", resp.Status)
+ }
+ stored, err := service.Get(context.Background(), resp.InvocationID)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if stored.Status != invocation.StatusFailed {
+ t.Fatalf("persisted status = %s, want failed (the row must not stay executing after a downstream disconnect)", stored.Status)
+ }
+}
diff --git a/internal/invocation/stream_sink.go b/internal/invocation/stream_sink.go
new file mode 100644
index 00000000..ae227bb2
--- /dev/null
+++ b/internal/invocation/stream_sink.go
@@ -0,0 +1,255 @@
+package invocation
+
+import (
+ "context"
+ "sync"
+ "sync/atomic"
+ "time"
+
+ "github.com/validmind/atryum/internal/mcp"
+)
+
+const (
+ streamAuditQueueCapacity = 512
+ streamAuditWorkerCount = 8
+ streamAuditWriteTimeout = 500 * time.Millisecond
+ streamAuditFlushTimeout = 2 * time.Second
+)
+
+type streamAuditWrite struct {
+ repo eventRepo
+ owner *auditingSink
+ event Event
+}
+
+type streamAuditDispatcher struct {
+ queues []chan streamAuditWrite
+ nextShard atomic.Uint64
+}
+
+func newStreamAuditDispatcher() *streamAuditDispatcher {
+ d := &streamAuditDispatcher{queues: make([]chan streamAuditWrite, streamAuditWorkerCount)}
+ // max(1, ...) guards a future non-multiple edit of the constants: an
+ // unbuffered queue would make enqueue drop nearly every write.
+ perWorkerCapacity := max(1, streamAuditQueueCapacity/streamAuditWorkerCount)
+ for i := range d.queues {
+ d.queues[i] = make(chan streamAuditWrite, perWorkerCapacity)
+ go d.runWorker(d.queues[i])
+ }
+ return d
+}
+
+func (d *streamAuditDispatcher) assignShard() int {
+ // Modulo in uint64 before converting: a wrapped counter converted to a
+ // negative int would panic on the queue index.
+ return int((d.nextShard.Add(1) - 1) % uint64(len(d.queues)))
+}
+
+func (d *streamAuditDispatcher) enqueue(shard int, write streamAuditWrite) bool {
+ select {
+ case d.queues[shard] <- write:
+ return true
+ default:
+ return false
+ }
+}
+
+func (d *streamAuditDispatcher) runWorker(queue <-chan streamAuditWrite) {
+ for write := range queue {
+ writeCtx, cancel := context.WithTimeout(context.Background(), streamAuditWriteTimeout)
+ err := write.repo.Create(writeCtx, write.event)
+ cancel()
+ write.owner.completeAuditWrite(err)
+ }
+}
+
+var sharedStreamAuditDispatcher = newStreamAuditDispatcher()
+
+// auditingSink wraps a caller-supplied mcp.StreamSink so every relayed event
+// and the call's outcome are recorded as invocation_events audit rows,
+// correlating the downstream request (requestID), the upstream server
+// (upstreamName), and this invocation (invocationID) — regardless of
+// whether the downstream write later fails. Writes happen here, in the
+// service layer, not the handler, precisely so they occur even if the
+// handler's write to the agent fails mid-stream.
+//
+// Audit writes run through one process-wide bounded worker pool rather than
+// creating a queue and goroutine for every invocation. A stalled repository
+// therefore cannot create unbounded relay goroutines or delay live delivery.
+type auditingSink struct {
+ inner mcp.StreamSink // may be nil: audit-only, no relay
+ events eventRepo
+ invocationID string
+ requestID *string
+ upstreamName string
+ limits StreamAuditLimits
+ auditShard int
+ persisted atomic.Int64
+ failed atomic.Int64
+ dropped atomic.Int64
+ standaloneDropped atomic.Int64
+ pendingMu sync.Mutex
+ pending int
+ closing bool
+ drained chan struct{}
+ drainOnce sync.Once
+
+ seq int
+ // downstreamErr is set once inner.Event returns an error — the
+ // downstream (agent) connection itself failed, as opposed to a
+ // transport-level or timeout failure from the upstream client.
+ downstreamErr error
+}
+
+func newAuditingSink(inner mcp.StreamSink, events eventRepo, invocationID string, requestID *string, upstreamName string, limits StreamAuditLimits) *auditingSink {
+ a := &auditingSink{
+ inner: inner,
+ events: events,
+ invocationID: invocationID,
+ requestID: requestID,
+ upstreamName: upstreamName,
+ limits: limits,
+ auditShard: sharedStreamAuditDispatcher.assignShard(),
+ drained: make(chan struct{}),
+ }
+ return a
+}
+
+func (a *auditingSink) StreamStarted() {
+ if a.inner != nil {
+ a.inner.StreamStarted()
+ }
+}
+
+func (a *auditingSink) Event(evt mcp.StreamEvent) error {
+ a.seq++
+ a.recordEvent(evt)
+ if a.inner == nil {
+ return nil
+ }
+ if err := a.inner.Event(evt); err != nil {
+ a.downstreamErr = err
+ return err
+ }
+ return nil
+}
+
+func (a *auditingSink) StreamStats(stats mcp.StreamStats) {
+ a.standaloneDropped.Add(stats.StandaloneEventsDropped)
+}
+
+func (a *auditingSink) recordEvent(evt mcp.StreamEvent) {
+ if a.events == nil {
+ return
+ }
+ if a.limits.MaxEvents > 0 && a.seq > a.limits.MaxEvents {
+ return // over the cap: still counted via seq, just not persisted
+ }
+ data := evt.Data
+ truncated := false
+ if a.limits.MaxEventBytes > 0 && len(data) > a.limits.MaxEventBytes {
+ data = data[:a.limits.MaxEventBytes]
+ truncated = true
+ }
+ payload := map[string]any{
+ "seq": a.seq,
+ "upstream_name": a.upstreamName,
+ "downstream_request_id": a.requestID,
+ "server_request": evt.ServerRequest,
+ "bytes": len(evt.Data),
+ // data is a plain string, not embedded raw JSON: a truncated payload
+ // is no longer valid JSON, and a string field survives that safely
+ // (encoding/json escapes it) where json.RawMessage would corrupt
+ // the whole audit row.
+ "data": string(data),
+ "truncated": truncated,
+ }
+ record := Event{
+ InvocationID: a.invocationID,
+ EventType: "invocation.stream_event",
+ Payload: mustJSON(payload),
+ CreatedAt: time.Now().UTC(),
+ }
+ a.pendingMu.Lock()
+ a.pending++
+ a.pendingMu.Unlock()
+ if !sharedStreamAuditDispatcher.enqueue(a.auditShard, streamAuditWrite{repo: a.events, owner: a, event: record}) {
+ a.dropped.Add(1)
+ a.completePending()
+ }
+}
+
+func (a *auditingSink) completeAuditWrite(err error) {
+ if err != nil {
+ a.failed.Add(1)
+ } else {
+ a.persisted.Add(1)
+ }
+ a.completePending()
+}
+
+func (a *auditingSink) completePending() {
+ a.pendingMu.Lock()
+ a.pending--
+ drained := a.closing && a.pending == 0
+ a.pendingMu.Unlock()
+ if drained {
+ a.drainOnce.Do(func() { close(a.drained) })
+ }
+}
+
+func (a *auditingSink) waitForAuditWrites() bool {
+ a.pendingMu.Lock()
+ a.closing = true
+ drained := a.pending == 0
+ a.pendingMu.Unlock()
+ if drained {
+ a.drainOnce.Do(func() { close(a.drained) })
+ }
+
+ timer := time.NewTimer(streamAuditFlushTimeout)
+ defer timer.Stop()
+ select {
+ case <-a.drained:
+ return true
+ case <-timer.C:
+ return false
+ }
+}
+
+// finish records the invocation.stream_completed totals row. terminal is
+// "succeeded", "failed", or "persistence_failed".
+//
+// finish blocks its caller — and therefore the agent's terminal frame, which
+// the handler writes only after InvokeStreaming returns — for up to
+// streamAuditFlushTimeout + streamAuditWriteTimeout when the audit store is
+// stalled. That is deliberate: the flush wait is what makes the persisted/
+// failed/dropped totals truthful, and writing stream_completed synchronously
+// here keeps it ordered before the invocation-level failed/succeeded event
+// (see finishExecutionStreaming's narrative-order comment). With a healthy
+// store the cost is a few milliseconds; with a stalled store the invocation
+// is already paying UpdateResult's own 5s bound, so the added tail is
+// accepted rather than trading away audit ordering.
+func (a *auditingSink) finish(completed time.Time, terminal string) {
+ if a.events == nil {
+ return
+ }
+ auditFlushed := a.waitForAuditWrites()
+ payload := map[string]any{
+ "events_total": a.seq,
+ "events_persisted": a.persisted.Load(),
+ "audit_write_failures": a.failed.Load(),
+ "audit_queue_dropped": a.dropped.Load(),
+ "standalone_events_dropped": a.standaloneDropped.Load(),
+ "audit_flush_timed_out": !auditFlushed,
+ "terminal": terminal,
+ }
+ writeCtx, cancel := context.WithTimeout(context.Background(), streamAuditWriteTimeout)
+ defer cancel()
+ _ = a.events.Create(writeCtx, Event{
+ InvocationID: a.invocationID,
+ EventType: "invocation.stream_completed",
+ Payload: mustJSON(payload),
+ CreatedAt: completed,
+ })
+}
diff --git a/internal/invocation/stream_sink_internal_test.go b/internal/invocation/stream_sink_internal_test.go
new file mode 100644
index 00000000..203d540b
--- /dev/null
+++ b/internal/invocation/stream_sink_internal_test.go
@@ -0,0 +1,180 @@
+package invocation
+
+import (
+ "context"
+ "encoding/json"
+ "runtime"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/validmind/atryum/internal/mcp"
+)
+
+type memoryStreamEventRepo struct {
+ mu sync.Mutex
+ events []Event
+}
+
+type orderedStreamEventRepo struct {
+ firstStarted chan struct{}
+ releaseFirst chan struct{}
+ secondStarted chan struct{}
+ firstOnce sync.Once
+ secondOnce sync.Once
+}
+
+func (r *orderedStreamEventRepo) Create(ctx context.Context, event Event) error {
+ if event.EventType != "invocation.stream_event" {
+ return nil
+ }
+ var payload struct {
+ Seq int `json:"seq"`
+ }
+ if err := json.Unmarshal(event.Payload, &payload); err != nil {
+ return err
+ }
+ switch payload.Seq {
+ case 1:
+ r.firstOnce.Do(func() { close(r.firstStarted) })
+ select {
+ case <-r.releaseFirst:
+ case <-ctx.Done():
+ return ctx.Err()
+ }
+ case 2:
+ r.secondOnce.Do(func() { close(r.secondStarted) })
+ }
+ return nil
+}
+
+func (r *orderedStreamEventRepo) ListByInvocation(context.Context, string, EventListFilter) ([]Event, int, error) {
+ return nil, 0, nil
+}
+
+func (r *memoryStreamEventRepo) Create(_ context.Context, event Event) error {
+ r.mu.Lock()
+ r.events = append(r.events, event)
+ r.mu.Unlock()
+ return nil
+}
+
+func (r *memoryStreamEventRepo) ListByInvocation(context.Context, string, EventListFilter) ([]Event, int, error) {
+ return nil, 0, nil
+}
+
+func (r *memoryStreamEventRepo) snapshot() []Event {
+ r.mu.Lock()
+ defer r.mu.Unlock()
+ return append([]Event(nil), r.events...)
+}
+
+func TestAuditingSinkCreationDoesNotStartPerInvocationGoroutine(t *testing.T) {
+ repo := &memoryStreamEventRepo{}
+ before := runtime.NumGoroutine()
+ sinks := make([]*auditingSink, 1000)
+ for i := range sinks {
+ sinks[i] = newAuditingSink(nil, repo, "inv", nil, "upstream", StreamAuditLimits{})
+ }
+ runtime.Gosched()
+ time.Sleep(20 * time.Millisecond)
+ after := runtime.NumGoroutine()
+ t.Logf("goroutine growth after creating 1000 sinks: %d", after-before)
+
+ if growth := after - before; growth > 16 {
+ t.Fatalf("creating 1000 audit sinks added %d goroutines; want service-level bounded workers", growth)
+ }
+ for _, sink := range sinks {
+ sink.finish(time.Now().UTC(), "succeeded")
+ }
+}
+
+func TestSharedAuditDispatcherHandlesOneThousandConcurrentSinks(t *testing.T) {
+ repo := &memoryStreamEventRepo{}
+ var wg sync.WaitGroup
+ wg.Add(1000)
+ for i := 0; i < 1000; i++ {
+ go func() {
+ defer wg.Done()
+ sink := newAuditingSink(nil, repo, "inv-load", nil, "upstream", StreamAuditLimits{MaxEvents: 1})
+ _ = sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress"}`)})
+ sink.finish(time.Now().UTC(), "succeeded")
+ }()
+ }
+ wg.Wait()
+
+ completed := 0
+ for _, event := range repo.snapshot() {
+ if event.EventType == "invocation.stream_completed" {
+ completed++
+ }
+ }
+ if completed != 1000 {
+ t.Fatalf("stream completion rows = %d, want 1000", completed)
+ }
+}
+
+func TestAuditingSinkIncludesStandaloneDropsInCompletion(t *testing.T) {
+ repo := &memoryStreamEventRepo{}
+ sink := newAuditingSink(nil, repo, "inv-drops", nil, "upstream", StreamAuditLimits{})
+ sink.StreamStats(mcp.StreamStats{StandaloneEventsDropped: 7})
+ sink.finish(time.Now().UTC(), "succeeded")
+
+ for _, event := range repo.snapshot() {
+ if event.EventType != "invocation.stream_completed" {
+ continue
+ }
+ var payload struct {
+ StandaloneEventsDropped int64 `json:"standalone_events_dropped"`
+ }
+ if err := json.Unmarshal(event.Payload, &payload); err != nil {
+ t.Fatal(err)
+ }
+ if payload.StandaloneEventsDropped != 7 {
+ t.Fatalf("standalone_events_dropped = %d, want 7", payload.StandaloneEventsDropped)
+ }
+ return
+ }
+ t.Fatal("missing invocation.stream_completed event")
+}
+
+func TestAuditingSinkPreservesEventOrderThroughSharedWorkers(t *testing.T) {
+ repo := &orderedStreamEventRepo{
+ firstStarted: make(chan struct{}),
+ releaseFirst: make(chan struct{}),
+ secondStarted: make(chan struct{}),
+ }
+ sink := newAuditingSink(nil, repo, "inv-order", nil, "upstream", StreamAuditLimits{})
+ _ = sink.Event(mcp.StreamEvent{Data: []byte(`{"progress":1}`)})
+ _ = sink.Event(mcp.StreamEvent{Data: []byte(`{"progress":2}`)})
+
+ select {
+ case <-repo.firstStarted:
+ case <-time.After(time.Second):
+ t.Fatal("first audit write did not start")
+ }
+ select {
+ case <-repo.secondStarted:
+ t.Fatal("second event started before the first event completed")
+ case <-time.After(30 * time.Millisecond):
+ }
+ close(repo.releaseFirst)
+ sink.finish(time.Now().UTC(), "succeeded")
+ select {
+ case <-repo.secondStarted:
+ default:
+ t.Fatal("second event never ran after the first completed")
+ }
+}
+
+func BenchmarkAuditingSinkDispatch(b *testing.B) {
+ repo := &memoryStreamEventRepo{}
+ b.ReportAllocs()
+ b.RunParallel(func(pb *testing.PB) {
+ for pb.Next() {
+ sink := newAuditingSink(nil, repo, "inv", nil, "upstream", StreamAuditLimits{MaxEvents: 1})
+ _ = sink.Event(mcp.StreamEvent{Data: []byte(`{"jsonrpc":"2.0","method":"notifications/progress"}`)})
+ sink.finish(time.Now().UTC(), "succeeded")
+ }
+ })
+}
diff --git a/internal/mcp/client.go b/internal/mcp/client.go
index c0c79cd2..938af063 100644
--- a/internal/mcp/client.go
+++ b/internal/mcp/client.go
@@ -6,6 +6,7 @@ import (
"context"
"database/sql"
"encoding/json"
+ "errors"
"fmt"
"io"
"log"
@@ -252,6 +253,12 @@ type Client struct {
sessionInitLocks map[string]*sync.Mutex
sessions map[string]string
sessionProtocols map[string]string
+
+ // standaloneStreams holds one shared standalone SSE GET per upstream
+ // session. Including the session ID in the key prevents a renewed session
+ // from reusing the dead connection associated with its predecessor.
+ standaloneMu sync.Mutex
+ standaloneStreams map[standaloneStreamKey]*standaloneStream
}
type InvokeResult struct {
@@ -260,6 +267,20 @@ type InvokeResult struct {
Failed bool
}
+// ErrStreamTimeout marks an error returned by InvokeStream as caused by the
+// header/idle/max-duration bound in StreamOptions firing, as opposed to a
+// transport failure or the sink itself returning an error. Callers can
+// distinguish it with errors.Is(err, ErrStreamTimeout).
+var ErrStreamTimeout = errors.New("stream timeout")
+
+// ErrStreamSessionRetryRefused marks an error returned by InvokeStream as
+// caused by an upstream reporting a missing/expired session after events
+// had already been relayed for this call. The normal missing-session
+// recovery (reinitialize and retry) is only safe when nothing has reached
+// the sink yet, since a retry after that point would relay a second copy
+// of everything already delivered — so InvokeStream fails instead.
+var ErrStreamSessionRetryRefused = errors.New("stream session retry refused: events already relayed")
+
type ForwardResult struct {
StatusCode int
Body []byte
@@ -310,7 +331,7 @@ func (r *Resolver) WithCredentials(credentials CredentialStore) *Resolver {
func NewHTTPClient() *Client {
debug := strings.EqualFold(os.Getenv("ATRYUM_MCP_DEBUG"), "1") || strings.EqualFold(os.Getenv("ATRYUM_MCP_DEBUG"), "true")
- return &Client{httpClient: &http.Client{}, debug: debug, sessionInitLocks: make(map[string]*sync.Mutex), sessions: make(map[string]string), sessionProtocols: make(map[string]string)}
+ return &Client{httpClient: &http.Client{}, debug: debug, sessionInitLocks: make(map[string]*sync.Mutex), sessions: make(map[string]string), sessionProtocols: make(map[string]string), standaloneStreams: make(map[standaloneStreamKey]*standaloneStream)}
}
func (r *Resolver) Resolve(name string) (Upstream, error) {
@@ -408,21 +429,45 @@ func (r *Resolver) BootstrapIfEmpty(ctx context.Context) error {
return nil
}
-func (c *Client) Invoke(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string) (InvokeResult, error) {
+// Invoke calls tool on upstream. meta carries the agent's own MCP
+// params._meta (e.g. progressToken) so upstreams that support progress
+// notifications can correlate them back to the agent's original request.
+func (c *Client) Invoke(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string, meta map[string]any) (InvokeResult, error) {
started := time.Now()
defer func() {
c.debugf("upstream invoke transport=%s server=%s tool=%s duration_ms=%d", upstream.Mode, upstream.Name, tool, time.Since(started).Milliseconds())
}()
switch upstream.Mode {
case UpstreamModeStdio:
- return c.invokeStdio(ctx, upstream, tool, input)
+ return c.invokeStdio(ctx, upstream, tool, input, requestID, meta)
case UpstreamModeHTTP, "":
- return c.invokeHTTP(ctx, upstream, tool, input, requestID)
+ return c.invokeHTTP(ctx, upstream, tool, input, requestID, meta)
default:
return InvokeResult{}, fmt.Errorf("unsupported upstream mode %q", upstream.Mode)
}
}
+// mergeRequestMeta merges the agent-supplied params._meta with the
+// atryumRequestId Atryum injects for its own audit correlation. Caller keys
+// (e.g. progressToken) are preserved; atryumRequestId always reflects the
+// current request, overriding any caller-supplied value under that key.
+// Returns nil when there is nothing to send, so callers can omit `_meta`
+// entirely rather than sending `{}`.
+func mergeRequestMeta(meta map[string]any, requestID *string) map[string]any {
+ hasRequestID := requestID != nil && *requestID != ""
+ if len(meta) == 0 && !hasRequestID {
+ return nil
+ }
+ merged := make(map[string]any, len(meta)+1)
+ for k, v := range meta {
+ merged[k] = v
+ }
+ if hasRequestID {
+ merged["atryumRequestId"] = *requestID
+ }
+ return merged
+}
+
func (c *Client) ListTools(ctx context.Context, upstream Upstream) ([]Tool, error) {
started := time.Now()
defer func() {
@@ -588,15 +633,65 @@ func (c *Client) TestConnection(ctx context.Context, upstream Upstream) Connecti
return result
}
-func (c *Client) invokeHTTP(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string) (InvokeResult, error) {
+// toolCallEnvelopeID is the fixed JSON-RPC id every Atryum-built tools/call
+// envelope carries. Everything that correlates a terminal response back to
+// the request — buffered decode and the streaming relays alike — must match
+// against this same value.
+var toolCallEnvelopeID = json.RawMessage("1")
+
+// marshalToolCallEnvelope builds the JSON-RPC tools/call request body Atryum
+// sends upstream. The envelope always uses toolCallEnvelopeID — Atryum's own
+// request id, if any, travels only in the caller-facing InvocationResponse,
+// not on the wire to the upstream.
+func marshalToolCallEnvelope(tool string, input map[string]any, requestID *string, meta map[string]any) ([]byte, error) {
+ return marshalToolCallEnvelopeWithMeta(tool, input, mergeRequestMeta(meta, requestID))
+}
+
+// marshalToolCallEnvelopeWithMeta builds the tools/call request body from an
+// already-fully-merged _meta map (see mergeRequestMeta), skipping that merge
+// step. Used by the streaming path, which may need to rewrite a caller's
+// progressToken after merging but before marshaling (see rewriteProgressToken).
+func marshalToolCallEnvelopeWithMeta(tool string, input map[string]any, meta map[string]any) ([]byte, error) {
+ params := map[string]any{"name": tool, "arguments": input}
+ if meta != nil {
+ params["_meta"] = meta
+ }
+ return json.Marshal(Envelope{JSONRPC: "2.0", ID: toolCallEnvelopeID, Method: "tools/call", Params: mustRawJSON(params)})
+}
+
+// toolCallResultFromRPCResponse maps an already-decoded tools/call JSON-RPC
+// response to the InvokeResult contract, applying the "ok" fallback body
+// when the upstream returns an empty/null result. The second return value
+// flags a missing-session RPC error so the caller can decide whether to
+// reinitialize and retry.
+func toolCallResultFromRPCResponse(rpcResp rpcResponse, statusCode int) (InvokeResult, bool) {
+ if len(rpcResp.Error) > 0 && string(rpcResp.Error) != "null" {
+ return InvokeResult{StatusCode: statusCode, Body: rpcResp.Error, Failed: true}, isMissingSessionRPCError(rpcResp.Error)
+ }
+ bodyBytes := rpcResp.Result
+ if len(bodyBytes) == 0 || string(bodyBytes) == "null" {
+ bodyBytes = []byte(`{"content":[{"type":"text","text":"ok"}]}`)
+ }
+ failed := statusCode >= http.StatusBadRequest || looksLikeToolError(bodyBytes)
+ return InvokeResult{StatusCode: statusCode, Body: bodyBytes, Failed: failed}, false
+}
+
+// toolCallResultFromForward decodes a raw tools/call ForwardResult (JSON or
+// SSE-wrapped) and maps it via toolCallResultFromRPCResponse.
+func toolCallResultFromForward(result ForwardResult) (InvokeResult, bool, error) {
+ rpcResp, err := decodeRPCResponse(result, toolCallEnvelopeID)
+ if err != nil {
+ return InvokeResult{}, false, err
+ }
+ invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, result.StatusCode)
+ return invoke, missingSession, nil
+}
+
+func (c *Client) invokeHTTP(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string, meta map[string]any) (InvokeResult, error) {
if err := c.ensureHTTPSession(ctx, upstream); err != nil {
return InvokeResult{}, err
}
- params := map[string]any{"name": tool, "arguments": input}
- if requestID != nil && *requestID != "" {
- params["_meta"] = map[string]any{"atryumRequestId": *requestID}
- }
- body, err := json.Marshal(Envelope{JSONRPC: "2.0", ID: json.RawMessage([]byte("1")), Method: "tools/call", Params: mustRawJSON(params)})
+ body, err := marshalToolCallEnvelope(tool, input, requestID, meta)
if err != nil {
return InvokeResult{}, err
}
@@ -604,40 +699,29 @@ func (c *Client) invokeHTTP(ctx context.Context, upstream Upstream, tool string,
if err != nil {
return InvokeResult{}, err
}
- rpcResp, err := decodeRPCResponse(result, json.RawMessage([]byte("1")))
+ invoke, missingSession, err := toolCallResultFromForward(result)
if err != nil {
return InvokeResult{}, err
}
- if len(rpcResp.Error) > 0 && string(rpcResp.Error) != "null" {
- if isMissingSessionRPCError(rpcResp.Error) {
- c.debugf("upstream http tools.call missing session server=%s session=%q status=%d error=%s", upstream.Name, result.SessionID, result.StatusCode, truncateForLog(rpcResp.Error, 600))
- if retryErr := c.reinitializeRequiredHTTPSession(ctx, upstream, result.SessionID); retryErr != nil {
- c.debugf("upstream http tools.call session retry init failed server=%s session=%q err=%v", upstream.Name, result.SessionID, retryErr)
- return InvokeResult{}, retryErr
- }
- result, err = c.doHTTPEnvelopeWithSessionRetry(ctx, upstream, body, DefaultMCPProtocolVersion)
- if err != nil {
- c.debugf("upstream http tools.call session retry transport failed server=%s err=%v", upstream.Name, err)
- return InvokeResult{}, err
- }
- rpcResp, err = decodeRPCResponse(result, json.RawMessage([]byte("1")))
- if err != nil {
- c.debugf("upstream http tools.call session retry decode failed server=%s status=%d err=%v", upstream.Name, result.StatusCode, err)
- return InvokeResult{}, err
- }
+ if missingSession {
+ c.debugf("upstream http tools.call missing session server=%s session=%q status=%d error=%s", upstream.Name, result.SessionID, result.StatusCode, truncateForLog(invoke.Body, 600))
+ if retryErr := c.reinitializeRequiredHTTPSession(ctx, upstream, result.SessionID); retryErr != nil {
+ c.debugf("upstream http tools.call session retry init failed server=%s session=%q err=%v", upstream.Name, result.SessionID, retryErr)
+ return InvokeResult{}, retryErr
+ }
+ result, err = c.doHTTPEnvelopeWithSessionRetry(ctx, upstream, body, DefaultMCPProtocolVersion)
+ if err != nil {
+ c.debugf("upstream http tools.call session retry transport failed server=%s err=%v", upstream.Name, err)
+ return InvokeResult{}, err
+ }
+ invoke, _, err = toolCallResultFromForward(result)
+ if err != nil {
+ c.debugf("upstream http tools.call session retry decode failed server=%s status=%d err=%v", upstream.Name, result.StatusCode, err)
+ return InvokeResult{}, err
}
}
- if len(rpcResp.Error) > 0 && string(rpcResp.Error) != "null" {
- c.debugf("upstream http tools.call rpc error server=%s status=%d error=%s", upstream.Name, result.StatusCode, truncateForLog(rpcResp.Error, 600))
- return InvokeResult{StatusCode: result.StatusCode, Body: rpcResp.Error, Failed: true}, nil
- }
- bodyBytes := rpcResp.Result
- if len(bodyBytes) == 0 || string(bodyBytes) == "null" {
- bodyBytes = []byte(`{"content":[{"type":"text","text":"ok"}]}`)
- }
- failed := result.StatusCode >= http.StatusBadRequest || looksLikeToolError(bodyBytes)
- c.debugf("upstream http tools.call server=%s status=%d failed=%t", upstream.Name, result.StatusCode, failed)
- return InvokeResult{StatusCode: result.StatusCode, Body: bodyBytes, Failed: failed}, nil
+ c.debugf("upstream http tools.call server=%s status=%d failed=%t", upstream.Name, result.StatusCode, invoke.Failed)
+ return invoke, nil
}
func (c *Client) listToolsHTTP(ctx context.Context, upstream Upstream) ([]Tool, error) {
@@ -692,10 +776,16 @@ func (c *Client) forwardEnvelopeHTTP(ctx context.Context, upstream Upstream, env
return c.doHTTPEnvelopeWithSessionRetry(ctx, upstream, body, protocolVersion)
}
-func (c *Client) doHTTPEnvelopeRaw(ctx context.Context, upstream Upstream, body []byte, protocolVersion, sessionID string) (*http.Response, error) {
+// doHTTPEnvelopeRaw sends one JSON-RPC envelope and returns the raw HTTP
+// response (headers received, body not yet consumed). streaming disables
+// the per-call upstream.Timeout wrapper: a streaming call's timing is
+// governed entirely by the header/idle/max-duration scheme in
+// invokeHTTPStream, not by a single fixed wall-clock budget that would
+// include however long the stream stays open.
+func (c *Client) doHTTPEnvelopeRaw(ctx context.Context, upstream Upstream, body []byte, protocolVersion, sessionID string, streaming bool) (*http.Response, error) {
endpoint := strings.TrimRight(upstream.BaseURL, "/")
client := c.httpClient
- if upstream.Timeout > 0 {
+ if !streaming && upstream.Timeout > 0 {
client = &http.Client{Timeout: upstream.Timeout}
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
@@ -714,7 +804,21 @@ func (c *Client) doHTTPEnvelopeRaw(ctx context.Context, upstream Upstream, body
return client.Do(req)
}
-func (c *Client) doHTTPEnvelope(ctx context.Context, upstream Upstream, body []byte, protocolVersion string) (ForwardResult, error) {
+// httpEnvelopeHeaders is an upstream HTTP response with session bookkeeping
+// already applied from its headers, but the body NOT yet consumed. Callers
+// own resp.Body: read (or discard) it and close it themselves. This split
+// exists so a streaming caller can inspect Content-Type and take over the
+// body incrementally instead of buffering it, while the ordinary buffered
+// path (doHTTPEnvelope) just reads it in one shot.
+type httpEnvelopeHeaders struct {
+ resp *http.Response
+ sessionID string
+ sessionExpired bool
+ protocolVersion string
+ contentType string
+}
+
+func (c *Client) doHTTPEnvelopeHeaders(ctx context.Context, upstream Upstream, body []byte, protocolVersion string, streaming bool) (httpEnvelopeHeaders, error) {
sessionID := c.getSession(upstream.Name)
effectiveProtocol := protocolVersion
if sessionID != "" {
@@ -722,11 +826,10 @@ func (c *Client) doHTTPEnvelope(ctx context.Context, upstream Upstream, body []b
effectiveProtocol = sessionProtocol
}
}
- resp, err := c.doHTTPEnvelopeRaw(ctx, upstream, body, effectiveProtocol, sessionID)
+ resp, err := c.doHTTPEnvelopeRaw(ctx, upstream, body, effectiveProtocol, sessionID, streaming)
if err != nil {
- return ForwardResult{}, err
+ return httpEnvelopeHeaders{}, err
}
- defer resp.Body.Close()
sessionExpired := false
// Capture/clear session id based on response. 404 with an existing
// session means the server forgot us; drop it so the next call inits.
@@ -736,17 +839,30 @@ func (c *Client) doHTTPEnvelope(ctx context.Context, upstream Upstream, body []b
c.clearSessionIfCurrent(upstream.Name, sessionID)
sessionExpired = true
}
- contentType := resp.Header.Get("Content-Type")
- if sessionExpired {
- bodyBytes, _ := io.ReadAll(resp.Body)
- return ForwardResult{StatusCode: resp.StatusCode, Body: bodyBytes, ContentType: contentType, ProtocolVersion: resp.Header.Get("MCP-Protocol-Version"), SessionExpired: true, SessionID: sessionID}, nil
+ return httpEnvelopeHeaders{
+ resp: resp,
+ sessionID: sessionID,
+ sessionExpired: sessionExpired,
+ protocolVersion: resp.Header.Get("MCP-Protocol-Version"),
+ contentType: resp.Header.Get("Content-Type"),
+ }, nil
+}
+
+func (c *Client) doHTTPEnvelope(ctx context.Context, upstream Upstream, body []byte, protocolVersion string) (ForwardResult, error) {
+ h, err := c.doHTTPEnvelopeHeaders(ctx, upstream, body, protocolVersion, false)
+ if err != nil {
+ return ForwardResult{}, err
+ }
+ defer h.resp.Body.Close()
+ if h.sessionExpired {
+ bodyBytes, _ := io.ReadAll(h.resp.Body)
+ return ForwardResult{StatusCode: h.resp.StatusCode, Body: bodyBytes, ContentType: h.contentType, ProtocolVersion: h.protocolVersion, SessionExpired: true, SessionID: h.sessionID}, nil
}
respBody := new(bytes.Buffer)
- _, err = respBody.ReadFrom(resp.Body)
- if err != nil {
+ if _, err := respBody.ReadFrom(h.resp.Body); err != nil {
return ForwardResult{}, err
}
- return ForwardResult{StatusCode: resp.StatusCode, Body: respBody.Bytes(), ContentType: contentType, ProtocolVersion: resp.Header.Get("MCP-Protocol-Version"), SessionID: sessionID}, nil
+ return ForwardResult{StatusCode: h.resp.StatusCode, Body: respBody.Bytes(), ContentType: h.contentType, ProtocolVersion: h.protocolVersion, SessionID: h.sessionID}, nil
}
func (c *Client) doHTTPEnvelopeWithSessionRetry(ctx context.Context, upstream Upstream, body []byte, protocolVersion string) (ForwardResult, error) {
@@ -972,7 +1088,65 @@ func extractForwardResultErrorDetail(result ForwardResult, expectedID json.RawMe
return extractErrorDetail(result.Body)
}
-func (c *Client) invokeStdio(ctx context.Context, upstream Upstream, tool string, input map[string]any) (InvokeResult, error) {
+// stdioStderrCap bounds how much of a stdio subprocess's stderr is
+// retained in memory for error diagnostics. Streaming calls can run far
+// longer than the old fixed request_timeout_seconds bound (up to
+// stream_max_duration_seconds, or unlimited if unset), so a verbose or
+// misbehaving upstream writing continuously to stderr for the life of a
+// long relay could otherwise grow this buffer without limit.
+const stdioStderrCap = 64 * 1024
+
+// boundedBuffer caps how many bytes Write retains, keeping the first
+// stdioStderrCap bytes and silently discarding the rest — matching this
+// file's existing truncateForLog convention (keep the head, not the tail)
+// for bounding diagnostic text. Write always reports success for the full
+// input, including the discarded portion: the subprocess's stderr pipe
+// must never see a short write or an error from this side.
+//
+// The mutex is required, not defensive: os/exec copies the stderr pipe into
+// this buffer on its own goroutine, which only stops once cmd.Wait reaps it —
+// and the error paths call Len/String before their deferred Wait runs, so a
+// subprocess still writing stderr as it dies races those reads.
+type boundedBuffer struct {
+ mu sync.Mutex
+ buf bytes.Buffer
+ limit int
+}
+
+func newBoundedBuffer(limit int) *boundedBuffer {
+ return &boundedBuffer{limit: limit}
+}
+
+func (b *boundedBuffer) Write(p []byte) (int, error) {
+ b.mu.Lock()
+ defer b.mu.Unlock()
+ remaining := b.limit - b.buf.Len()
+ if remaining <= 0 {
+ return len(p), nil
+ }
+ keep := p
+ if len(keep) > remaining {
+ keep = keep[:remaining]
+ }
+ if _, err := b.buf.Write(keep); err != nil {
+ return 0, err
+ }
+ return len(p), nil
+}
+
+func (b *boundedBuffer) Len() int {
+ b.mu.Lock()
+ defer b.mu.Unlock()
+ return b.buf.Len()
+}
+
+func (b *boundedBuffer) String() string {
+ b.mu.Lock()
+ defer b.mu.Unlock()
+ return b.buf.String()
+}
+
+func (c *Client) invokeStdio(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string, meta map[string]any) (InvokeResult, error) {
if upstream.Command == "" {
return InvokeResult{}, fmt.Errorf("stdio upstream %q missing command", upstream.Name)
}
@@ -981,6 +1155,7 @@ func (c *Client) invokeStdio(ctx context.Context, upstream Upstream, tool string
for k, v := range upstream.Env {
cmd.Env = append(cmd.Env, k+"="+v)
}
+ configureStdioProcessGroup(cmd)
stdin, err := cmd.StdinPipe()
if err != nil {
return InvokeResult{}, err
@@ -989,7 +1164,7 @@ func (c *Client) invokeStdio(ctx context.Context, upstream Upstream, tool string
if err != nil {
return InvokeResult{}, err
}
- stderr := new(bytes.Buffer)
+ stderr := newBoundedBuffer(stdioStderrCap)
cmd.Stderr = stderr
if err := cmd.Start(); err != nil {
return InvokeResult{}, err
@@ -1000,24 +1175,27 @@ func (c *Client) invokeStdio(ctx context.Context, upstream Upstream, tool string
}()
reader := bufio.NewReader(stdout)
- if err := writeRPC(stdin, c.nextRPCID(), "initialize", map[string]any{
+ initID := c.nextRPCID()
+ if err := writeRPC(stdin, initID, "initialize", map[string]any{
"protocolVersion": DefaultMCPProtocolVersion,
"clientInfo": map[string]any{"name": "atryum", "version": version.Version},
"capabilities": map[string]any{},
}); err != nil {
return InvokeResult{}, err
}
- if _, err := readRPC(reader); err != nil {
+ if _, err := readRPC(reader, rpcIDMessage(initID)); err != nil {
return InvokeResult{}, err
}
_ = writeRPC(stdin, c.nextRPCID(), "notifications/initialized", map[string]any{})
- if err := writeRPC(stdin, c.nextRPCID(), "tools/call", map[string]any{
- "name": tool,
- "arguments": input,
- }); err != nil {
+ callParams := map[string]any{"name": tool, "arguments": input}
+ if merged := mergeRequestMeta(meta, requestID); merged != nil {
+ callParams["_meta"] = merged
+ }
+ callID := c.nextRPCID()
+ if err := writeRPC(stdin, callID, "tools/call", callParams); err != nil {
return InvokeResult{}, err
}
- resp, err := readRPC(reader)
+ resp, err := readRPC(reader, rpcIDMessage(callID))
if err != nil {
if stderr.Len() > 0 {
return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String()))
@@ -1043,6 +1221,7 @@ func (c *Client) listToolsStdio(ctx context.Context, upstream Upstream) ([]Tool,
for k, v := range upstream.Env {
cmd.Env = append(cmd.Env, k+"="+v)
}
+ configureStdioProcessGroup(cmd)
stdin, err := cmd.StdinPipe()
if err != nil {
return nil, err
@@ -1051,7 +1230,7 @@ func (c *Client) listToolsStdio(ctx context.Context, upstream Upstream) ([]Tool,
if err != nil {
return nil, err
}
- stderr := new(bytes.Buffer)
+ stderr := newBoundedBuffer(stdioStderrCap)
cmd.Stderr = stderr
if err := cmd.Start(); err != nil {
return nil, err
@@ -1061,21 +1240,23 @@ func (c *Client) listToolsStdio(ctx context.Context, upstream Upstream) ([]Tool,
_ = cmd.Wait()
}()
reader := bufio.NewReader(stdout)
- if err := writeRPC(stdin, c.nextRPCID(), "initialize", map[string]any{
+ initID := c.nextRPCID()
+ if err := writeRPC(stdin, initID, "initialize", map[string]any{
"protocolVersion": DefaultMCPProtocolVersion,
"clientInfo": map[string]any{"name": "atryum", "version": version.Version},
"capabilities": map[string]any{},
}); err != nil {
return nil, err
}
- if _, err := readRPC(reader); err != nil {
+ if _, err := readRPC(reader, rpcIDMessage(initID)); err != nil {
return nil, err
}
_ = writeRPC(stdin, c.nextRPCID(), "notifications/initialized", map[string]any{})
- if err := writeRPC(stdin, c.nextRPCID(), "tools/list", map[string]any{}); err != nil {
+ listID := c.nextRPCID()
+ if err := writeRPC(stdin, listID, "tools/list", map[string]any{}); err != nil {
return nil, err
}
- resp, err := readRPC(reader)
+ resp, err := readRPC(reader, rpcIDMessage(listID))
if err != nil {
if stderr.Len() > 0 {
return nil, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String()))
@@ -1250,6 +1431,7 @@ func (c *Client) testStdio(ctx context.Context, upstream Upstream) ConnectionTes
for k, v := range upstream.Env {
cmd.Env = append(cmd.Env, k+"="+v)
}
+ configureStdioProcessGroup(cmd)
stdin, err := cmd.StdinPipe()
if err != nil {
message := err.Error()
@@ -1260,7 +1442,7 @@ func (c *Client) testStdio(ctx context.Context, upstream Upstream) ConnectionTes
message := err.Error()
return ConnectionTestResult{Ok: false, Message: message, ConnectionStatus: ConnectionStatusUnreachable, AuthStatus: AuthStatusUnknown, LastCheckOK: false, LastErrorSummary: &message}
}
- stderr := new(bytes.Buffer)
+ stderr := newBoundedBuffer(stdioStderrCap)
cmd.Stderr = stderr
if err := cmd.Start(); err != nil {
message := err.Error()
@@ -1271,7 +1453,8 @@ func (c *Client) testStdio(ctx context.Context, upstream Upstream) ConnectionTes
_ = cmd.Wait()
}()
reader := bufio.NewReader(stdout)
- if err := writeRPC(stdin, c.nextRPCID(), "initialize", map[string]any{
+ initID := c.nextRPCID()
+ if err := writeRPC(stdin, initID, "initialize", map[string]any{
"protocolVersion": DefaultMCPProtocolVersion,
"clientInfo": map[string]any{"name": "atryum", "version": version.Version},
"capabilities": map[string]any{},
@@ -1279,7 +1462,7 @@ func (c *Client) testStdio(ctx context.Context, upstream Upstream) ConnectionTes
message := err.Error()
return ConnectionTestResult{Ok: false, Message: message, ConnectionStatus: ConnectionStatusUnreachable, AuthStatus: AuthStatusUnknown, LastCheckOK: false, LastErrorSummary: &message}
}
- if _, err := readRPC(reader); err != nil {
+ if _, err := readRPC(reader, rpcIDMessage(initID)); err != nil {
message := err.Error()
if stderr.Len() > 0 {
message = strings.TrimSpace(stderr.String())
@@ -1291,53 +1474,6 @@ func (c *Client) testStdio(ctx context.Context, upstream Upstream) ConnectionTes
return ConnectionTestResult{Ok: true, Message: "stdio initialize ok", ConnectionStatus: ConnectionStatusReady, AuthStatus: AuthStatusReady, ReauthNeeded: false, LastCheckOK: true}
}
-func extractSSEJSONRPCResponse(r io.Reader, expectedID json.RawMessage) ([]byte, error) {
- scanner := bufio.NewScanner(r)
- scanner.Buffer(make([]byte, 1024*1024), 4*1024*1024)
- var dataLines []string
- flush := func() ([]byte, bool) {
- if len(dataLines) == 0 {
- return nil, false
- }
- payload := []byte(strings.Join(dataLines, "\n"))
- dataLines = nil
- match := classifyJSONRPCResponsePayload(payload, expectedID)
- if match == jsonRPCResponseIDMatch || match == jsonRPCResponseNullIDError {
- return payload, true
- }
- return nil, false
- }
- for scanner.Scan() {
- line := scanner.Text()
- if line == "" {
- if payload, ok := flush(); ok {
- return payload, nil
- }
- continue
- }
- if strings.HasPrefix(line, ":") {
- continue
- }
- field, value, ok := strings.Cut(line, ":")
- if !ok {
- field = line
- value = ""
- } else if strings.HasPrefix(value, " ") {
- value = strings.TrimPrefix(value, " ")
- }
- if field == "data" {
- dataLines = append(dataLines, value)
- }
- }
- if err := scanner.Err(); err != nil {
- return nil, err
- }
- if payload, ok := flush(); ok {
- return payload, nil
- }
- return nil, fmt.Errorf("no JSON-RPC response in SSE stream")
-}
-
type jsonRPCResponseMatch int
const (
@@ -1544,9 +1680,20 @@ func writeRPC(w interface{ Write([]byte) (int, error) }, id int64, method string
return err
}
-func readRPC(reader *bufio.Reader) (rpcResponse, error) {
+// readRPC reads newline-delimited JSON-RPC messages from reader until it
+// finds the response whose id matches expectedID, skipping everything else
+// (notifications, stray server-to-client requests, responses to some other
+// id). This replaces a historical "first message with any id/result/error
+// wins" read: without correlation, a genuine incoming server-to-client
+// request (which also carries an id) could be misread as the answer to our
+// own call, since it was indistinguishable from a response by that check.
+func readRPC(reader *bufio.Reader, expectedID json.RawMessage) (rpcResponse, error) {
+ return readRPCWithLimit(reader, expectedID, defaultStreamMaxMessageBytes)
+}
+
+func readRPCWithLimit(reader *bufio.Reader, expectedID json.RawMessage, maxMessageBytes int) (rpcResponse, error) {
for {
- line, err := reader.ReadBytes('\n')
+ line, err := readLineLimited(reader, maxMessageBytes)
if err != nil {
return rpcResponse{}, err
}
@@ -1554,17 +1701,37 @@ func readRPC(reader *bufio.Reader) (rpcResponse, error) {
if len(line) == 0 {
continue
}
- var resp rpcResponse
- if err := json.Unmarshal(line, &resp); err != nil {
+ if classifyRPCMessage(line, expectedID) != rpcMessageTerminalResponse {
continue
}
- if len(resp.ID) == 0 && len(resp.Result) == 0 && len(resp.Error) == 0 {
+ var resp rpcResponse
+ if err := json.Unmarshal(line, &resp); err != nil {
continue
}
return resp, nil
}
}
+func readLineLimited(reader *bufio.Reader, maxBytes int) ([]byte, error) {
+ if maxBytes <= 0 {
+ maxBytes = defaultStreamMaxMessageBytes
+ }
+ line := make([]byte, 0, min(maxBytes, 64*1024))
+ for {
+ fragment, err := reader.ReadSlice('\n')
+ if len(fragment) > maxBytes-len(line) {
+ return nil, ErrStreamMessageTooLarge
+ }
+ line = append(line, fragment...)
+ if err == nil {
+ return line, nil
+ }
+ if err != bufio.ErrBufferFull {
+ return nil, err
+ }
+ }
+}
+
func looksLikeToolError(body []byte) bool {
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
@@ -1608,6 +1775,12 @@ func (c *Client) nextRPCID() int64 {
return c.nextID.Add(1)
}
+// rpcIDMessage renders a stdio request id (an int64 from nextRPCID) as the
+// json.RawMessage form readRPC/classifyRPCMessage compare ids against.
+func rpcIDMessage(id int64) json.RawMessage {
+ return json.RawMessage(strconv.FormatInt(id, 10))
+}
+
func cloneMap(in map[string]string) map[string]string {
if len(in) == 0 {
return nil
diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go
index 422c1a63..25df1af7 100644
--- a/internal/mcp/client_test.go
+++ b/internal/mcp/client_test.go
@@ -467,7 +467,7 @@ func TestInvokeSkipsSSENotificationBeforeResponse(t *testing.T) {
client := NewHTTPClient()
client.httpClient = server.Client()
- result, err := client.Invoke(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil)
+ result, err := client.Invoke(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil)
if err != nil {
t.Fatalf("Invoke returned error: %v", err)
}
@@ -476,7 +476,12 @@ func TestInvokeSkipsSSENotificationBeforeResponse(t *testing.T) {
}
}
-func TestListToolsDecodesMultilineSSEData(t *testing.T) {
+func TestInvokeForwardsCallerMetaAndAtryumRequestID(t *testing.T) {
+ var capturedParams struct {
+ Name string `json:"name"`
+ Arguments map[string]any `json:"arguments"`
+ Meta map[string]any `json:"_meta"`
+ }
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var req Envelope
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
@@ -484,16 +489,15 @@ func TestListToolsDecodesMultilineSSEData(t *testing.T) {
}
switch req.Method {
case "initialize":
- w.Header().Set("Mcp-Session-Id", "sid-multiline")
+ w.Header().Set("Mcp-Session-Id", "sid-meta")
writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
case "notifications/initialized":
w.WriteHeader(http.StatusAccepted)
- case "tools/list":
- writeTestSSEEvents(w, []string{
- `{"jsonrpc":"2.0",`,
- `"id":1,`,
- `"result":{"tools":[{"name":"stories.multiline"}]}}`,
- })
+ case "tools/call":
+ if err := json.Unmarshal(req.Params, &capturedParams); err != nil {
+ t.Fatalf("decode tools/call params: %v", err)
+ }
+ writeTestRPC(w, req.ID, map[string]any{"content": []any{map[string]any{"type": "text", "text": "ok"}}}, nil)
default:
t.Fatalf("unexpected method %q", req.Method)
}
@@ -502,12 +506,104 @@ func TestListToolsDecodesMultilineSSEData(t *testing.T) {
client := NewHTTPClient()
client.httpClient = server.Client()
- tools, err := client.ListTools(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL})
- if err != nil {
- t.Fatalf("ListTools returned error: %v", err)
+ requestID := "req-42"
+ meta := map[string]any{"progressToken": "tok-7"}
+ if _, err := client.Invoke(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, &requestID, meta); err != nil {
+ t.Fatalf("Invoke returned error: %v", err)
}
- if len(tools) != 1 || tools[0].Name != "stories.multiline" {
- t.Fatalf("unexpected tools: %#v", tools)
+
+ if got := capturedParams.Meta["progressToken"]; got != "tok-7" {
+ t.Fatalf("expected caller progressToken preserved in upstream _meta, got %#v", capturedParams.Meta)
+ }
+ if got := capturedParams.Meta["atryumRequestId"]; got != "req-42" {
+ t.Fatalf("expected atryumRequestId injected into upstream _meta, got %#v", capturedParams.Meta)
+ }
+}
+
+func TestInvokeOmitsMetaWhenCallerMetaAndRequestIDAreEmpty(t *testing.T) {
+ var sawMeta bool
+ var rawParams json.RawMessage
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-nometa")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ rawParams = req.Params
+ var decoded map[string]any
+ if err := json.Unmarshal(req.Params, &decoded); err != nil {
+ t.Fatalf("decode tools/call params: %v", err)
+ }
+ _, sawMeta = decoded["_meta"]
+ writeTestRPC(w, req.ID, map[string]any{"content": []any{map[string]any{"type": "text", "text": "ok"}}}, nil)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ if _, err := client.Invoke(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil); err != nil {
+ t.Fatalf("Invoke returned error: %v", err)
+ }
+
+ if sawMeta {
+ t.Fatalf("expected no _meta field when caller meta and requestID are both empty, got params %s", rawParams)
+ }
+}
+
+func TestBoundedBufferCapsRetainedBytesWithoutErroringWrites(t *testing.T) {
+ b := newBoundedBuffer(8)
+ n, err := b.Write([]byte("1234"))
+ if err != nil || n != 4 {
+ t.Fatalf("Write(1234) = %d, %v, want 4, nil", n, err)
+ }
+ n, err = b.Write([]byte("567890")) // 4 + 6 = 10, exceeds the 8-byte cap
+ if err != nil || n != 6 {
+ t.Fatalf("Write(567890) = %d, %v, want 6, nil (full length reported even though truncated)", n, err)
+ }
+ if got := b.String(); got != "12345678" {
+ t.Fatalf("String() = %q, want the first 8 bytes \"12345678\"", got)
+ }
+ if b.Len() != 8 {
+ t.Fatalf("Len() = %d, want 8", b.Len())
+ }
+ // Further writes past the cap must still report full success (the
+ // subprocess's stderr pipe must never see a short write or an error).
+ n, err = b.Write([]byte("more data"))
+ if err != nil || n != len("more data") {
+ t.Fatalf("Write past cap = %d, %v, want %d, nil", n, err, len("more data"))
+ }
+ if b.Len() != 8 {
+ t.Fatalf("Len() after writing past cap = %d, want 8 (unchanged)", b.Len())
+ }
+}
+
+func TestMergeRequestMeta(t *testing.T) {
+ reqID := "req-1"
+
+ if got := mergeRequestMeta(nil, nil); got != nil {
+ t.Fatalf("expected nil for empty meta and nil requestID, got %#v", got)
+ }
+ empty := ""
+ if got := mergeRequestMeta(nil, &empty); got != nil {
+ t.Fatalf("expected nil for empty meta and empty requestID, got %#v", got)
+ }
+ if got := mergeRequestMeta(map[string]any{"progressToken": "tok"}, nil); got["progressToken"] != "tok" || got["atryumRequestId"] != nil {
+ t.Fatalf("expected caller meta preserved without atryumRequestId, got %#v", got)
+ }
+ if got := mergeRequestMeta(nil, &reqID); got["atryumRequestId"] != "req-1" {
+ t.Fatalf("expected atryumRequestId-only meta, got %#v", got)
+ }
+ if got := mergeRequestMeta(map[string]any{"progressToken": "tok", "atryumRequestId": "spoofed"}, &reqID); got["progressToken"] != "tok" || got["atryumRequestId"] != "req-1" {
+ t.Fatalf("expected atryumRequestId to win over a caller-supplied value, got %#v", got)
}
}
@@ -821,3 +917,34 @@ func writeTestSSEEvents(w http.ResponseWriter, events ...[]string) {
_, _ = w.Write([]byte("\n"))
}
}
+
+// writeTestSSEEventFlush writes and flushes one SSE event immediately, so a
+// test server can hold a stream open between events (unlike writeTestSSEEvents,
+// which writes every event in one shot with no flush in between).
+func writeTestSSEEventFlush(w http.ResponseWriter, flusher http.Flusher, data string) {
+ _, _ = w.Write([]byte("event: message\ndata: " + data + "\n\n"))
+ flusher.Flush()
+}
+
+// fakeStreamSink is a test StreamSink that records what it received. onEvent,
+// when set, lets a test hook into delivery (e.g. to unblock a fake upstream
+// only after confirming an event was actually delivered incrementally).
+type fakeStreamSink struct {
+ started bool
+ startedCount int
+ events []StreamEvent
+ onEvent func(StreamEvent) error
+}
+
+func (s *fakeStreamSink) StreamStarted() {
+ s.started = true
+ s.startedCount++
+}
+
+func (s *fakeStreamSink) Event(evt StreamEvent) error {
+ s.events = append(s.events, evt)
+ if s.onEvent != nil {
+ return s.onEvent(evt)
+ }
+ return nil
+}
diff --git a/internal/mcp/http_stream.go b/internal/mcp/http_stream.go
new file mode 100644
index 00000000..cdb41c63
--- /dev/null
+++ b/internal/mcp/http_stream.go
@@ -0,0 +1,584 @@
+package mcp
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "io"
+ "math/rand/v2"
+ "net/http"
+ "strings"
+ "sync"
+ "sync/atomic"
+ "time"
+)
+
+// This file merges two independent sources of upstream messages into one
+// ordered sequence for a single tools/call, while enforcing timeouts and
+// supporting reconnection. Read it in this order:
+//
+// 1. invokeHTTPStream, the entry point. It mints a per-call progress token
+// (rewriteProgressToken, in standalone_stream.go) so the shared
+// standalone stream can attribute progress to this specific call, then
+// calls doHTTPToolCallStream and retries once on a missing-session
+// response.
+// 2. doHTTPToolCallStream sends the request. A plain JSON response is
+// handled right there and never reaches step 3; an SSE response is
+// handed to relaySSEToolCall.
+// 3. relaySSEToolCall is the merge point: one select loop reading two
+// channels — postStreamPump's msgs (this call's own POST response,
+// including any resume after a disconnect) and progressCh (this call's
+// slice of the standalone stream, routed by standalone_stream.go's
+// routeStandaloneEvent). Every sink call happens on this one goroutine,
+// so the sink itself never needs its own locking.
+// 4. postStreamPump owns the POST response's read loop on its own
+// goroutine, purely so relaySSEToolCall's select never blocks on it. It
+// is the only place that resumes a disconnected stream.
+//
+// The three time limits live in callTimeoutGuard (stream_timeout.go). The
+// stdio equivalent of this whole file is stdio_stream.go, and is much
+// simpler: one reader, no shared connection, no resume.
+
+// streamCallOutcome is the result of one attempt to send a streaming
+// tools/call request. missingSession mirrors doHTTPEnvelope's
+// SessionExpired signal so invokeHTTPStream can apply the same
+// reinitialize-and-retry-once policy invokeHTTP uses — but only when
+// eventsRelayed is 0: once anything has reached the sink, the downstream
+// has already seen stream bytes, so retrying would relay a second copy of
+// everything. In that case the caller fails instead of retrying.
+type streamCallOutcome struct {
+ invoke InvokeResult
+ missingSession bool
+ sessionID string
+ eventsRelayed int
+}
+
+// doHTTPToolCallStream sends one tools/call request and either reads a
+// plain JSON body (mapped exactly like the buffered path) or, for an SSE
+// response, relays intermediate events to sink live and returns once the
+// terminal response is read.
+func (c *Client) doHTTPToolCallStream(ctx context.Context, upstream Upstream, body []byte, sink StreamSink, progressCh <-chan StreamEvent, opts StreamOptions) (streamCallOutcome, error) {
+ guard := newCallTimeoutGuard(ctx)
+ defer guard.stop()
+ guard.armSetupTimeout(opts.HeaderTimeout)
+
+ h, err := c.doHTTPEnvelopeHeaders(guard.ctx, upstream, body, DefaultMCPProtocolVersion, true)
+ guard.disarmSetupTimeout()
+ if err != nil {
+ return streamCallOutcome{}, guard.timeoutErr(upstream.Name, "", err)
+ }
+ resp := h.resp
+
+ // Armed as soon as headers are back, before any body is read — the
+ // header timeout only ever bounded waiting for headers, so every
+ // body-reading branch below (including the two early returns, not just
+ // the SSE relay) needs its own bound. Without this, a slow/hanging body
+ // on a 404-session-expired or plain-JSON response during a streaming
+ // call would be unbounded: the per-call http.Client.Timeout that would
+ // normally catch this is deliberately skipped in streaming mode (see
+ // doHTTPEnvelopeRaw's streaming param).
+ guard.armBodyTimeouts(opts.IdleTimeout, opts.MaxDuration)
+
+ if h.sessionExpired {
+ defer resp.Body.Close()
+ // Drain (bounded) so the connection can be reused. A session-expired
+ // body is a small error payload; cap it rather than trust the
+ // upstream, since the guard's timers are the only other bound here
+ // and both are disableable by configuration.
+ _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 64*1024))
+ return streamCallOutcome{missingSession: true, sessionID: h.sessionID}, nil
+ }
+
+ if !strings.Contains(strings.ToLower(h.contentType), "text/event-stream") {
+ defer resp.Body.Close()
+ bodyBytes, err := readAllLimited(resp.Body, opts.maxMessageBytes())
+ if err != nil {
+ return streamCallOutcome{}, guard.timeoutErr(upstream.Name, "", err)
+ }
+ forward := ForwardResult{StatusCode: resp.StatusCode, Body: bodyBytes, ContentType: h.contentType, ProtocolVersion: h.protocolVersion, SessionID: h.sessionID}
+ invoke, missingSession, err := toolCallResultFromForward(forward)
+ if err != nil {
+ return streamCallOutcome{}, err
+ }
+ return streamCallOutcome{invoke: invoke, missingSession: missingSession, sessionID: h.sessionID}, nil
+ }
+
+ // relaySSEToolCall owns resp.Body because it may replace this response
+ // with one or more resumed GET streams before the terminal response.
+ return c.relaySSEToolCall(resp, sink, progressCh, guard, upstream, h.sessionID, opts.maxMessageBytes())
+}
+
+func readAllLimited(r io.Reader, maxBytes int) ([]byte, error) {
+ if maxBytes <= 0 {
+ maxBytes = defaultStreamMaxMessageBytes
+ }
+ const maxInt64 = int64(^uint64(0) >> 1)
+ limit := int64(maxBytes)
+ if limit < maxInt64 {
+ limit++
+ }
+ body, err := io.ReadAll(io.LimitReader(r, limit))
+ if err != nil {
+ return nil, err
+ }
+ if len(body) > maxBytes {
+ return nil, ErrStreamMessageTooLarge
+ }
+ return body, nil
+}
+
+// postStreamMsg is one message pumped from a tools/call POST response by
+// postStreamPump: either a data-bearing JSON-RPC payload (data != nil), or a
+// terminal error ending the stream (err != nil).
+type postStreamMsg struct {
+ data []byte
+ err error
+}
+
+// postStreamPump owns the tools/call POST response's read loop — including
+// SSE resumption — on its own goroutine, feeding relaySSEToolCall with only
+// the data-bearing JSON-RPC payloads (or a final error) through msgs. This
+// lets relaySSEToolCall select between this stream and a per-call
+// standalone-stream channel (progressCh) without either blocking the
+// other, so a call is only ever done reading (and only ever returns to its
+// caller) once both are accounted for — see progressWaiter for why that
+// matters.
+type postStreamPump struct {
+ msgs chan postStreamMsg
+
+ mu sync.Mutex
+ current *http.Response
+ stopped bool
+ stopOnce sync.Once
+ done chan struct{}
+}
+
+func newPostStreamPump(c *Client, guard *callTimeoutGuard, upstream Upstream, resp *http.Response, maxMessageBytes int) *postStreamPump {
+ p := &postStreamPump{msgs: make(chan postStreamMsg), current: resp, done: make(chan struct{})}
+ go p.run(c, guard, upstream, maxMessageBytes)
+ return p
+}
+
+// stop closes the currently-active response body, if any — causing a
+// blocked Read to return promptly — and marks the pump stopped so it exits
+// instead of trying to resume. Safe to call more than once; only the first
+// call has any effect. Always safe to call even if the pump has already
+// finished on its own.
+func (p *postStreamPump) stop() {
+ p.stopOnce.Do(func() {
+ p.mu.Lock()
+ p.stopped = true
+ cur := p.current
+ p.mu.Unlock()
+ close(p.done)
+ if cur != nil {
+ _ = cur.Body.Close()
+ }
+ })
+}
+
+// setCurrent installs resp as the response the pump is currently reading from
+// after a resume and closes the response it supersedes. Returns false — and
+// leaves resp to the caller to close — if stop was already called, so a resume
+// racing a stop can't resurrect a pump that's supposed to be shutting down.
+func (p *postStreamPump) setCurrent(resp *http.Response) bool {
+ p.mu.Lock()
+ if p.stopped {
+ p.mu.Unlock()
+ return false
+ }
+ previous := p.current
+ p.current = resp
+ p.mu.Unlock()
+
+ // Close outside p.mu: a custom response body may perform blocking cleanup,
+ // and stop must remain able to acquire the lock and close the active body.
+ if previous != nil && previous != resp {
+ _ = previous.Body.Close()
+ }
+ return true
+}
+
+// send delivers msg, or exits early if stop is called while blocked trying
+// to (msgs is unbuffered: without this, a caller that stops reading msgs
+// after its own terminal response — see relaySSEToolCall — would otherwise
+// leave this goroutine permanently blocked on a send nobody will ever
+// receive).
+func (p *postStreamPump) send(msg postStreamMsg) {
+ select {
+ case p.msgs <- msg:
+ case <-p.done:
+ }
+}
+
+const (
+ minSSEReconnectDelay = 200 * time.Millisecond
+ maxSSEReconnectDelay = 30 * time.Second
+)
+
+func sseReconnectJitter(base time.Duration) time.Duration {
+ return base / 5
+}
+
+func sseReconnectDelay(serverDelay time.Duration, attempt int) time.Duration {
+ base := serverDelay
+ if base <= 0 {
+ base = minSSEReconnectDelay
+ for range min(attempt, 8) {
+ if base >= maxSSEReconnectDelay/2 {
+ base = maxSSEReconnectDelay
+ break
+ }
+ base *= 2
+ }
+ }
+ base = max(base, minSSEReconnectDelay)
+ base = min(base, maxSSEReconnectDelay)
+ jitter := sseReconnectJitter(base)
+ if jitter == 0 || base == maxSSEReconnectDelay {
+ return base
+ }
+ return min(base+time.Duration(rand.Int64N(int64(jitter)+1)), maxSSEReconnectDelay)
+}
+
+func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstream, maxMessageBytes int) {
+ defer close(p.msgs)
+ // Per-line activity (not per-event) resets the idle timer, so an
+ // upstream keepalive comment during a long-running tool counts as
+ // liveness. This subsumes a per-event reset: every event is lines.
+ reader := newSSEEventReaderWithLimit(p.current.Body, maxMessageBytes)
+ reader.onActivity = guard.resetIdle
+ lastEventID := ""
+ retryDelay := time.Duration(0)
+ reconnectAttempt := 0
+ // resumedFrom holds, after a resume, the cursor id the Last-Event-ID
+ // header carried. Replay semantics are exclusive of the cursor, but the
+ // classic server off-by-one replays it inclusively — without this guard
+ // the cursor event's data would be relayed to the agent a second time.
+ // The guard window closes at the first event bearing any other id, so a
+ // server legitimately reusing the id much later is unaffected.
+ resumedFrom := ""
+ for {
+ evt, err := reader.NextEvent()
+ if err != nil {
+ p.mu.Lock()
+ stopped := p.stopped
+ p.mu.Unlock()
+ if stopped {
+ return
+ }
+ if timeoutErr := guard.timeoutErr(upstream.Name, "", nil); timeoutErr != nil {
+ p.send(postStreamMsg{err: timeoutErr})
+ return
+ }
+ if err != io.EOF {
+ p.send(postStreamMsg{err: err})
+ return
+ }
+ if lastEventID == "" {
+ p.send(postStreamMsg{err: fmt.Errorf("upstream %q closed the stream without a JSON-RPC response or resumable event id", upstream.Name)})
+ return
+ }
+ delay := sseReconnectDelay(retryDelay, reconnectAttempt)
+ reconnectAttempt++
+ if err := waitForSSEReconnect(guard.ctx, delay); err != nil {
+ p.send(postStreamMsg{err: guard.timeoutErr(upstream.Name, "while waiting to resume", err)})
+ return
+ }
+ resumed, err := c.resumeSSEStream(guard.ctx, upstream, lastEventID)
+ if err != nil {
+ p.send(postStreamMsg{err: guard.timeoutErr(upstream.Name, "while resuming", err)})
+ return
+ }
+ if !p.setCurrent(resumed) {
+ _ = resumed.Body.Close()
+ return
+ }
+ reader = newSSEEventReaderWithLimit(resumed.Body, maxMessageBytes)
+ reader.onActivity = guard.resetIdle
+ resumedFrom = lastEventID
+ continue
+ }
+ reconnectAttempt = 0
+ if evt.HasRetry {
+ retryDelay = evt.Retry
+ }
+ if evt.HasID {
+ if resumedFrom != "" && evt.ID == resumedFrom {
+ // Inclusive replay of the cursor event we already relayed
+ // before the disconnect: keep the bookkeeping, skip the data.
+ lastEventID = evt.ID
+ continue
+ }
+ resumedFrom = ""
+ lastEventID = evt.ID
+ }
+ if !evt.HasData {
+ continue
+ }
+ p.send(postStreamMsg{data: evt.Data})
+ }
+}
+
+// relaySSEToolCall selects between the POST response and this call's
+// standalone progress channel, keeping all sink calls on one goroutine. The
+// pump owns the response body, including resumed responses. StreamStarted is
+// withheld only when a zero-event missing-session response will be retried.
+func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progressCh <-chan StreamEvent, guard *callTimeoutGuard, upstream Upstream, sessionID string, maxMessageBytes int) (streamCallOutcome, error) {
+ expectedID := toolCallEnvelopeID
+ statusCode := resp.StatusCode
+ relayed := 0
+ started := false
+ ensureStarted := func() {
+ if !started {
+ started = true
+ sink.StreamStarted()
+ }
+ }
+ deliver := func(evt StreamEvent) error {
+ guard.resetIdle()
+ ensureStarted()
+ relayed++
+ return sink.Event(evt)
+ }
+ // fail builds the outcome every early-return-with-error site below
+ // shares: relayed and sessionID as they stood at the moment of failure,
+ // paired with whatever error caused it.
+ fail := func(err error) (streamCallOutcome, error) {
+ return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err
+ }
+
+ pump := newPostStreamPump(c, guard, upstream, resp, maxMessageBytes)
+ defer pump.stop()
+
+ for {
+ select {
+ case evt, ok := <-progressCh:
+ if !ok {
+ // Never actually closed (its registration outlives this
+ // call — see invokeHTTPStream's grace period), but nil this
+ // out defensively so a closed channel can't busy-loop.
+ progressCh = nil
+ continue
+ }
+ if err := deliver(evt); err != nil {
+ return fail(err)
+ }
+ case msg, ok := <-pump.msgs:
+ if !ok {
+ return fail(fmt.Errorf("upstream %q: stream ended unexpectedly", upstream.Name))
+ }
+ if msg.err != nil {
+ return fail(msg.err)
+ }
+ payload := msg.data
+ switch classifyRPCMessage(payload, expectedID) {
+ case rpcMessageTerminalResponse:
+ var rpcResp rpcResponse
+ if err := json.Unmarshal(payload, &rpcResp); err != nil {
+ return fail(err)
+ }
+ invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, statusCode)
+ if progressCh != nil {
+ if err := drainTrailingProgress(guard.ctx, progressCh, deliver, terminalSettleWindow); err != nil {
+ return fail(guard.timeoutErr(upstream.Name, "while draining trailing progress", err))
+ }
+ }
+ if !(missingSession && relayed == 0) {
+ ensureStarted()
+ }
+ return streamCallOutcome{invoke: invoke, missingSession: missingSession, eventsRelayed: relayed, sessionID: sessionID}, nil
+ case rpcMessageServerRequest:
+ if err := deliver(StreamEvent{Data: payload, ServerRequest: true}); err != nil {
+ return fail(err)
+ }
+ case rpcMessageNotification:
+ if err := deliver(StreamEvent{Data: payload}); err != nil {
+ return fail(err)
+ }
+ default:
+ // Unrecognized payload shape (e.g. a response to some other id).
+ // Not ours to interpret; ignore and keep reading.
+ }
+ }
+ }
+}
+
+// drainTrailingProgress gives standalone-stream notifications already in
+// flight a brief, absolutely bounded chance to arrive before the terminal
+// response is finalized. Arrivals do not extend the window, and cancellation
+// ends the drain immediately: an upstream that keeps sending after its terminal
+// response must not keep the invocation open indefinitely.
+func drainTrailingProgress(ctx context.Context, progressCh <-chan StreamEvent, deliver func(StreamEvent) error, window time.Duration) error {
+ if err := ctx.Err(); err != nil {
+ return err
+ }
+ settle := time.NewTimer(window)
+ defer settle.Stop()
+ for {
+ select {
+ case <-ctx.Done():
+ return ctx.Err()
+ case evt, ok := <-progressCh:
+ if !ok {
+ return nil
+ }
+ if err := ctx.Err(); err != nil {
+ return err
+ }
+ if err := deliver(evt); err != nil {
+ return err
+ }
+ case <-settle.C:
+ return nil
+ }
+ }
+}
+
+func waitForSSEReconnect(ctx context.Context, delay time.Duration) error {
+ if delay <= 0 {
+ return nil
+ }
+ timer := time.NewTimer(delay)
+ defer timer.Stop()
+ select {
+ case <-timer.C:
+ return nil
+ case <-ctx.Done():
+ return ctx.Err()
+ }
+}
+
+// resumeSSEStream continues a server-closed Streamable HTTP response. The
+// MCP transport specifies a GET to the same endpoint carrying Last-Event-ID;
+// session, protocol, and authentication headers must match the original
+// connection so the upstream can locate the pending request.
+func (c *Client) resumeSSEStream(ctx context.Context, upstream Upstream, lastEventID string) (*http.Response, error) {
+ endpoint := strings.TrimRight(upstream.BaseURL, "/")
+ req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
+ if err != nil {
+ return nil, err
+ }
+ req.Header.Set("Accept", "text/event-stream")
+ req.Header.Set("Last-Event-ID", lastEventID)
+ if protocol := c.getSessionProtocol(upstream.Name); protocol != "" {
+ req.Header.Set("MCP-Protocol-Version", protocol)
+ }
+ if sessionID := c.getSession(upstream.Name); sessionID != "" {
+ req.Header.Set("Mcp-Session-Id", sessionID)
+ }
+ applyAuthHeaders(req, upstream)
+ resp, err := c.httpClient.Do(req)
+ if err != nil {
+ return nil, err
+ }
+ if resp.StatusCode >= http.StatusBadRequest {
+ defer resp.Body.Close()
+ body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
+ return nil, fmt.Errorf("upstream %q resume failed with HTTP %d: %s", upstream.Name, resp.StatusCode, extractErrorDetail(body))
+ }
+ if !strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
+ defer resp.Body.Close()
+ return nil, fmt.Errorf("upstream %q resume returned content type %q, want text/event-stream", upstream.Name, resp.Header.Get("Content-Type"))
+ }
+ if newSession := strings.TrimSpace(resp.Header.Get("Mcp-Session-Id")); newSession != "" {
+ protocol := c.getSessionProtocol(upstream.Name)
+ c.setSession(upstream.Name, newSession, protocol)
+ }
+ return resp, nil
+}
+
+// runSessionInitBounded runs fn (a session initialize/reinitialize step)
+// under opts.HeaderTimeout. The session-init POSTs happen before the
+// streaming call proper, so doHTTPToolCallStream's own header timeout never
+// covers them; and in streaming mode the caller's ctx carries no deadline
+// (the fixed request timeout is deliberately not applied — that's the whole
+// point of the streaming timeout scheme). Without this bound, an upstream
+// with no per-server timeout_seconds configured that hangs during
+// initialize would block the call indefinitely.
+func runSessionInitBounded(ctx context.Context, upstream Upstream, opts StreamOptions, fn func(context.Context) error) error {
+ initCtx := ctx
+ if opts.HeaderTimeout > 0 {
+ var cancel context.CancelFunc
+ initCtx, cancel = context.WithTimeout(ctx, opts.HeaderTimeout)
+ defer cancel()
+ }
+ err := fn(initCtx)
+ if err != nil && initCtx.Err() != nil && ctx.Err() == nil {
+ // The bound we imposed fired (not the caller's own ctx): surface it
+ // as the same typed timeout the rest of the streaming path uses.
+ return fmt.Errorf("upstream %q: timed out initializing session before streaming call: %w", upstream.Name, ErrStreamTimeout)
+ }
+ return err
+}
+
+func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string, meta map[string]any, sink StreamSink, opts StreamOptions) (InvokeResult, error) {
+ if err := runSessionInitBounded(ctx, upstream, opts, func(initCtx context.Context) error {
+ return c.ensureHTTPSession(initCtx, upstream)
+ }); err != nil {
+ return InvokeResult{}, err
+ }
+
+ merged := mergeRequestMeta(meta, requestID)
+ // A callSink restores the caller's progress token for messages arriving
+ // from either the POST response or the standalone stream.
+ effectiveSink := sink
+ var progressCh chan StreamEvent
+ var standalone *standaloneStream
+ var standaloneDropped atomic.Int64
+ var wireToken string
+ if rewritten, token, original, ok := c.rewriteProgressToken(merged); ok {
+ wireToken = token
+ merged = rewritten
+ effectiveSink = newCallSink(sink, wireToken, original)
+ progressCh = make(chan StreamEvent, standaloneWaiterEventBuffer)
+ standalone = c.acquireStandaloneStreamWithLimit(upstream, opts.maxMessageBytes())
+ standalone.registerWaiter(wireToken, progressWaiter{events: progressCh, dropped: &standaloneDropped})
+ waitForStandaloneReady(ctx, standalone, opts.HeaderTimeout)
+ defer func() {
+ current := standalone
+ current.unregisterWaiter(wireToken)
+ c.releaseStandaloneStream(current)
+ if statsSink, ok := effectiveSink.(StreamStatsSink); ok {
+ statsSink.StreamStats(StreamStats{StandaloneEventsDropped: standaloneDropped.Load()})
+ }
+ }()
+ }
+
+ body, err := marshalToolCallEnvelopeWithMeta(tool, input, merged)
+ if err != nil {
+ return InvokeResult{}, err
+ }
+
+ outcome, err := c.doHTTPToolCallStream(ctx, upstream, body, effectiveSink, progressCh, opts)
+ if err != nil {
+ return InvokeResult{}, err
+ }
+ if outcome.missingSession {
+ if outcome.eventsRelayed > 0 {
+ return InvokeResult{}, fmt.Errorf("upstream %q reported a missing session after the stream had already relayed %d event(s): %w", upstream.Name, outcome.eventsRelayed, ErrStreamSessionRetryRefused)
+ }
+ c.debugf("upstream http tools.call stream missing session server=%s session=%q", upstream.Name, outcome.sessionID)
+ if retryErr := runSessionInitBounded(ctx, upstream, opts, func(initCtx context.Context) error {
+ return c.reinitializeRequiredHTTPSession(initCtx, upstream, outcome.sessionID)
+ }); retryErr != nil {
+ return InvokeResult{}, retryErr
+ }
+ if standalone != nil {
+ standalone.unregisterWaiter(wireToken)
+ c.releaseStandaloneStream(standalone)
+ standalone = c.acquireStandaloneStreamWithLimit(upstream, opts.maxMessageBytes())
+ standalone.registerWaiter(wireToken, progressWaiter{events: progressCh, dropped: &standaloneDropped})
+ waitForStandaloneReady(ctx, standalone, opts.HeaderTimeout)
+ }
+ outcome, err = c.doHTTPToolCallStream(ctx, upstream, body, effectiveSink, progressCh, opts)
+ if err != nil {
+ return InvokeResult{}, err
+ }
+ if outcome.missingSession {
+ return InvokeResult{}, fmt.Errorf("upstream %q rejected session after reinitialize", upstream.Name)
+ }
+ }
+ c.debugf("upstream http tools.call stream server=%s status=%d failed=%t events_relayed=%d", upstream.Name, outcome.invoke.StatusCode, outcome.invoke.Failed, outcome.eventsRelayed)
+ return outcome.invoke, nil
+}
diff --git a/internal/mcp/http_stream_test.go b/internal/mcp/http_stream_test.go
new file mode 100644
index 00000000..b8b43bdf
--- /dev/null
+++ b/internal/mcp/http_stream_test.go
@@ -0,0 +1,1031 @@
+package mcp
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "io"
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+)
+
+// invokeStreamTestServer builds the initialize/notifications.initialized
+// scaffolding shared by the InvokeStream tests below, dispatching tools/call
+// to callHandler. Failures inside the handler use t.Errorf plus an HTTP
+// error, never t.Fatal: FailNow from a non-test goroutine only kills the
+// handler, and the never-answered client would hang the test until the
+// suite timeout instead of failing it cleanly.
+func invokeStreamTestServer(t *testing.T, sessionID string, callHandler func(w http.ResponseWriter, r *http.Request, req Envelope)) *httptest.Server {
+ t.Helper()
+ return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ // The standalone SSE stream a progressToken-bearing call opens
+ // alongside its tools/call POST. This fake upstream doesn't
+ // support it — a legitimate, spec-allowed response — so tests
+ // using a progressToken don't need every callHandler to be
+ // GET-aware.
+ http.Error(w, "not found", http.StatusNotFound)
+ return
+ }
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Errorf("decode request: %v", err)
+ http.Error(w, "bad request body", http.StatusBadRequest)
+ return
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", sessionID)
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ callHandler(w, r, req)
+ default:
+ t.Errorf("unexpected method %q", req.Method)
+ http.Error(w, "unexpected method", http.StatusInternalServerError)
+ }
+ }))
+}
+
+func TestInvokeStreamRelaysEventsBeforeTerminalResponseExists(t *testing.T) {
+ release := make(chan struct{})
+ server := invokeStreamTestServer(t, "sid-incremental", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":"tok-7","progress":1}}`)
+ <-release // the terminal response cannot be written until the test has observed the event above
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{onEvent: func(StreamEvent) error {
+ close(release)
+ return nil
+ }}
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !sink.started {
+ t.Fatal("expected StreamStarted to fire")
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected exactly one relayed event, got %d", len(sink.events))
+ }
+ if !strings.Contains(string(sink.events[0].Data), "notifications/progress") {
+ t.Fatalf("expected progress notification relayed, got %s", sink.events[0].Data)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+}
+
+func TestInvokeStreamResumesAfterUpstreamClosesSSEBeforeTerminalResponse(t *testing.T) {
+ // Handler goroutines and the test goroutine share this with no
+ // happens-before edge the race detector recognizes; keep it atomic.
+ var resumeRequests atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ resumeRequests.Add(1)
+ if got := r.Header.Get("Last-Event-ID"); got != "evt-1" {
+ t.Fatalf("resume Last-Event-ID = %q, want evt-1", got)
+ }
+ if got := r.Header.Get("Mcp-Session-Id"); got != "sid-resume" {
+ t.Fatalf("resume Mcp-Session-Id = %q, want sid-resume", got)
+ }
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ _, _ = io.WriteString(w, "id: evt-2\ndata: {\"jsonrpc\":\"2.0\",\"id\":\"1\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"done after resume\"}]}}\n\n")
+ flusher.Flush()
+ return
+ }
+
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-resume")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ _, _ = io.WriteString(w, "id: evt-1\nretry: 1\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}\n\n")
+ flusher.Flush()
+ // End this HTTP response without the terminal JSON-RPC response.
+ // A resumable MCP stream continues through a GET with Last-Event-ID.
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+ result, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "stories.get", map[string]any{}, nil, nil, sink,
+ StreamOptions{IdleTimeout: time.Second, MaxDuration: 5 * time.Second},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if got := resumeRequests.Load(); got != 1 {
+ t.Fatalf("resume request count = %d, want 1", got)
+ }
+ if len(sink.events) != 1 || !strings.Contains(string(sink.events[0].Data), "notifications/progress") {
+ t.Fatalf("expected exactly the pre-disconnect progress event, got %#v", sink.events)
+ }
+ if !strings.Contains(string(result.Body), "done after resume") {
+ t.Fatalf("expected terminal response from resumed stream, got %s", result.Body)
+ }
+}
+
+type closeTrackingBody struct {
+ io.Reader
+ closeCount atomic.Int32
+}
+
+func (b *closeTrackingBody) Close() error {
+ b.closeCount.Add(1)
+ return nil
+}
+
+func TestPostStreamPumpClosesSupersededResponseBody(t *testing.T) {
+ previous := &closeTrackingBody{Reader: strings.NewReader("")}
+ current := &closeTrackingBody{Reader: strings.NewReader("")}
+ pump := &postStreamPump{
+ current: &http.Response{Body: previous},
+ done: make(chan struct{}),
+ }
+ t.Cleanup(pump.stop)
+
+ if !pump.setCurrent(&http.Response{Body: current}) {
+ t.Fatal("setCurrent unexpectedly rejected the resumed response")
+ }
+ if got := previous.closeCount.Load(); got != 1 {
+ t.Fatalf("superseded response body close count = %d, want 1", got)
+ }
+ if got := current.closeCount.Load(); got != 0 {
+ t.Fatalf("current response body was closed before the pump stopped: count=%d", got)
+ }
+
+ pump.stop()
+ if got := current.closeCount.Load(); got != 1 {
+ t.Fatalf("current response body close count after stop = %d, want 1", got)
+ }
+}
+
+func TestSSEReconnectDelayUsesBoundedExponentialBackoff(t *testing.T) {
+ first := sseReconnectDelay(0, 0)
+ second := sseReconnectDelay(0, 1)
+
+ if first < minSSEReconnectDelay || first > minSSEReconnectDelay+sseReconnectJitter(minSSEReconnectDelay) {
+ t.Fatalf("first reconnect delay = %s, want bounded minimum-delay jitter", first)
+ }
+ if second < 2*minSSEReconnectDelay || second > 2*minSSEReconnectDelay+sseReconnectJitter(2*minSSEReconnectDelay) {
+ t.Fatalf("second reconnect delay = %s, want bounded exponential-delay jitter", second)
+ }
+ if got := sseReconnectDelay(time.Nanosecond, 0); got < minSSEReconnectDelay {
+ t.Fatalf("server retry below minimum produced %s, want at least %s", got, minSSEReconnectDelay)
+ }
+ if got := sseReconnectDelay(24*time.Hour, 0); got > maxSSEReconnectDelay {
+ t.Fatalf("server retry above maximum produced %s, want at most %s", got, maxSSEReconnectDelay)
+ }
+ for range 100 {
+ if got := sseReconnectDelay(maxSSEReconnectDelay-time.Second, 0); got > maxSSEReconnectDelay {
+ t.Fatalf("jitter pushed reconnect delay above maximum: %s", got)
+ }
+ }
+}
+
+func TestInvokeStreamRejectsOversizedPlainJSONResponse(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-large-json", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ writeTestRPC(w, req.ID, map[string]any{"content": strings.Repeat("x", 256)}, nil)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ _, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "large-json", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "demo",
+ map[string]any{},
+ nil,
+ nil,
+ &fakeStreamSink{},
+ StreamOptions{MaxMessageBytes: 128},
+ )
+ if !errors.Is(err, ErrStreamMessageTooLarge) {
+ t.Fatalf("InvokeStream error = %v, want ErrStreamMessageTooLarge", err)
+ }
+}
+
+// TestInvokeStreamResumeSkipsInclusivelyReplayedCursorEvent is a regression
+// test for reconnect duplicate delivery: Last-Event-ID replay is exclusive
+// of the cursor, but the classic server off-by-one replays the cursor event
+// itself again. That event's data already reached the agent before the
+// disconnect — relaying it twice would deliver a duplicate notification.
+func TestInvokeStreamResumeSkipsInclusivelyReplayedCursorEvent(t *testing.T) {
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ if got := r.Header.Get("Last-Event-ID"); got != "evt-1" {
+ t.Fatalf("resume Last-Event-ID = %q, want evt-1", got)
+ }
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ // Buggy inclusive replay: evt-1 again, then genuinely new events.
+ _, _ = io.WriteString(w, "id: evt-1\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}\n\n")
+ _, _ = io.WriteString(w, "id: evt-2\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":2}}\n\n")
+ _, _ = io.WriteString(w, "id: evt-3\ndata: {\"jsonrpc\":\"2.0\",\"id\":\"1\",\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"done after resume\"}]}}\n\n")
+ flusher.Flush()
+ return
+ }
+
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-resume-dupe")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ _, _ = io.WriteString(w, "id: evt-1\nretry: 1\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}\n\n")
+ flusher.Flush()
+ // Close without the terminal response → client resumes via GET.
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+ result, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "stories.get", map[string]any{}, nil, nil, sink,
+ StreamOptions{IdleTimeout: time.Second, MaxDuration: 5 * time.Second},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if len(sink.events) != 2 {
+ t.Fatalf("expected exactly 2 relayed notifications (evt-1 once + evt-2, cursor replay deduplicated), got %d: %#v", len(sink.events), sink.events)
+ }
+ if !strings.Contains(string(sink.events[0].Data), `"progress":1`) || !strings.Contains(string(sink.events[1].Data), `"progress":2`) {
+ t.Fatalf("expected progress 1 then progress 2, got %#v", sink.events)
+ }
+ if !strings.Contains(string(result.Body), "done after resume") {
+ t.Fatalf("expected terminal response from resumed stream, got %s", result.Body)
+ }
+}
+
+func TestInvokeStreamRelaysNotificationAndServerRequest(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-mixed", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"halfway"}}`)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"srv-1","method":"sampling/createMessage","params":{}}`)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if len(sink.events) != 2 {
+ t.Fatalf("expected 2 relayed events (notification + server request), got %d: %#v", len(sink.events), sink.events)
+ }
+ if sink.events[0].ServerRequest {
+ t.Fatalf("expected first event to be a notification, got %#v", sink.events[0])
+ }
+ if !sink.events[1].ServerRequest {
+ t.Fatalf("expected second event to be flagged as a server request, got %#v", sink.events[1])
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+}
+
+func TestInvokeStreamJSONResponseNeverTouchesSink(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-json", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ writeTestRPC(w, req.ID, map[string]any{"content": []any{map[string]any{"type": "text", "text": "ok"}}}, nil)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if sink.started || len(sink.events) != 0 {
+ t.Fatalf("expected sink to never be touched for a JSON response, got started=%t events=%#v", sink.started, sink.events)
+ }
+ if !strings.Contains(string(result.Body), `"text":"ok"`) {
+ t.Fatalf("expected plain JSON result body, got %s", result.Body)
+ }
+}
+
+func TestInvokeStreamSSETerminalOnlyResponseStartsSink(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-terminal-only-sse", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ writeTestSSEEventFlush(
+ w,
+ w.(http.Flusher),
+ `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`,
+ )
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "terminal-only-sse", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "demo",
+ map[string]any{},
+ nil,
+ nil,
+ sink,
+ StreamOptions{},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if sink.startedCount != 1 {
+ t.Fatalf("expected terminal-only HTTP SSE to start the sink exactly once, got %d", sink.startedCount)
+ }
+ if len(sink.events) != 0 {
+ t.Fatalf("expected no intermediate events, got %d", len(sink.events))
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+}
+
+// TestInvokeStreamHangingJSONBodyBoundedByIdleTimeout is a regression test:
+// StreamOptions' idle/max-duration bounds must apply to every body-reading
+// branch of doHTTPToolCallStream, not just the SSE relay. The per-call
+// http.Client.Timeout that would normally catch a hanging JSON body is
+// deliberately skipped in streaming mode (see doHTTPEnvelopeRaw's streaming
+// param), so without this a slow-to-complete JSON response during a
+// streaming call attempt would hang forever.
+func TestInvokeStreamHangingJSONBodyBoundedByIdleTimeout(t *testing.T) {
+ blockUntilTestDone := make(chan struct{})
+ server := invokeStreamTestServer(t, "sid-slow-json", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "application/json")
+ flusher := w.(http.Flusher)
+ _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"1","result":`)) // deliberately incomplete
+ flusher.Flush()
+ <-blockUntilTestDone // never completes the body
+ })
+ t.Cleanup(func() {
+ close(blockUntilTestDone)
+ server.Close()
+ })
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ start := time.Now()
+ _, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{IdleTimeout: 50 * time.Millisecond})
+ elapsed := time.Since(start)
+
+ if err == nil {
+ t.Fatal("expected a timeout error for the hanging JSON body read")
+ }
+ if !errors.Is(err, ErrStreamTimeout) {
+ t.Fatalf("expected errors.Is(err, ErrStreamTimeout), got %v", err)
+ }
+ if elapsed > 5*time.Second {
+ t.Fatalf("took too long to abort: %s", elapsed)
+ }
+}
+
+// TestInvokeStreamHangingSessionInitBoundedByHeaderTimeout is a regression
+// test: the session-initialize POST happens before doHTTPToolCallStream's
+// own header timeout is armed, and in streaming mode neither the per-call
+// http.Client timeout (deliberately skipped) nor the caller's ctx (no
+// deadline) bounds it. An upstream with no per-server timeout configured
+// that hangs on initialize would block the call forever without
+// runSessionInitBounded.
+func TestInvokeStreamHangingSessionInitBoundedByHeaderTimeout(t *testing.T) {
+ blockUntilTestDone := make(chan struct{})
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ if req.Method != "initialize" {
+ t.Fatalf("unexpected method %q before initialize completed", req.Method)
+ }
+ <-blockUntilTestDone // hang the initialize response forever
+ }))
+ t.Cleanup(func() {
+ close(blockUntilTestDone)
+ server.Close()
+ })
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ start := time.Now()
+ // upstream.Timeout deliberately zero: no per-server bound to fall back on.
+ _, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{HeaderTimeout: 50 * time.Millisecond})
+ elapsed := time.Since(start)
+
+ if err == nil {
+ t.Fatal("expected a session-init timeout error")
+ }
+ if !errors.Is(err, ErrStreamTimeout) {
+ t.Fatalf("expected errors.Is(err, ErrStreamTimeout) for the hung session init, got %v", err)
+ }
+ if elapsed > 5*time.Second {
+ t.Fatalf("session-init timeout took too long to abort: %s", elapsed)
+ }
+ if sink.started || len(sink.events) != 0 {
+ t.Fatalf("expected the sink never to be touched during a failed session init, got started=%t events=%d", sink.started, len(sink.events))
+ }
+}
+
+func TestInvokeStreamMapsTerminalRPCErrorAfterRelayedEvents(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-error", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"tool exploded"}}`)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !result.Failed {
+ t.Fatalf("expected Failed result, got %#v", result)
+ }
+ if !strings.Contains(string(result.Body), "tool exploded") {
+ t.Fatalf("expected error body, got %s", result.Body)
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected the progress notification to have been relayed before the terminal error, got %d", len(sink.events))
+ }
+}
+
+func TestInvokeStreamRetriesOnceWhenMissingSessionBeforeAnyEventRelayed(t *testing.T) {
+ // Handler goroutines and the test goroutine share these with no
+ // happens-before edge the race detector recognizes; guard them.
+ var stateMu sync.Mutex
+ var sessions []string
+ var toolsCallCount int
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ stateMu.Lock()
+ sessionID := "sid-1"
+ if len(sessions) > 0 {
+ sessionID = "sid-2"
+ }
+ sessions = append(sessions, sessionID)
+ stateMu.Unlock()
+ w.Header().Set("Mcp-Session-Id", sessionID)
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ stateMu.Lock()
+ toolsCallCount++
+ call := toolsCallCount
+ stateMu.Unlock()
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ if call == 1 {
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"No session ID provided for non-initialization request"}}`)
+ return
+ }
+ if got := r.Header.Get("Mcp-Session-Id"); got != "sid-2" {
+ t.Fatalf("retry tools/call used session %q, want sid-2", got)
+ }
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ stateMu.Lock()
+ gotCalls, gotSessions := toolsCallCount, len(sessions)
+ stateMu.Unlock()
+ if gotCalls != 2 {
+ t.Fatalf("tools/call count = %d, want 2", gotCalls)
+ }
+ if gotSessions != 2 {
+ t.Fatalf("initialize session count = %d, want 2", gotSessions)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body after retry, got %s", result.Body)
+ }
+ if sink.startedCount != 1 {
+ t.Fatalf("expected StreamStarted to fire exactly once (for the successful retry, not the discarded missing-session attempt), got %d", sink.startedCount)
+ }
+}
+
+func TestInvokeStreamRefusesRetryAfterEventsAlreadyRelayed(t *testing.T) {
+ // Atomics: shared between handler goroutines and the test goroutine.
+ var initializeCount atomic.Int32
+ var toolsCallCount atomic.Int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ initializeCount.Add(1)
+ w.Header().Set("Mcp-Session-Id", "sid-1")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ toolsCallCount.Add(1)
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"No session ID provided for non-initialization request"}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ _, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err == nil {
+ t.Fatal("expected an error refusing to retry mid-stream")
+ }
+ if !strings.Contains(err.Error(), "already relayed") {
+ t.Fatalf("expected a mid-stream retry refusal error, got %v", err)
+ }
+ if !errors.Is(err, ErrStreamSessionRetryRefused) {
+ t.Fatalf("expected errors.Is(err, ErrStreamSessionRetryRefused) to hold, got %v", err)
+ }
+ if got := toolsCallCount.Load(); got != 1 {
+ t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed)", got)
+ }
+ if got := initializeCount.Load(); got != 1 {
+ t.Fatalf("initialize count = %d, want 1 (no reinitialize attempt)", got)
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected the one notification before the terminal error to have been relayed, got %d", len(sink.events))
+ }
+}
+
+func TestInvokeStreamIdleTimeoutAbortsRead(t *testing.T) {
+ blockUntilTestDone := make(chan struct{})
+ server := invokeStreamTestServer(t, "sid-idle", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progress":1}}`)
+ <-blockUntilTestDone // never send the terminal event
+ })
+ t.Cleanup(func() {
+ close(blockUntilTestDone)
+ server.Close()
+ })
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ start := time.Now()
+ _, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, nil, sink, StreamOptions{IdleTimeout: 50 * time.Millisecond})
+ elapsed := time.Since(start)
+
+ if err == nil {
+ t.Fatal("expected an idle timeout error")
+ }
+ if !strings.Contains(err.Error(), "idle timeout") {
+ t.Fatalf("expected an idle timeout error, got %v", err)
+ }
+ if !errors.Is(err, ErrStreamTimeout) {
+ t.Fatalf("expected errors.Is(err, ErrStreamTimeout) to hold so callers can distinguish it from other failures, got %v", err)
+ }
+ if elapsed > 5*time.Second {
+ t.Fatalf("idle timeout took too long to abort: %s", elapsed)
+ }
+ if !sink.started {
+ t.Fatal("expected StreamStarted before the timeout")
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected exactly one relayed event before the timeout, got %d", len(sink.events))
+ }
+}
+
+func TestDrainTrailingProgressDrainsBufferedBurstThenStopsAtClose(t *testing.T) {
+ progressCh := make(chan StreamEvent, 3)
+ progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)}
+ progressCh <- StreamEvent{Data: []byte(`{"progress":2}`)}
+ progressCh <- StreamEvent{Data: []byte(`{"progress":3}`)}
+ close(progressCh)
+
+ var delivered []StreamEvent
+ deliver := func(evt StreamEvent) error {
+ delivered = append(delivered, evt)
+ return nil
+ }
+ // An hour-long window cannot elapse during the test: returning at all
+ // proves the closed channel — not the timer — ended the drain, after the
+ // full buffered burst was delivered.
+ if err := drainTrailingProgress(context.Background(), progressCh, deliver, time.Hour); err != nil {
+ t.Fatalf("drainTrailingProgress returned error: %v", err)
+ }
+ if len(delivered) != 3 {
+ t.Fatalf("delivered %d events, want the full burst of 3", len(delivered))
+ }
+ if !strings.Contains(string(delivered[2].Data), `"progress":3`) {
+ t.Fatalf("expected in-order burst delivery, got %#v", delivered)
+ }
+}
+
+func TestDrainTrailingProgressStopsAtDeliverError(t *testing.T) {
+ progressCh := make(chan StreamEvent, 2)
+ progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)}
+ progressCh <- StreamEvent{Data: []byte(`{"progress":2}`)}
+
+ sinkErr := errors.New("sink rejected event")
+ deliverCalls := 0
+ deliver := func(StreamEvent) error {
+ deliverCalls++
+ return sinkErr
+ }
+ if err := drainTrailingProgress(context.Background(), progressCh, deliver, time.Hour); !errors.Is(err, sinkErr) {
+ t.Fatalf("drainTrailingProgress error = %v, want the deliver error", err)
+ }
+ if deliverCalls != 1 {
+ t.Fatalf("deliver called %d times, want 1 (abort on first failure)", deliverCalls)
+ }
+}
+
+func TestDrainTrailingProgressReturnsOnceWindowElapsesWithNoArrival(t *testing.T) {
+ progressCh := make(chan StreamEvent) // open, never receives anything
+ deliver := func(StreamEvent) error {
+ t.Error("deliver must not be called when nothing arrives")
+ return nil
+ }
+ start := time.Now()
+ if err := drainTrailingProgress(context.Background(), progressCh, deliver, 20*time.Millisecond); err != nil {
+ t.Fatalf("drainTrailingProgress returned error: %v", err)
+ }
+ // Generous bound: only pins that the timer path returns at all rather
+ // than blocking on the open channel forever.
+ if elapsed := time.Since(start); elapsed > 5*time.Second {
+ t.Fatalf("window-elapse return took %s", elapsed)
+ }
+}
+
+// TestDrainTrailingProgressUsesAbsoluteWindowDuringContinuousArrivals proves
+// the trailing-progress window is an absolute deadline, not one that resets
+// on every arrival: a sender goroutine floods progressCh every gap (2ms),
+// far faster than window (30ms), so if an arrival reset the window it would
+// never elapse and drainTrailingProgress would hang.
+//
+// sender goroutine drainTrailingProgress goroutine
+// ----------------- -------------------------------
+// ticker fires every gap (2ms)
+// progressCh <- event ------------> delivered; window (30ms) is
+// progressCh <- event ------------> armed once at the start and
+// progressCh <- event ------------> never extended by an arrival
+// ... (keeps sending until told to stop)
+// window elapses
+// result <- nil
+// |
+// main: case err = <-result: <--------------+ must win the race against
+// close(senderStop); <-senderDone <-time.After(5*window), the
+// test's own bail-out path
+//
+// If drainTrailingProgress instead reset the window on every arrival, result
+// would never fire, the 5*window bail-out would trip, and the test would
+// fail with "continuous progress extended the terminal drain past its
+// absolute window".
+func TestDrainTrailingProgressUsesAbsoluteWindowDuringContinuousArrivals(t *testing.T) {
+ const (
+ window = 30 * time.Millisecond
+ gap = 2 * time.Millisecond
+ )
+
+ progressCh := make(chan StreamEvent, 1)
+ senderStop := make(chan struct{})
+ senderDone := make(chan struct{})
+ go func() {
+ defer close(senderDone)
+ ticker := time.NewTicker(gap)
+ defer ticker.Stop()
+ for {
+ select {
+ case <-ticker.C:
+ select {
+ case progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)}:
+ case <-senderStop:
+ return
+ }
+ case <-senderStop:
+ return
+ }
+ }
+ }()
+
+ result := make(chan error, 1)
+ go func() {
+ result <- drainTrailingProgress(context.Background(), progressCh, func(StreamEvent) error { return nil }, window)
+ }()
+
+ var err error
+ select {
+ case err = <-result:
+ case <-time.After(5 * window):
+ close(senderStop)
+ <-senderDone
+ close(progressCh)
+ <-result
+ t.Fatal("continuous progress extended the terminal drain past its absolute window")
+ }
+ close(senderStop)
+ <-senderDone
+ if err != nil {
+ t.Fatalf("drainTrailingProgress returned error: %v", err)
+ }
+}
+
+func TestDrainTrailingProgressStopsWhenCallIsCanceled(t *testing.T) {
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+
+ err := drainTrailingProgress(
+ ctx,
+ make(chan StreamEvent),
+ func(StreamEvent) error {
+ t.Fatal("deliver must not run after cancellation")
+ return nil
+ },
+ time.Hour,
+ )
+ if !errors.Is(err, context.Canceled) {
+ t.Fatalf("drainTrailingProgress error = %v, want context.Canceled", err)
+ }
+}
+
+// TestInvokeStreamMaxDurationAbortsStreamThatNeverGoesIdle pins the one
+// behavior that distinguishes MaxDuration from IdleTimeout: an upstream
+// emitting events frequently enough that the idle bound never fires must
+// still be cut off once the total response-reading phase exceeds
+// MaxDuration.
+//
+// Two clocks race inside doHTTPToolCallStream; this server's behavior lets
+// only one of them ever fire:
+//
+// upstream (every 25ms) idle deadline (IdleTimeout=2s) max-duration deadline (300ms)
+// ---------------------- ------------------------------ -----------------------------
+// t=25ms progress event -------> reset to t=2025ms
+// t=50ms progress event -------> reset to t=2050ms
+// t=75ms progress event -------> reset to t=2075ms
+// ... (never a 2s gap, so the
+// idle deadline keeps
+// sliding into the future)
+// fixed at t=300ms, never reset
+// t=300ms -------------------------------------------------------> fires first
+// InvokeStream returns ErrStreamTimeout
+// ("max stream duration")
+//
+// serverDone (closed in t.Cleanup) then stops the emitting loop.
+func TestInvokeStreamMaxDurationAbortsStreamThatNeverGoesIdle(t *testing.T) {
+ serverDone := make(chan struct{})
+ server := invokeStreamTestServer(t, "sid-max-duration", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ for {
+ select {
+ case <-serverDone:
+ return
+ case <-time.After(25 * time.Millisecond):
+ }
+ if _, err := io.WriteString(w, "data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}\n\n"); err != nil {
+ return // client gave up; stop emitting
+ }
+ flusher.Flush()
+ }
+ })
+ t.Cleanup(func() {
+ close(serverDone)
+ server.Close()
+ })
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ start := time.Now()
+ _, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "stories.get", map[string]any{}, nil, nil, sink,
+ StreamOptions{IdleTimeout: 2 * time.Second, MaxDuration: 300 * time.Millisecond},
+ )
+ elapsed := time.Since(start)
+
+ if err == nil {
+ t.Fatal("expected a max-duration timeout error")
+ }
+ if !strings.Contains(err.Error(), "max stream duration") {
+ t.Fatalf("expected a max stream duration error, got %v", err)
+ }
+ if !errors.Is(err, ErrStreamTimeout) {
+ t.Fatalf("expected errors.Is(err, ErrStreamTimeout) to hold, got %v", err)
+ }
+ if elapsed > 5*time.Second {
+ t.Fatalf("max-duration abort took too long: %s", elapsed)
+ }
+ if len(sink.events) == 0 {
+ t.Fatal("expected events to have been relayed before the max-duration cutoff")
+ }
+}
+
+// TestInvokeStreamSSECommentKeepalivesResetIdleTimeout pins that comment
+// (":keepalive") lines count as upstream activity for the idle bound: a
+// long-quiet tool heartbeating through SSE comments must not be cut off,
+// even though no event arrives for longer than IdleTimeout.
+func TestInvokeStreamSSECommentKeepalivesResetIdleTimeout(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-keepalive", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ // 600ms of comment-only heartbeats against a 250ms idle timeout:
+ // no event arrives until well past the idle bound, so only the
+ // comments can be keeping the call alive.
+ for range 12 {
+ time.Sleep(50 * time.Millisecond)
+ if _, err := io.WriteString(w, ": keepalive\n"); err != nil {
+ return
+ }
+ flusher.Flush()
+ }
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done after heartbeats"}]}}`)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "stories.get", map[string]any{}, nil, nil, sink,
+ StreamOptions{IdleTimeout: 250 * time.Millisecond},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !strings.Contains(string(result.Body), "done after heartbeats") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+ if len(sink.events) != 0 {
+ t.Fatalf("comments are not events; expected none relayed, got %d", len(sink.events))
+ }
+}
+
+// resumeFailureTestServer serves a tools/call SSE response that ends without
+// a terminal message (forcing a resume) and dispatches the resume GET to
+// getHandler. "retry: 1" keeps the reconnect delay at its 200ms floor.
+func resumeFailureTestServer(t *testing.T, getHandler func(w http.ResponseWriter, r *http.Request)) *httptest.Server {
+ t.Helper()
+ return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ getHandler(w, r)
+ return
+ }
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Errorf("decode request: %v", err)
+ return
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-resume-fail")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ _, _ = io.WriteString(w, "id: evt-1\nretry: 1\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}\n\n")
+ flusher.Flush()
+ // Close without the terminal response → client resumes via GET.
+ default:
+ t.Errorf("unexpected method %q", req.Method)
+ }
+ }))
+}
+
+func TestInvokeStreamResumeFailureSurfacesUpstreamHTTPError(t *testing.T) {
+ server := resumeFailureTestServer(t, func(w http.ResponseWriter, r *http.Request) {
+ http.Error(w, "resume rejected", http.StatusInternalServerError)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ _, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "stories.get", map[string]any{}, nil, nil, sink,
+ StreamOptions{IdleTimeout: 2 * time.Second, MaxDuration: 5 * time.Second},
+ )
+ if err == nil {
+ t.Fatal("expected an error when the resume GET fails")
+ }
+ if !strings.Contains(err.Error(), "resume failed with HTTP 500") {
+ t.Fatalf("expected the resume HTTP status surfaced, got %v", err)
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected the pre-disconnect event to have been relayed, got %d", len(sink.events))
+ }
+}
+
+func TestInvokeStreamResumeRejectsNonSSEContentType(t *testing.T) {
+ server := resumeFailureTestServer(t, func(w http.ResponseWriter, r *http.Request) {
+ w.Header().Set("Content-Type", "application/json")
+ _, _ = io.WriteString(w, `{"unexpected":"plain body"}`)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &fakeStreamSink{}
+
+ _, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "stories.get", map[string]any{}, nil, nil, sink,
+ StreamOptions{IdleTimeout: 2 * time.Second, MaxDuration: 5 * time.Second},
+ )
+ if err == nil {
+ t.Fatal("expected an error when the resume response is not SSE")
+ }
+ if !strings.Contains(err.Error(), "want text/event-stream") {
+ t.Fatalf("expected the content-type mismatch surfaced, got %v", err)
+ }
+}
diff --git a/internal/mcp/sse_reader.go b/internal/mcp/sse_reader.go
new file mode 100644
index 00000000..9329b3e7
--- /dev/null
+++ b/internal/mcp/sse_reader.go
@@ -0,0 +1,177 @@
+package mcp
+
+import (
+ "bufio"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "strconv"
+ "strings"
+ "time"
+)
+
+// sseEventReader incrementally parses a Server-Sent Events body, returning
+// one joined "data:" payload per event via Next. It is the shared parser
+// behind both the buffered SSE consumers (extractSSEJSONRPCResponse, used by
+// tools/list, initialize, and the default forward path) and the incremental
+// streaming relay (relaySSEToolCall) — one parser, two ways of consuming it.
+type sseEventReader struct {
+ scanner *bufio.Scanner
+ maxBytes int
+ eventSize int
+ dataLines []string
+ eventID string
+ retry time.Duration
+ hasData bool
+ hasID bool
+ hasRetry bool
+ // onActivity, when set, fires once per line read from the stream —
+ // including comment (":keepalive") and other control lines that never
+ // surface as events. Streaming consumers hook their idle-timeout reset
+ // here so an upstream heartbeating via SSE comments during a long tool
+ // run counts as liveness even though no event arrives.
+ onActivity func()
+}
+
+type sseWireEvent struct {
+ Data []byte
+ ID string
+ Retry time.Duration
+ HasData bool
+ HasID bool
+ HasRetry bool
+}
+
+func newSSEEventReader(r io.Reader) *sseEventReader {
+ return newSSEEventReaderWithLimit(r, defaultStreamMaxMessageBytes)
+}
+
+func newSSEEventReaderWithLimit(r io.Reader, maxBytes int) *sseEventReader {
+ if maxBytes <= 0 {
+ maxBytes = defaultStreamMaxMessageBytes
+ }
+ scanner := bufio.NewScanner(r)
+ initialSize := min(maxBytes, 64*1024)
+ scanner.Buffer(make([]byte, initialSize), maxBytes)
+ return &sseEventReader{scanner: scanner, maxBytes: maxBytes}
+}
+
+// NextEvent returns one complete SSE event, including the id/retry fields
+// needed to resume a Streamable HTTP response after the upstream closes it.
+func (r *sseEventReader) NextEvent() (sseWireEvent, error) {
+ for r.scanner.Scan() {
+ if r.onActivity != nil {
+ r.onActivity()
+ }
+ line := r.scanner.Text()
+ if line == "" {
+ if !r.hasData && !r.hasID && !r.hasRetry {
+ r.eventSize = 0
+ continue
+ }
+ return r.takeEvent(), nil
+ }
+ if len(line)+1 > r.maxBytes-r.eventSize {
+ return sseWireEvent{}, ErrStreamMessageTooLarge
+ }
+ r.eventSize += len(line) + 1
+ if strings.HasPrefix(line, ":") {
+ continue
+ }
+ field, value, ok := strings.Cut(line, ":")
+ if !ok {
+ field = line
+ value = ""
+ } else if strings.HasPrefix(value, " ") {
+ value = strings.TrimPrefix(value, " ")
+ }
+ switch field {
+ case "data":
+ r.dataLines = append(r.dataLines, value)
+ r.hasData = true
+ case "id":
+ // The SSE specification ignores id values containing NUL.
+ if !strings.ContainsRune(value, '\x00') {
+ r.eventID = value
+ r.hasID = true
+ }
+ case "retry":
+ millis, err := strconv.ParseInt(value, 10, 64)
+ if err == nil && millis >= 0 {
+ const maxRetryMillis = int64((time.Duration(1<<63 - 1)) / time.Millisecond)
+ if millis > maxRetryMillis {
+ r.retry = time.Duration(1<<63 - 1)
+ } else {
+ r.retry = time.Duration(millis) * time.Millisecond
+ }
+ r.hasRetry = true
+ }
+ }
+ }
+ if err := r.scanner.Err(); err != nil {
+ if errors.Is(err, bufio.ErrTooLong) {
+ return sseWireEvent{}, ErrStreamMessageTooLarge
+ }
+ return sseWireEvent{}, err
+ }
+ if r.hasData || r.hasID || r.hasRetry {
+ return r.takeEvent(), nil
+ }
+ return sseWireEvent{}, io.EOF
+}
+
+func (r *sseEventReader) takeEvent() sseWireEvent {
+ evt := sseWireEvent{
+ Data: []byte(strings.Join(r.dataLines, "\n")),
+ ID: r.eventID,
+ Retry: r.retry,
+ HasData: r.hasData,
+ HasID: r.hasID,
+ HasRetry: r.hasRetry,
+ }
+ r.dataLines = nil
+ r.eventID = ""
+ r.retry = 0
+ r.hasData = false
+ r.hasID = false
+ r.hasRetry = false
+ r.eventSize = 0
+ return evt
+}
+
+// Next is the payload-only view used by buffered consumers. Control-only
+// events (id/retry with no data) are skipped because they carry no JSON-RPC
+// message for those callers to decode.
+func (r *sseEventReader) Next() ([]byte, error) {
+ for {
+ evt, err := r.NextEvent()
+ if err != nil {
+ return nil, err
+ }
+ if evt.HasData {
+ return evt.Data, nil
+ }
+ }
+}
+
+// extractSSEJSONRPCResponse scans an SSE body for the one event that is
+// either the response matching expectedID or the null-id error JSON-RPC
+// uses when a server can't identify which request an error belongs to,
+// skipping everything else (notifications, unrelated responses).
+func extractSSEJSONRPCResponse(r io.Reader, expectedID json.RawMessage) ([]byte, error) {
+ reader := newSSEEventReader(r)
+ for {
+ payload, err := reader.Next()
+ if err == io.EOF {
+ return nil, fmt.Errorf("no JSON-RPC response in SSE stream")
+ }
+ if err != nil {
+ return nil, err
+ }
+ match := classifyJSONRPCResponsePayload(payload, expectedID)
+ if match == jsonRPCResponseIDMatch || match == jsonRPCResponseNullIDError {
+ return payload, nil
+ }
+ }
+}
diff --git a/internal/mcp/sse_reader_test.go b/internal/mcp/sse_reader_test.go
new file mode 100644
index 00000000..fa20b50f
--- /dev/null
+++ b/internal/mcp/sse_reader_test.go
@@ -0,0 +1,118 @@
+package mcp
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+)
+
+func TestListToolsDecodesMultilineSSEData(t *testing.T) {
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-multiline")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/list":
+ writeTestSSEEvents(w, []string{
+ `{"jsonrpc":"2.0",`,
+ `"id":1,`,
+ `"result":{"tools":[{"name":"stories.multiline"}]}}`,
+ })
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ tools, err := client.ListTools(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL})
+ if err != nil {
+ t.Fatalf("ListTools returned error: %v", err)
+ }
+ if len(tools) != 1 || tools[0].Name != "stories.multiline" {
+ t.Fatalf("unexpected tools: %#v", tools)
+ }
+}
+
+func TestSSEEventReaderSaturatesOversizedRetryWithoutOverflow(t *testing.T) {
+ reader := newSSEEventReader(strings.NewReader("retry: 9223372036854775807\n\n"))
+ evt, err := reader.NextEvent()
+ if err != nil {
+ t.Fatal(err)
+ }
+ if !evt.HasRetry {
+ t.Fatal("expected retry field to be parsed")
+ }
+ if evt.Retry <= 0 {
+ t.Fatalf("oversized retry overflowed to %s; want a positive saturated duration", evt.Retry)
+ }
+}
+
+func TestSSEEventReaderRejectsAggregateEventOverLimit(t *testing.T) {
+ reader := newSSEEventReaderWithLimit(
+ strings.NewReader("data: 12345678\ndata: 12345678\n\n"),
+ 16,
+ )
+
+ _, err := reader.NextEvent()
+ if !errors.Is(err, ErrStreamMessageTooLarge) {
+ t.Fatalf("NextEvent error = %v, want ErrStreamMessageTooLarge", err)
+ }
+}
+
+func TestSSEEventReaderRejectsSingleLineOverLimitWithTypedError(t *testing.T) {
+ reader := newSSEEventReaderWithLimit(strings.NewReader("data: 1234567890\n\n"), 12)
+
+ _, err := reader.NextEvent()
+ if !errors.Is(err, ErrStreamMessageTooLarge) {
+ t.Fatalf("NextEvent error = %v, want ErrStreamMessageTooLarge", err)
+ }
+}
+
+func TestSSEEventReaderDoesNotAccumulateCommentBytesAcrossEvents(t *testing.T) {
+ reader := newSSEEventReaderWithLimit(
+ strings.NewReader(": ping\n\n: ping\n\n: ping\n\ndata: ok\n\n"),
+ 12,
+ )
+
+ evt, err := reader.NextEvent()
+ if err != nil {
+ t.Fatalf("NextEvent returned error after bounded comment events: %v", err)
+ }
+ if string(evt.Data) != "ok" {
+ t.Fatalf("event data = %q, want ok", evt.Data)
+ }
+}
+
+// TestSSEEventReaderOnActivityFiresPerLineIncludingComments pins the
+// liveness hook streaming consumers rely on for idle-timeout resets: every
+// line read counts as activity, including comment/keepalive lines that never
+// surface as events.
+func TestSSEEventReaderOnActivityFiresPerLineIncludingComments(t *testing.T) {
+ reader := newSSEEventReader(strings.NewReader(": keepalive\n: keepalive\ndata: {\"a\":1}\n\n"))
+ activity := 0
+ reader.onActivity = func() { activity++ }
+
+ evt, err := reader.NextEvent()
+ if err != nil {
+ t.Fatalf("NextEvent returned error: %v", err)
+ }
+ if !evt.HasData {
+ t.Fatalf("expected the data event, got %#v", evt)
+ }
+ // Two comment lines + one data line + the blank event terminator.
+ if activity != 4 {
+ t.Fatalf("onActivity fired %d times, want 4 (comments must count as activity)", activity)
+ }
+}
diff --git a/internal/mcp/standalone_stream.go b/internal/mcp/standalone_stream.go
new file mode 100644
index 00000000..6e904ed1
--- /dev/null
+++ b/internal/mcp/standalone_stream.go
@@ -0,0 +1,418 @@
+package mcp
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "net/http"
+ "strings"
+ "sync"
+ "sync/atomic"
+ "time"
+)
+
+// progressWaiter routes a shared-stream message to the call goroutine that
+// owns the sink. This prevents shared-reader writes from racing the call's
+// terminal response.
+type progressWaiter struct {
+ events chan StreamEvent
+ dropped *atomic.Int64
+}
+
+// A full waiter buffer drops that call's event rather than blocking the one
+// reader shared by every active call.
+const standaloneWaiterEventBuffer = 32
+
+// standaloneStream manages the shared SSE GET used for server-initiated
+// messages. It is needed for SDKs that send progress without a related
+// request ID, placing it on this connection instead of the tools/call POST.
+// terminalSettleWindow briefly drains progress that races the terminal across
+// the independent POST and standalone connections. It is an absolute bound:
+// progress arriving after the terminal response cannot extend the call.
+const terminalSettleWindow = 25 * time.Millisecond
+
+type standaloneStream struct {
+ key standaloneStreamKey
+ mu sync.Mutex
+ refCount int
+ cancel context.CancelFunc
+ done chan struct{}
+ waiters map[string]progressWaiter
+ // unsupported is set once opening the connection fails outright (e.g. a
+ // 404/405, which some upstreams legitimately return for this endpoint
+ // per spec). It stops every later acquire from re-attempting a doomed
+ // connection on every single streaming call; it resets naturally the
+ // next time refCount drops to zero and this entry is evicted.
+ unsupported bool
+ maxBytes int
+ // ready closes once the stream's first connect attempt (success or
+ // failure) has returned. A caller about to send a request that can
+ // trigger progress on this stream must wait on it first — otherwise the
+ // request can reach the upstream and complete its first progress update
+ // before the GET has even been issued, silently dropping that update.
+ ready chan struct{}
+}
+
+type standaloneStreamKey struct {
+ upstreamName string
+ sessionID string
+ protocol string
+}
+
+type standaloneUnsupportedError struct {
+ err error
+}
+
+func (e *standaloneUnsupportedError) Error() string {
+ return e.err.Error()
+}
+
+func (e *standaloneUnsupportedError) Unwrap() error {
+ return e.err
+}
+
+// acquireStandaloneStream returns the shared standaloneStream for upstream,
+// creating it and starting its reader goroutine if this is the first
+// waiter. Callers must pair this with exactly one releaseStandaloneStream.
+func (c *Client) acquireStandaloneStream(upstream Upstream) *standaloneStream {
+ return c.acquireStandaloneStreamWithLimit(upstream, defaultStreamMaxMessageBytes)
+}
+
+func (c *Client) acquireStandaloneStreamWithLimit(upstream Upstream, maxMessageBytes int) *standaloneStream {
+ key := standaloneStreamKey{
+ upstreamName: upstream.Name,
+ sessionID: c.getSession(upstream.Name),
+ protocol: c.getSessionProtocol(upstream.Name),
+ }
+ c.standaloneMu.Lock()
+ s := c.standaloneStreams[key]
+ if s == nil {
+ s = &standaloneStream{
+ key: key,
+ waiters: make(map[string]progressWaiter),
+ maxBytes: maxMessageBytes,
+ ready: make(chan struct{}),
+ }
+ c.standaloneStreams[key] = s
+ }
+ s.mu.Lock()
+ s.refCount++
+ start := s.refCount == 1 && !s.unsupported
+ if start {
+ streamCtx, cancel := context.WithCancel(context.Background())
+ done := make(chan struct{})
+ s.cancel = cancel
+ s.done = done
+ go c.runStandaloneStream(streamCtx, upstream, s, done)
+ }
+ s.mu.Unlock()
+ c.standaloneMu.Unlock()
+ return s
+}
+
+// releaseStandaloneStream drops one reference acquired via
+// acquireStandaloneStream. The map and reference count change atomically so a
+// concurrent acquire cannot reuse an entry that this release is about to
+// evict. Once the last reference is gone, it cancels the reader goroutine and
+// waits for it to fully exit.
+func (c *Client) releaseStandaloneStream(s *standaloneStream) {
+ c.standaloneMu.Lock()
+ s.mu.Lock()
+ s.refCount--
+ last := s.refCount <= 0
+ var cancel context.CancelFunc
+ var done chan struct{}
+ if last {
+ cancel = s.cancel
+ done = s.done
+ s.cancel = nil
+ s.done = nil
+ if c.standaloneStreams[s.key] == s {
+ delete(c.standaloneStreams, s.key)
+ }
+ }
+ s.mu.Unlock()
+ c.standaloneMu.Unlock()
+ if cancel != nil {
+ cancel()
+ <-done
+ }
+}
+
+// defaultStandaloneReadyTimeout bounds waitForStandaloneReady when the caller
+// supplies no header timeout, so a hung upstream can't block a call forever
+// waiting on a connection attempt that will never resolve.
+const defaultStandaloneReadyTimeout = 10 * time.Second
+
+// waitForStandaloneReady blocks until s's first connect attempt completes, so
+// the tools/call request sent right after this returns cannot race ahead of
+// the standalone GET subscribing upstream. Bounded by headerTimeout (or
+// defaultStandaloneReadyTimeout when unset) and ctx, so a slow or hanging
+// upstream degrades to the pre-fix race rather than blocking indefinitely.
+func waitForStandaloneReady(ctx context.Context, s *standaloneStream, headerTimeout time.Duration) {
+ if headerTimeout <= 0 {
+ headerTimeout = defaultStandaloneReadyTimeout
+ }
+ timer := time.NewTimer(headerTimeout)
+ defer timer.Stop()
+ select {
+ case <-s.ready:
+ case <-ctx.Done():
+ case <-timer.C:
+ }
+}
+
+func (s *standaloneStream) registerWaiter(token string, w progressWaiter) {
+ s.mu.Lock()
+ s.waiters[token] = w
+ s.mu.Unlock()
+}
+
+func (s *standaloneStream) unregisterWaiter(token string) {
+ s.mu.Lock()
+ delete(s.waiters, token)
+ s.mu.Unlock()
+}
+
+// openStandaloneGET opens the standalone SSE stream: a bare GET carrying the
+// session's headers, no Last-Event-ID (see standaloneStream doc comment).
+// Mirrors resumeSSEStream's header handling.
+func (c *Client) openStandaloneGET(ctx context.Context, upstream Upstream, sessionID, protocol string) (*http.Response, error) {
+ endpoint := strings.TrimRight(upstream.BaseURL, "/")
+ req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
+ if err != nil {
+ return nil, err
+ }
+ req.Header.Set("Accept", "text/event-stream")
+ if protocol != "" {
+ req.Header.Set("MCP-Protocol-Version", protocol)
+ }
+ if sessionID != "" {
+ req.Header.Set("Mcp-Session-Id", sessionID)
+ }
+ applyAuthHeaders(req, upstream)
+ resp, err := c.httpClient.Do(req)
+ if err != nil {
+ return nil, err
+ }
+ if resp.StatusCode >= http.StatusBadRequest {
+ defer resp.Body.Close()
+ body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
+ err := fmt.Errorf("upstream %q standalone stream returned HTTP %d: %s", upstream.Name, resp.StatusCode, extractErrorDetail(body))
+ if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusMethodNotAllowed {
+ return nil, &standaloneUnsupportedError{err: err}
+ }
+ return nil, err
+ }
+ if !strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") {
+ defer resp.Body.Close()
+ return nil, &standaloneUnsupportedError{err: fmt.Errorf("upstream %q standalone stream returned content type %q, want text/event-stream", upstream.Name, resp.Header.Get("Content-Type"))}
+ }
+ return resp, nil
+}
+
+func (c *Client) runStandaloneStream(ctx context.Context, upstream Upstream, s *standaloneStream, done chan<- struct{}) {
+ defer close(done)
+ attempt := 0
+ firstAttempt := true
+ for {
+ resp, err := c.openStandaloneGET(ctx, upstream, s.key.sessionID, s.key.protocol)
+ if firstAttempt {
+ firstAttempt = false
+ close(s.ready)
+ }
+ if err != nil {
+ if ctx.Err() != nil {
+ return
+ }
+ var unsupported *standaloneUnsupportedError
+ if errors.As(err, &unsupported) {
+ c.debugf("standalone stream unsupported server=%s session=%q err=%v", upstream.Name, s.key.sessionID, err)
+ s.mu.Lock()
+ s.unsupported = true
+ s.mu.Unlock()
+ return
+ }
+ c.debugf("standalone stream open failed; retrying server=%s session=%q err=%v", upstream.Name, s.key.sessionID, err)
+ if waitForSSEReconnect(ctx, sseReconnectDelay(0, attempt)) != nil {
+ return
+ }
+ attempt++
+ continue
+ }
+
+ reader := newSSEEventReaderWithLimit(resp.Body, s.maxBytes)
+ for {
+ evt, readErr := reader.NextEvent()
+ if readErr != nil {
+ _ = resp.Body.Close()
+ if ctx.Err() != nil {
+ return
+ }
+ c.debugf("standalone stream disconnected; retrying server=%s session=%q err=%v", upstream.Name, s.key.sessionID, readErr)
+ break
+ }
+ attempt = 0
+ if evt.HasData {
+ c.routeStandaloneEvent(s, evt.Data)
+ }
+ }
+ if waitForSSEReconnect(ctx, sseReconnectDelay(0, attempt)) != nil {
+ return
+ }
+ attempt++
+ }
+}
+
+// routeStandaloneEvent attributes one standalone-stream message to whichever
+// registered call it belongs to. Progress notifications carry the token
+// Atryum minted for that call (see rewriteProgressToken) in
+// params.progressToken, giving an unambiguous match. Anything else (e.g. a
+// logging notification) carries no per-call correlator at all; it is
+// delivered only when exactly one call is currently waiting on this stream,
+// since there is no way to attribute it correctly when several calls are
+// in flight concurrently — and silently guessing wrong would leak one
+// caller's message to another.
+func (c *Client) routeStandaloneEvent(s *standaloneStream, payload []byte) {
+ var message map[string]json.RawMessage
+ if err := json.Unmarshal(payload, &message); err != nil {
+ return
+ }
+ if _, hasMethod := message["method"]; !hasMethod {
+ return
+ }
+ // method plus id is a server-to-client request, not a notification — the
+ // same distinction classifyRPCMessage draws on the POST stream. Flag it
+ // so a sampling/elicitation request arriving here is audited as what it
+ // is instead of being mislabeled a notification.
+ _, isServerRequest := message["id"]
+ wireToken, hasToken := extractProgressToken(message)
+
+ s.mu.Lock()
+ var waiter progressWaiter
+ var ok bool
+ if hasToken {
+ waiter, ok = s.waiters[wireToken]
+ } else if len(s.waiters) == 1 {
+ for _, w := range s.waiters {
+ waiter, ok = w, true
+ }
+ }
+ s.mu.Unlock()
+ if !ok {
+ return
+ }
+
+ // Handed off to the matching call's own goroutine via its channel — see
+ // progressWaiter for why this indirection matters. callSink.Event (on
+ // the receiving end) restores the caller's original progressToken
+ // itself (matching on its own wireToken), so the raw payload is sent
+ // through unmodified here.
+ select {
+ case waiter.events <- StreamEvent{Data: payload, ServerRequest: isServerRequest}:
+ default:
+ // Buffer full, or the receiving call already stopped draining it —
+ // drop rather than block this shared reader goroutine, which also
+ // serves every other call currently sharing this connection.
+ if waiter.dropped != nil {
+ waiter.dropped.Add(1)
+ }
+ }
+}
+
+// extractProgressToken reads params.progressToken from an already-decoded
+// JSON-RPC message, normalizing it to a bare string for map lookup
+// regardless of whether the upstream echoed it back as a JSON string or a
+// number.
+func extractProgressToken(message map[string]json.RawMessage) (string, bool) {
+ paramsRaw, ok := message["params"]
+ if !ok {
+ return "", false
+ }
+ var params struct {
+ ProgressToken json.RawMessage `json:"progressToken"`
+ }
+ if err := json.Unmarshal(paramsRaw, ¶ms); err != nil || len(params.ProgressToken) == 0 {
+ return "", false
+ }
+ return strings.Trim(string(params.ProgressToken), `"`), true
+}
+
+// rewriteProgressTokenInPayload replaces params.progressToken in an
+// already-wire-formatted JSON-RPC message with originalToken, restoring the
+// value the caller actually supplied before relaying the message onward.
+func rewriteProgressTokenInPayload(payload []byte, originalToken any) ([]byte, error) {
+ var generic map[string]any
+ if err := json.Unmarshal(payload, &generic); err != nil {
+ return nil, err
+ }
+ params, ok := generic["params"].(map[string]any)
+ if !ok {
+ return nil, fmt.Errorf("message has no params object")
+ }
+ params["progressToken"] = originalToken
+ generic["params"] = params
+ return json.Marshal(generic)
+}
+
+// rewriteProgressToken replaces meta's progressToken, if any, with a value
+// unique to this specific call, returning the rewritten meta, that wire
+// token, and the caller's original token. Atryum multiplexes every
+// downstream caller of a given upstream onto one shared session, so two
+// unrelated concurrent calls could independently pick the same
+// caller-supplied progressToken; rewriting to a per-call value here is what
+// lets routeStandaloneEvent attribute a notification to the right call
+// instead of risking a cross-call delivery.
+func (c *Client) rewriteProgressToken(meta map[string]any) (rewritten map[string]any, wireToken string, original any, ok bool) {
+ if meta == nil {
+ return meta, "", nil, false
+ }
+ original, ok = meta["progressToken"]
+ if !ok {
+ return meta, "", nil, false
+ }
+ wireToken = fmt.Sprintf("atryum-pt-%d", c.nextID.Add(1))
+ rewritten = make(map[string]any, len(meta))
+ for k, v := range meta {
+ rewritten[k] = v
+ }
+ rewritten["progressToken"] = wireToken
+ return rewritten, wireToken, original, true
+}
+
+// callSink restores the caller's progressToken after Atryum replaces it with
+// a per-call wire token. It wraps both POST and standalone-stream delivery so
+// the internal token cannot leak through either path.
+type callSink struct {
+ inner StreamSink
+ wireToken string
+ originalToken any
+}
+
+func newCallSink(inner StreamSink, wireToken string, originalToken any) *callSink {
+ return &callSink{inner: inner, wireToken: wireToken, originalToken: originalToken}
+}
+
+func (s *callSink) StreamStarted() {
+ s.inner.StreamStarted()
+}
+
+func (s *callSink) Event(evt StreamEvent) error {
+ var message map[string]json.RawMessage
+ if err := json.Unmarshal(evt.Data, &message); err == nil {
+ if token, ok := extractProgressToken(message); ok && token == s.wireToken {
+ if rewritten, err := rewriteProgressTokenInPayload(evt.Data, s.originalToken); err == nil {
+ evt.Data = rewritten
+ }
+ }
+ }
+ return s.inner.Event(evt)
+}
+
+func (s *callSink) StreamStats(stats StreamStats) {
+ if sink, ok := s.inner.(StreamStatsSink); ok {
+ sink.StreamStats(stats)
+ }
+}
diff --git a/internal/mcp/standalone_stream_test.go b/internal/mcp/standalone_stream_test.go
new file mode 100644
index 00000000..8ea26e2e
--- /dev/null
+++ b/internal/mcp/standalone_stream_test.go
@@ -0,0 +1,999 @@
+package mcp
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "sync"
+ "sync/atomic"
+ "testing"
+ "time"
+)
+
+// syncFakeStreamSink is fakeStreamSink's mutex-protected counterpart.
+// Production delivery is single-goroutine by contract (see StreamSink's doc
+// comment — routeStandaloneEvent hands events to the call goroutine via a
+// channel, it never touches the sink itself), so the mutex is not covering
+// a production race. It exists for the tests here whose *assertion*
+// goroutine reads events while InvokeStream may still be delivering on its
+// own goroutine — the plain fakeStreamSink (client_test.go) would race in
+// that test-side pattern.
+type syncFakeStreamSink struct {
+ mu sync.Mutex
+ started bool
+ events []StreamEvent
+}
+
+func newSyncFakeStreamSink() *syncFakeStreamSink {
+ return &syncFakeStreamSink{}
+}
+
+func (s *syncFakeStreamSink) StreamStarted() {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.started = true
+}
+
+func (s *syncFakeStreamSink) Event(evt StreamEvent) error {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.events = append(s.events, evt)
+ return nil
+}
+
+func (s *syncFakeStreamSink) wasStarted() bool {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ return s.started
+}
+
+func (s *syncFakeStreamSink) snapshotEvents() []StreamEvent {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ out := make([]StreamEvent, len(s.events))
+ copy(out, s.events)
+ return out
+}
+
+func TestRewriteProgressToken(t *testing.T) {
+ client := NewHTTPClient()
+ if _, _, _, ok := client.rewriteProgressToken(nil); ok {
+ t.Fatal("expected no rewrite when meta is nil")
+ }
+ if _, _, _, ok := client.rewriteProgressToken(map[string]any{"atryumRequestId": "x"}); ok {
+ t.Fatal("expected no rewrite when meta has no progressToken")
+ }
+
+ rewritten, wireToken, original, ok := client.rewriteProgressToken(map[string]any{"progressToken": float64(7), "atryumRequestId": "req-1"})
+ if !ok {
+ t.Fatal("expected a rewrite when progressToken is present")
+ }
+ if original != float64(7) {
+ t.Fatalf("expected original token 7, got %#v", original)
+ }
+ if rewritten["progressToken"] != wireToken {
+ t.Fatalf("expected rewritten meta to carry the wire token, got %#v", rewritten["progressToken"])
+ }
+ if rewritten["atryumRequestId"] != "req-1" {
+ t.Fatal("expected other meta keys preserved")
+ }
+
+ _, wireToken2, _, _ := client.rewriteProgressToken(map[string]any{"progressToken": "other"})
+ if wireToken2 == wireToken {
+ t.Fatal("expected distinct wire tokens across calls, so concurrent callers can't collide")
+ }
+}
+
+func TestExtractAndRewriteProgressTokenInPayload(t *testing.T) {
+ payload := []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":"atryum-pt-3","progress":1,"total":3}}`)
+ var msg map[string]json.RawMessage
+ if err := json.Unmarshal(payload, &msg); err != nil {
+ t.Fatalf("unmarshal: %v", err)
+ }
+ token, ok := extractProgressToken(msg)
+ if !ok || token != "atryum-pt-3" {
+ t.Fatalf("extractProgressToken = (%q, %v), want (atryum-pt-3, true)", token, ok)
+ }
+
+ rewritten, err := rewriteProgressTokenInPayload(payload, float64(42))
+ if err != nil {
+ t.Fatalf("rewriteProgressTokenInPayload: %v", err)
+ }
+ if !strings.Contains(string(rewritten), `"progressToken":42`) {
+ t.Fatalf("expected original numeric token restored, got %s", rewritten)
+ }
+ if !strings.Contains(string(rewritten), `"progress":1`) {
+ t.Fatalf("expected other params fields preserved, got %s", rewritten)
+ }
+
+ if _, ok := extractProgressToken(map[string]json.RawMessage{"method": json.RawMessage(`"notifications/message"`)}); ok {
+ t.Fatal("expected no token when params is absent")
+ }
+}
+
+// TestRouteStandaloneEventAttributesTokenlessNotificationOnlyWhenUnambiguous
+// covers routeStandaloneEvent's fallback for messages with no progressToken
+// (e.g. a plain logging notification): deliverable only when exactly one
+// call is waiting on the stream, since guessing with several concurrent
+// waiters would leak one caller's message to another.
+func TestRouteStandaloneEventAttributesTokenlessNotificationOnlyWhenUnambiguous(t *testing.T) {
+ client := NewHTTPClient()
+ payload := []byte(`{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"hello"}}`)
+
+ chA := make(chan StreamEvent, 1)
+ lone := &standaloneStream{waiters: map[string]progressWaiter{"tok-a": {events: chA}}}
+ client.routeStandaloneEvent(lone, payload)
+ select {
+ case <-chA:
+ default:
+ t.Fatal("expected the lone waiter to receive a tokenless notification")
+ }
+
+ chB, chC := make(chan StreamEvent, 1), make(chan StreamEvent, 1)
+ ambiguous := &standaloneStream{waiters: map[string]progressWaiter{
+ "tok-b": {events: chB},
+ "tok-c": {events: chC},
+ }}
+ client.routeStandaloneEvent(ambiguous, payload)
+ select {
+ case <-chB:
+ t.Fatal("expected a tokenless notification to be dropped, not guessed, with multiple concurrent waiters")
+ case <-chC:
+ t.Fatal("expected a tokenless notification to be dropped, not guessed, with multiple concurrent waiters")
+ default:
+ }
+}
+
+func TestRouteStandaloneEventCountsDropsForFullWaiter(t *testing.T) {
+ client := NewHTTPClient()
+ events := make(chan StreamEvent, 1)
+ events <- StreamEvent{Data: []byte(`{"already":"queued"}`)}
+ var dropped atomic.Int64
+ stream := &standaloneStream{waiters: map[string]progressWaiter{
+ "tok": {events: events, dropped: &dropped},
+ }}
+
+ client.routeStandaloneEvent(
+ stream,
+ []byte(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":"tok","progress":2}}`),
+ )
+
+ if got := dropped.Load(); got != 1 {
+ t.Fatalf("standalone dropped count = %d, want 1", got)
+ }
+}
+
+// TestInvokeStreamStandaloneStreamRelaysProgressNotification is the
+// regression test for the real end-to-end gap this feature fixes: the
+// reference MCP Python SDK's Context.report_progress sends progress
+// notifications on the standalone GET stream, never on the tools/call POST
+// response body, because it doesn't attribute the notification to the
+// request that triggered it. Without a standalone-stream reader, Atryum
+// would relay zero progress notifications for such a server even though the
+// terminal response arrives correctly.
+func TestInvokeStreamStandaloneStreamRelaysProgressNotification(t *testing.T) {
+ tokenCh := make(chan string, 1)
+ notifSent := make(chan struct{})
+
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ if got := r.Header.Get("Last-Event-ID"); got != "" {
+ t.Fatalf("standalone GET unexpectedly carried Last-Event-ID=%q (that's the resume path)", got)
+ }
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ token := <-tokenCh
+ writeTestSSEEventFlush(w, flusher, fmt.Sprintf(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":%q,"progress":1}}`, token))
+ close(notifSent)
+ <-r.Context().Done()
+ return
+ }
+
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-standalone")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ var params struct {
+ Meta struct {
+ ProgressToken string `json:"progressToken"`
+ } `json:"_meta"`
+ }
+ _ = json.Unmarshal(req.Params, ¶ms)
+ tokenCh <- params.Meta.ProgressToken
+ <-notifSent
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := newSyncFakeStreamSink()
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "real", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "slow_streaming_task", map[string]any{}, nil, map[string]any{"progressToken": "caller-token"}, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+ if !sink.wasStarted() {
+ t.Fatal("expected StreamStarted to fire for a notification delivered only via the standalone stream")
+ }
+ events := sink.snapshotEvents()
+ if len(events) != 1 {
+ t.Fatalf("expected exactly one relayed event, got %d: %#v", len(events), events)
+ }
+ if !strings.Contains(string(events[0].Data), `"progressToken":"caller-token"`) {
+ t.Fatalf("expected the caller's original progressToken restored, got %s", events[0].Data)
+ }
+}
+
+// TestInvokeStreamStandaloneProgressResetsIdleTimeout proves progress
+// notifications delivered only on the standalone GET stream still count as
+// activity for the call's IdleTimeout, even though the tools/call POST
+// connection itself stays silent while waiting on them.
+//
+// standalone GET tools/call POST
+// -------------- ----------------
+// tokenCh <- progressToken
+// token := <-tokenCh
+// sleep 60ms; send progress=1 ---------------------> (idle timer reset;
+// sleep 60ms; send progress=2 ---------------------> 150ms IdleTimeout
+// sleep 60ms; send progress=3 ---------------------> never sees a gap
+// sleep 60ms; send progress=4 ---------------------> wider than 60ms)
+// close(progressComplete) ---------------------> <-progressComplete
+// <-r.Context().Done() write terminal "done"
+//
+// Total elapsed (~240ms) exceeds IdleTimeout (150ms), but no single gap
+// between events does — if standalone-stream activity didn't reset the
+// call's idle timer, this call would time out despite the upstream still
+// actively working.
+func TestInvokeStreamStandaloneProgressResetsIdleTimeout(t *testing.T) {
+ tokenCh := make(chan string, 1)
+ progressComplete := make(chan struct{})
+
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ token := <-tokenCh
+ for progress := 1; progress <= 4; progress++ {
+ time.Sleep(60 * time.Millisecond)
+ writeTestSSEEventFlush(w, flusher, fmt.Sprintf(
+ `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":%q,"progress":%d}}`,
+ token,
+ progress,
+ ))
+ }
+ close(progressComplete)
+ <-r.Context().Done()
+ return
+ }
+
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-standalone-idle")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ var params struct {
+ Meta struct {
+ ProgressToken string `json:"progressToken"`
+ } `json:"_meta"`
+ }
+ _ = json.Unmarshal(req.Params, ¶ms)
+ tokenCh <- params.Meta.ProgressToken
+
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ _, _ = w.Write([]byte(": stream ready\n\n"))
+ flusher.Flush()
+ <-progressComplete
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := newSyncFakeStreamSink()
+
+ result, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "standalone-idle", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "slow_streaming_task",
+ map[string]any{},
+ nil,
+ map[string]any{"progressToken": "caller-token"},
+ sink,
+ StreamOptions{IdleTimeout: 150 * time.Millisecond},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error while standalone progress remained active: %v", err)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+ if events := sink.snapshotEvents(); len(events) != 4 {
+ t.Fatalf("expected four relayed progress events, got %d", len(events))
+ }
+}
+
+// TestInvokeStreamRebindsStandaloneStreamAfterSessionRenewal proves that
+// when a session goes stale mid-call, the standalone stream reconnects
+// under the renewed session rather than staying bound to the old one.
+//
+// tools/call POST standalone GET
+// ---------------- --------------
+// initialize -> sid-1 (initializeCount=1)
+// connects with
+// Mcp-Session-Id: sid-1
+// closes sid1Connected,
+// blocks on ctx.Done()
+// tools/call #1: <-sid1Connected <---------------------
+// responds: error "No session ID
+// provided..." (session looks stale)
+// client re-initializes -> sid-2 (initializeCount=2)
+// reconnects with
+// Mcp-Session-Id: sid-2
+// closes sid2Connected,
+// blocks on <-tokenForRetry
+// tools/call #2 (Mcp-Session-Id: sid-2)
+// tokenForRetry <- progressToken ---------------------->
+// receives token, sends
+// progress=2 notification,
+// closes progressSent
+// <-sid2Connected, then <-progressSent <----------------
+// writes terminal "done"
+//
+// The one relayed event must come from the sid-2 standalone stream, with the
+// caller's original progressToken restored, proving the rebind — not the
+// original sid-1 connection, which never got to deliver anything.
+func TestInvokeStreamRebindsStandaloneStreamAfterSessionRenewal(t *testing.T) {
+ var initializeCount atomic.Int32
+ var toolsCallCount atomic.Int32
+ sid1Connected := make(chan struct{})
+ sid2Connected := make(chan struct{})
+ tokenForRetry := make(chan string, 1)
+ progressSent := make(chan struct{})
+ var closeSID1, closeSID2, closeProgress sync.Once
+
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ w.Header().Set("Content-Type", "text/event-stream")
+ w.(http.Flusher).Flush()
+ switch r.Header.Get("Mcp-Session-Id") {
+ case "sid-1":
+ closeSID1.Do(func() { close(sid1Connected) })
+ case "sid-2":
+ closeSID2.Do(func() { close(sid2Connected) })
+ token := <-tokenForRetry
+ writeTestSSEEventFlush(w, w.(http.Flusher), fmt.Sprintf(
+ `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":%q,"progress":2}}`,
+ token,
+ ))
+ closeProgress.Do(func() { close(progressSent) })
+ }
+ <-r.Context().Done()
+ return
+ }
+
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ count := initializeCount.Add(1)
+ w.Header().Set("Mcp-Session-Id", fmt.Sprintf("sid-%d", count))
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ count := toolsCallCount.Add(1)
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ if count == 1 {
+ <-sid1Connected
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"No session ID provided for non-initialization request"}}`)
+ return
+ }
+ if got := r.Header.Get("Mcp-Session-Id"); got != "sid-2" {
+ t.Errorf("retry tools/call session = %q, want sid-2", got)
+ }
+ var params struct {
+ Meta struct {
+ ProgressToken string `json:"progressToken"`
+ } `json:"_meta"`
+ }
+ _ = json.Unmarshal(req.Params, ¶ms)
+ tokenForRetry <- params.Meta.ProgressToken
+ select {
+ case <-sid2Connected:
+ case <-time.After(2 * time.Second):
+ t.Error("standalone stream never reconnected with sid-2")
+ }
+ select {
+ case <-progressSent:
+ case <-time.After(2 * time.Second):
+ t.Error("sid-2 standalone stream never delivered progress")
+ }
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := newSyncFakeStreamSink()
+ result, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "renew-session", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "tool",
+ map[string]any{},
+ nil,
+ map[string]any{"progressToken": "caller-token"},
+ sink,
+ StreamOptions{},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("terminal result = %s, want done", result.Body)
+ }
+ events := sink.snapshotEvents()
+ if len(events) != 1 || !strings.Contains(string(events[0].Data), `"progressToken":"caller-token"`) {
+ t.Fatalf("relayed events = %#v, want sid-2 progress with restored token", events)
+ }
+}
+
+func TestInvokeStreamStandaloneRetriesTransientOpenFailure(t *testing.T) {
+ var getCount atomic.Int32
+ tokenCh := make(chan string, 1)
+ progressSent := make(chan struct{})
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ if getCount.Add(1) == 1 {
+ http.Error(w, "temporary failure", http.StatusInternalServerError)
+ return
+ }
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ token := <-tokenCh
+ writeTestSSEEventFlush(w, flusher, fmt.Sprintf(
+ `{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":%q,"progress":1}}`,
+ token,
+ ))
+ close(progressSent)
+ <-r.Context().Done()
+ return
+ }
+
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-transient")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ var params struct {
+ Meta struct {
+ ProgressToken string `json:"progressToken"`
+ } `json:"_meta"`
+ }
+ _ = json.Unmarshal(req.Params, ¶ms)
+ tokenCh <- params.Meta.ProgressToken
+ select {
+ case <-progressSent:
+ case <-time.After(3 * time.Second):
+ t.Error("standalone stream did not retry its transient 500 response")
+ }
+ w.Header().Set("Content-Type", "text/event-stream")
+ writeTestSSEEventFlush(w, w.(http.Flusher), `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := newSyncFakeStreamSink()
+ _, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "transient", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "tool",
+ map[string]any{},
+ nil,
+ map[string]any{"progressToken": "caller-token"},
+ sink,
+ StreamOptions{},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if got := getCount.Load(); got < 2 {
+ t.Fatalf("standalone GET count = %d, want a retry after the transient failure", got)
+ }
+ if len(sink.snapshotEvents()) != 1 {
+ t.Fatalf("relayed events = %#v, want progress from retried standalone stream", sink.snapshotEvents())
+ }
+}
+
+// TestInvokeStreamRestoresCallerProgressTokenOnPOSTResponseStreamToo is a
+// regression test: some upstreams echo a call's progress notifications on
+// the tools/call POST response itself, not the standalone stream — that's
+// actually the more spec-typical case for a request-scoped notification.
+// The caller's original progressToken must be restored there too, not only
+// on notifications that happen to arrive via the standalone stream.
+func TestInvokeStreamRestoresCallerProgressTokenOnPOSTResponseStreamToo(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-post-echo", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ var params struct {
+ Meta struct {
+ ProgressToken string `json:"progressToken"`
+ } `json:"_meta"`
+ }
+ _ = json.Unmarshal(req.Params, ¶ms)
+
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, fmt.Sprintf(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":%q,"progress":1}}`, params.Meta.ProgressToken))
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := newSyncFakeStreamSink()
+
+ _, err := client.InvokeStream(context.Background(), Upstream{Name: "shortcut", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "stories.get", map[string]any{}, nil, map[string]any{"progressToken": "caller-token"}, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ events := sink.snapshotEvents()
+ if len(events) != 1 {
+ t.Fatalf("expected exactly one relayed event, got %d: %#v", len(events), events)
+ }
+ if !strings.Contains(string(events[0].Data), `"progressToken":"caller-token"`) {
+ t.Fatalf("expected the caller's original progressToken restored on the POST-response stream, got %s", events[0].Data)
+ }
+ if strings.Contains(string(events[0].Data), "atryum-pt-") {
+ t.Fatalf("expected Atryum's internal wire token never to leak to the agent, got %s", events[0].Data)
+ }
+}
+
+// TestInvokeStreamStandaloneStreamAvoidsProgressTokenCollisionAcrossCalls
+// proves two concurrent callers who happen to pick the same progressToken
+// don't cross-deliver: Atryum multiplexes every caller of an upstream onto
+// one shared session, so the standalone stream is shared too, and the only
+// thing preventing a collision is the per-call wire-token rewrite.
+//
+// Both callers ask for progressToken=1. Three gates force the same
+// interleaving every run, regardless of goroutine scheduling:
+//
+// tool-a POST ---\
+// +--> both block on <-getConnected
+// tool-b POST ---/ |
+// v
+// standalone GET -----> closes getConnected, then blocks on <-gotBothTokens
+// |
+// tool-a records its wireTokenA |
+// tool-b records its wireTokenB |
+// (postCount==2) ------------+---> closes gotBothTokens
+// |
+// standalone GET <-------------+ unblocks, then:
+// sends progress=1 on wireTokenA
+// sends progress=2 on wireTokenB
+// closes notifsDone
+// |
+// tool-a, tool-b <-------------+ both unblock, each sends its own
+// terminal response
+//
+// Both callers asked for progressToken=1, but each only ever sees its own
+// wire token restored back to 1, with its own progress value (A=1, B=2) —
+// never the other call's notification.
+func TestInvokeStreamStandaloneStreamAvoidsProgressTokenCollisionAcrossCalls(t *testing.T) {
+ var mu sync.Mutex
+ tokenFor := map[string]string{}
+ postCount := 0
+ getConnected := make(chan struct{})
+ gotBothTokens := make(chan struct{})
+ notifsDone := make(chan struct{})
+ var closeGetConnectedOnce, closeGotBothTokensOnce sync.Once
+
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ closeGetConnectedOnce.Do(func() { close(getConnected) })
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ <-gotBothTokens
+ mu.Lock()
+ tokA, tokB := tokenFor["tool-a"], tokenFor["tool-b"]
+ mu.Unlock()
+ writeTestSSEEventFlush(w, flusher, fmt.Sprintf(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":%q,"progress":1}}`, tokA))
+ writeTestSSEEventFlush(w, flusher, fmt.Sprintf(`{"jsonrpc":"2.0","method":"notifications/progress","params":{"progressToken":%q,"progress":2}}`, tokB))
+ close(notifsDone)
+ <-r.Context().Done()
+ return
+ }
+
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-collision")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ var params struct {
+ Name string `json:"name"`
+ Meta struct {
+ ProgressToken string `json:"progressToken"`
+ } `json:"_meta"`
+ }
+ _ = json.Unmarshal(req.Params, ¶ms)
+ <-getConnected
+ mu.Lock()
+ tokenFor[params.Name] = params.Meta.ProgressToken
+ postCount++
+ ready := postCount == 2
+ mu.Unlock()
+ if ready {
+ closeGotBothTokensOnce.Do(func() { close(gotBothTokens) })
+ }
+ <-notifsDone
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, fmt.Sprintf(`{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done-%s"}]}}`, params.Name))
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ upstream := Upstream{Name: "real", Mode: UpstreamModeHTTP, BaseURL: server.URL}
+
+ sinkA := newSyncFakeStreamSink()
+ sinkB := newSyncFakeStreamSink()
+
+ var wg sync.WaitGroup
+ var errA, errB error
+ wg.Add(2)
+ go func() {
+ defer wg.Done()
+ _, errA = client.InvokeStream(context.Background(), upstream, "tool-a", map[string]any{}, nil, map[string]any{"progressToken": float64(1)}, sinkA, StreamOptions{})
+ }()
+ go func() {
+ defer wg.Done()
+ _, errB = client.InvokeStream(context.Background(), upstream, "tool-b", map[string]any{}, nil, map[string]any{"progressToken": float64(1)}, sinkB, StreamOptions{})
+ }()
+ wg.Wait()
+
+ if errA != nil {
+ t.Fatalf("call A error: %v", errA)
+ }
+ if errB != nil {
+ t.Fatalf("call B error: %v", errB)
+ }
+
+ eventsA, eventsB := sinkA.snapshotEvents(), sinkB.snapshotEvents()
+ if len(eventsA) != 1 {
+ t.Fatalf("call A: expected exactly 1 relayed event, got %d: %#v", len(eventsA), eventsA)
+ }
+ if len(eventsB) != 1 {
+ t.Fatalf("call B: expected exactly 1 relayed event, got %d: %#v", len(eventsB), eventsB)
+ }
+ if !strings.Contains(string(eventsA[0].Data), `"progress":1`) || !strings.Contains(string(eventsA[0].Data), `"progressToken":1`) {
+ t.Fatalf("call A got the wrong notification or token, want its own progress=1/token=1, got %s", eventsA[0].Data)
+ }
+ if !strings.Contains(string(eventsB[0].Data), `"progress":2`) || !strings.Contains(string(eventsB[0].Data), `"progressToken":1`) {
+ t.Fatalf("call B got the wrong notification or token, want its own progress=2/token=1, got %s", eventsB[0].Data)
+ }
+}
+
+func TestStandaloneStreamRefcountsSharedConnection(t *testing.T) {
+ var connections int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method != http.MethodGet {
+ t.Fatalf("unexpected method %q", r.Method)
+ }
+ atomic.AddInt32(&connections, 1)
+ w.Header().Set("Content-Type", "text/event-stream")
+ w.(http.Flusher).Flush()
+ <-r.Context().Done()
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ upstream := Upstream{Name: "real", Mode: UpstreamModeHTTP, BaseURL: server.URL}
+
+ s1 := client.acquireStandaloneStream(upstream)
+ s2 := client.acquireStandaloneStream(upstream)
+ if s1 != s2 {
+ t.Fatal("expected the second acquire to reuse the same standaloneStream while the first is still active")
+ }
+
+ deadline := time.Now().Add(2 * time.Second)
+ for atomic.LoadInt32(&connections) < 1 && time.Now().Before(deadline) {
+ time.Sleep(10 * time.Millisecond)
+ }
+ if got := atomic.LoadInt32(&connections); got != 1 {
+ t.Fatalf("expected exactly 1 standalone connection while both waiters are active, got %d", got)
+ }
+
+ client.releaseStandaloneStream(s1)
+ if got := atomic.LoadInt32(&connections); got != 1 {
+ t.Fatalf("releasing one of two references should not close the connection yet, got %d", got)
+ }
+ client.releaseStandaloneStream(s2)
+
+ s3 := client.acquireStandaloneStream(upstream)
+ if s3 == s1 {
+ t.Fatal("expected a fresh standaloneStream after the previous one was fully released")
+ }
+ deadline = time.Now().Add(2 * time.Second)
+ for atomic.LoadInt32(&connections) < 2 && time.Now().Before(deadline) {
+ time.Sleep(10 * time.Millisecond)
+ }
+ if got := atomic.LoadInt32(&connections); got != 2 {
+ t.Fatalf("expected a new connection after full release + reacquire, got %d", got)
+ }
+ client.releaseStandaloneStream(s3)
+}
+
+// TestStandaloneStreamConcurrentLastReleaseAndAcquireKeepsLiveEntry is a
+// 1000-iteration stress test for the release/acquire race on the shared
+// standaloneStreams registry: releasing the last reference to a stream and
+// acquiring a new one for the same upstream key, at the same instant, must
+// never leave the registry pointing at neither.
+//
+// release goroutine acquire goroutine
+// ------------------ ------------------
+// <-start <-start
+// releaseStandaloneStream(current) next = acquireStandaloneStream(upstream)
+// (refcount-- on current; if it
+// hits 0, evict current from
+// standaloneStreams[key])
+//
+// close(start) releases both goroutines at once, so which one the scheduler
+// runs first varies per iteration: if acquire wins, it finds current still
+// live and reuses it (next == current); if release wins first, it evicts
+// current and acquire must create a fresh entry. Either outcome is fine —
+// the assertion after wg.Wait() is that standaloneStreams[next.key] is
+// always next, never the evicted stream and never nothing.
+func TestStandaloneStreamConcurrentLastReleaseAndAcquireKeepsLiveEntry(t *testing.T) {
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "release-acquire-race", Mode: UpstreamModeHTTP}
+
+ for i := 0; i < 1000; i++ {
+ current := client.acquireStandaloneStream(upstream)
+ start := make(chan struct{})
+ var next *standaloneStream
+ var wg sync.WaitGroup
+ wg.Add(2)
+ go func() {
+ defer wg.Done()
+ <-start
+ client.releaseStandaloneStream(current)
+ }()
+ go func() {
+ defer wg.Done()
+ <-start
+ next = client.acquireStandaloneStream(upstream)
+ }()
+ close(start)
+ wg.Wait()
+
+ client.standaloneMu.Lock()
+ live := client.standaloneStreams[next.key]
+ client.standaloneMu.Unlock()
+ if live != next {
+ t.Fatalf("iteration %d: newly acquired stream was evicted by concurrent release", i)
+ }
+ client.releaseStandaloneStream(next)
+ }
+}
+
+// TestInvokeStreamStandaloneStreamUnsupportedDoesNotFailCall covers an
+// upstream that returns a plain error (e.g. 404/405, which some servers
+// legitimately return for this endpoint per spec) for the standalone GET:
+// the tools/call itself must still succeed normally via its own POST
+// response stream.
+func TestInvokeStreamStandaloneStreamUnsupportedDoesNotFailCall(t *testing.T) {
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ http.Error(w, "not found", http.StatusNotFound)
+ return
+ }
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-unsupported")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := newSyncFakeStreamSink()
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "real", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "tool", map[string]any{}, nil, map[string]any{"progressToken": "tok"}, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+ if len(sink.snapshotEvents()) != 0 {
+ t.Fatalf("expected no relayed events when the standalone stream is unsupported, got %#v", sink.snapshotEvents())
+ }
+}
+
+// TestInvokeStreamStandaloneStreamWrongContentTypeDoesNotFailCall covers the
+// other shape of "unsupported": a 200 response that isn't actually SSE
+// (some servers, on a bare GET, just serve something unrelated rather than
+// the expected 404/405). openStandaloneGET must reject it the same way it
+// rejects an outright error status, without affecting the call itself.
+func TestInvokeStreamStandaloneStreamWrongContentTypeDoesNotFailCall(t *testing.T) {
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if r.Method == http.MethodGet {
+ w.Header().Set("Content-Type", "text/html")
+ _, _ = w.Write([]byte("not an SSE stream"))
+ return
+ }
+ var req Envelope
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("decode request: %v", err)
+ }
+ switch req.Method {
+ case "initialize":
+ w.Header().Set("Mcp-Session-Id", "sid-wrong-content-type")
+ writeTestRPC(w, req.ID, map[string]any{"protocolVersion": r.Header.Get("MCP-Protocol-Version"), "capabilities": map[string]any{}}, nil)
+ case "notifications/initialized":
+ w.WriteHeader(http.StatusAccepted)
+ case "tools/call":
+ w.Header().Set("Content-Type", "text/event-stream")
+ flusher := w.(http.Flusher)
+ writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ default:
+ t.Fatalf("unexpected method %q", req.Method)
+ }
+ }))
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := newSyncFakeStreamSink()
+
+ result, err := client.InvokeStream(context.Background(), Upstream{Name: "real", Mode: UpstreamModeHTTP, BaseURL: server.URL}, "tool", map[string]any{}, nil, map[string]any{"progressToken": "tok"}, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+ if len(sink.snapshotEvents()) != 0 {
+ t.Fatalf("expected no relayed events when the standalone stream has the wrong content type, got %#v", sink.snapshotEvents())
+ }
+}
+
+// TestRouteStandaloneEventFlagsServerRequests pins that a server-to-client
+// request (id + method) routed via the standalone stream carries
+// ServerRequest — the same distinction relaySSEToolCall draws on the POST
+// stream — so audit doesn't mislabel a sampling/elicitation request as a
+// notification just because of which connection carried it.
+func TestRouteStandaloneEventFlagsServerRequests(t *testing.T) {
+ client := NewHTTPClient()
+ events := make(chan StreamEvent, 2)
+ stream := &standaloneStream{waiters: map[string]progressWaiter{
+ "tok": {events: events},
+ }}
+
+ client.routeStandaloneEvent(stream, []byte(`{"jsonrpc":"2.0","id":"srv-9","method":"sampling/createMessage","params":{}}`))
+ client.routeStandaloneEvent(stream, []byte(`{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"hi"}}`))
+
+ request := <-events
+ if !request.ServerRequest {
+ t.Fatalf("expected a standalone-routed server request to be flagged, got %#v", request)
+ }
+ notification := <-events
+ if notification.ServerRequest {
+ t.Fatalf("expected a standalone-routed notification to stay unflagged, got %#v", notification)
+ }
+}
+
+// statsRecordingSink is a fakeStreamSink that also implements
+// StreamStatsSink, so tests can pin the wiring that hands standalone-stream
+// drop accounting to the sink when the call finishes.
+type statsRecordingSink struct {
+ fakeStreamSink
+ statsCalls int
+ lastStats StreamStats
+}
+
+func (s *statsRecordingSink) StreamStats(stats StreamStats) {
+ s.statsCalls++
+ s.lastStats = stats
+}
+
+// TestInvokeStreamDeliversStreamStatsToStatsSink pins the delivery path for
+// standalone-stream statistics: a progressToken-bearing call must report
+// StreamStats exactly once to a sink that implements StreamStatsSink —
+// through the callSink wrapper — before InvokeStream returns. A broken type
+// assertion or defer wiring in invokeHTTPStream would otherwise pass every
+// other test silently.
+func TestInvokeStreamDeliversStreamStatsToStatsSink(t *testing.T) {
+ server := invokeStreamTestServer(t, "sid-stats", func(w http.ResponseWriter, r *http.Request, req Envelope) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ writeTestSSEEventFlush(w, w.(http.Flusher), `{"jsonrpc":"2.0","id":"1","result":{"content":[{"type":"text","text":"done"}]}}`)
+ })
+ defer server.Close()
+
+ client := NewHTTPClient()
+ client.httpClient = server.Client()
+ sink := &statsRecordingSink{}
+
+ _, err := client.InvokeStream(
+ context.Background(),
+ Upstream{Name: "stats", Mode: UpstreamModeHTTP, BaseURL: server.URL},
+ "demo", map[string]any{}, nil,
+ map[string]any{"progressToken": "caller-token"},
+ sink, StreamOptions{},
+ )
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if sink.statsCalls != 1 {
+ t.Fatalf("StreamStats called %d times, want exactly 1", sink.statsCalls)
+ }
+ if sink.lastStats.StandaloneEventsDropped != 0 {
+ t.Fatalf("StandaloneEventsDropped = %d, want 0 (nothing was dropped)", sink.lastStats.StandaloneEventsDropped)
+ }
+}
diff --git a/internal/mcp/stdio_process_unix.go b/internal/mcp/stdio_process_unix.go
new file mode 100644
index 00000000..4f76922d
--- /dev/null
+++ b/internal/mcp/stdio_process_unix.go
@@ -0,0 +1,58 @@
+//go:build !windows
+
+package mcp
+
+import (
+ "os/exec"
+ "syscall"
+ "time"
+)
+
+// configureStdioProcessGroup puts cmd in its own process group (Setpgid) so
+// killStdioProcessGroup can terminate it and any descendants it spawns
+// (e.g. a wrapper script's own child process) in one signal, and sets a
+// bounded WaitDelay so Wait() doesn't hang forever if the kill somehow
+// doesn't land. Without this, canceling the command's context only kills
+// the directly-spawned process: a grandchild that inherited the stdout
+// pipe's write end keeps it open, so the parent's read never sees EOF and
+// a cancellation-based timeout (idle/max-duration) never actually unblocks
+// the read it was supposed to abort.
+func configureStdioProcessGroup(cmd *exec.Cmd) {
+ cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+ cmd.WaitDelay = 2 * time.Second
+ cmd.Cancel = func() error {
+ if cmd.Process == nil {
+ return nil
+ }
+ return killStdioProcessGroup(cmd.Process.Pid)
+ }
+}
+
+// killStdioProcessGroup sends SIGKILL to the process group headed by pid.
+//
+// os/exec's default Cancel (os.Process.Kill) is safe against PID reuse
+// because os.Process tracks internally whether it has already reaped the
+// process, and refuses to signal after that point. Our Cancel bypasses
+// os.Process entirely — a raw syscall.Kill(-pid, ...) is the only way to
+// reach the whole group, not just the one process — so it doesn't get that
+// same protection for free. Getpgid closes most of that gap: since
+// Setpgid made our child its own group leader at spawn time (pgid == pid),
+// verifying that still holds immediately before signaling means we skip
+// the kill if the process has already exited (Getpgid returns ESRCH) or if
+// pid has been recycled by an unrelated process that is not a matching
+// group leader. This narrows the residual race to the syscall gap between
+// the Getpgid check and the Kill call, rather than the much larger window
+// between process exit and this function running.
+func killStdioProcessGroup(pid int) error {
+ pgid, err := syscall.Getpgid(pid)
+ if err != nil {
+ // Already gone, or otherwise unqueryable — nothing safe to kill.
+ return nil
+ }
+ if pgid != pid {
+ // pid no longer heads the process group we spawned it into; refuse
+ // to signal a group we don't recognize.
+ return nil
+ }
+ return syscall.Kill(-pid, syscall.SIGKILL)
+}
diff --git a/internal/mcp/stdio_process_unix_test.go b/internal/mcp/stdio_process_unix_test.go
new file mode 100644
index 00000000..d63ed3dd
--- /dev/null
+++ b/internal/mcp/stdio_process_unix_test.go
@@ -0,0 +1,51 @@
+//go:build !windows
+
+package mcp
+
+import (
+ "context"
+ "os/exec"
+ "testing"
+ "time"
+)
+
+func TestKillStdioProcessGroupTerminatesRunningProcess(t *testing.T) {
+ cmd := exec.CommandContext(context.Background(), "sleep", "30")
+ configureStdioProcessGroup(cmd)
+ if err := cmd.Start(); err != nil {
+ t.Fatalf("start: %v", err)
+ }
+ done := make(chan error, 1)
+ go func() { done <- cmd.Wait() }()
+
+ if err := killStdioProcessGroup(cmd.Process.Pid); err != nil {
+ t.Fatalf("killStdioProcessGroup: %v", err)
+ }
+
+ select {
+ case <-done:
+ case <-time.After(5 * time.Second):
+ t.Fatal("process was not killed within 5s")
+ }
+}
+
+// TestKillStdioProcessGroupIsNoOpForAlreadyExitedProcess proves the
+// Getpgid-verification guard: once a process has exited and been reaped,
+// killStdioProcessGroup must not blindly signal its (now-recycled-eligible)
+// pid. Getpgid on an already-reaped pid returns an error (no such
+// process), so the function returns nil without calling Kill at all.
+func TestKillStdioProcessGroupIsNoOpForAlreadyExitedProcess(t *testing.T) {
+ cmd := exec.CommandContext(context.Background(), "true")
+ configureStdioProcessGroup(cmd)
+ if err := cmd.Start(); err != nil {
+ t.Fatalf("start: %v", err)
+ }
+ pid := cmd.Process.Pid
+ if err := cmd.Wait(); err != nil {
+ t.Fatalf("wait: %v", err)
+ }
+
+ if err := killStdioProcessGroup(pid); err != nil {
+ t.Fatalf("expected a no-op (nil error) for an already-exited pid, got %v", err)
+ }
+}
diff --git a/internal/mcp/stdio_process_windows.go b/internal/mcp/stdio_process_windows.go
new file mode 100644
index 00000000..a095a396
--- /dev/null
+++ b/internal/mcp/stdio_process_windows.go
@@ -0,0 +1,12 @@
+//go:build windows
+
+package mcp
+
+import "os/exec"
+
+// configureStdioProcessGroup is a no-op on Windows: process-group signalling
+// is POSIX-specific (see stdio_process_unix.go). Atryum's release targets
+// are darwin/linux only; this stub exists solely so the package still
+// builds on Windows, at today's level of process cleanup (Cmd's default
+// context-cancellation behavior, which only kills the direct child).
+func configureStdioProcessGroup(cmd *exec.Cmd) {}
diff --git a/internal/mcp/stdio_stream.go b/internal/mcp/stdio_stream.go
new file mode 100644
index 00000000..82421df1
--- /dev/null
+++ b/internal/mcp/stdio_stream.go
@@ -0,0 +1,161 @@
+package mcp
+
+import (
+ "bufio"
+ "bytes"
+ "context"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "os"
+ "os/exec"
+ "strings"
+
+ "github.com/validmind/atryum/internal/version"
+)
+
+// invokeStdioStream is invokeStdio's live-relay counterpart. Stdio has no
+// Content-Type or other signal that selects streaming, so StreamStarted fires
+// only before the first intermediate notification or server request. A
+// terminal-only response never touches the sink.
+func (c *Client) invokeStdioStream(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string, meta map[string]any, sink StreamSink, opts StreamOptions) (InvokeResult, error) {
+ if upstream.Command == "" {
+ return InvokeResult{}, fmt.Errorf("stdio upstream %q missing command", upstream.Name)
+ }
+ guard := newCallTimeoutGuard(ctx)
+ defer guard.stop()
+
+ cmd := exec.CommandContext(guard.ctx, upstream.Command, upstream.Args...)
+ cmd.Env = os.Environ()
+ for k, v := range upstream.Env {
+ cmd.Env = append(cmd.Env, k+"="+v)
+ }
+ configureStdioProcessGroup(cmd)
+ stdin, err := cmd.StdinPipe()
+ if err != nil {
+ return InvokeResult{}, err
+ }
+ stdout, err := cmd.StdoutPipe()
+ if err != nil {
+ return InvokeResult{}, err
+ }
+ stderr := newBoundedBuffer(stdioStderrCap)
+ cmd.Stderr = stderr
+ if err := cmd.Start(); err != nil {
+ return InvokeResult{}, err
+ }
+ defer func() {
+ // guard.stop() MUST run before cmd.Wait(): stopping cancels the
+ // guard context, which triggers the process-group kill, which is
+ // what makes Wait return. Deferring these separately would run
+ // them in LIFO order — Wait before stop — and a stdio server that
+ // keeps running after answering (normal for long-lived servers)
+ // or that ignores stdin-close would then block Wait forever on
+ // every return path where no timeout had fired (sink error, or a
+ // successfully received terminal response). The process is
+ // per-call and disposable, so killing it once we have our answer
+ // (or have given up) is the correct lifecycle.
+ guard.stop()
+ _ = stdin.Close()
+ _ = cmd.Wait()
+ }()
+
+ reader := bufio.NewReader(stdout)
+ // The initialize handshake gets the same header-phase bound the HTTP
+ // path applies before its response headers arrive: without it, a
+ // subprocess that starts but never answers initialize would block
+ // readRPC with no bound of its own (only the caller's ctx).
+ guard.armSetupTimeout(opts.HeaderTimeout)
+ initID := c.nextRPCID()
+ if err := writeRPC(stdin, initID, "initialize", map[string]any{
+ "protocolVersion": DefaultMCPProtocolVersion,
+ "clientInfo": map[string]any{"name": "atryum", "version": version.Version},
+ "capabilities": map[string]any{},
+ }); err != nil {
+ return InvokeResult{}, err
+ }
+ if _, err := readRPCWithLimit(reader, rpcIDMessage(initID), opts.maxMessageBytes()); err != nil {
+ if timeoutErr := guard.timeoutErr(upstream.Name, "during stdio initialize", nil); timeoutErr != nil {
+ return InvokeResult{}, timeoutErr
+ }
+ if stderr.Len() > 0 {
+ return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String()))
+ }
+ return InvokeResult{}, err
+ }
+ guard.disarmSetupTimeout()
+ _ = writeRPC(stdin, c.nextRPCID(), "notifications/initialized", map[string]any{})
+ callParams := map[string]any{"name": tool, "arguments": input}
+ if merged := mergeRequestMeta(meta, requestID); merged != nil {
+ callParams["_meta"] = merged
+ }
+ callID := c.nextRPCID()
+ if err := writeRPC(stdin, callID, "tools/call", callParams); err != nil {
+ return InvokeResult{}, err
+ }
+
+ guard.armBodyTimeouts(opts.IdleTimeout, opts.MaxDuration)
+ return c.relayStdioToolCall(reader, sink, guard, upstream, callID, stderr, opts.maxMessageBytes())
+}
+
+// relayStdioToolCall reads reader's newline-delimited JSON-RPC messages,
+// relaying every intermediate (non-terminal) message to sink as it arrives,
+// and returns once the terminal response for callID is read.
+func (c *Client) relayStdioToolCall(reader *bufio.Reader, sink StreamSink, guard *callTimeoutGuard, upstream Upstream, callID int64, stderr *boundedBuffer, maxMessageBytes int) (InvokeResult, error) {
+ expectedID := rpcIDMessage(callID)
+ started := false
+ ensureStarted := func() {
+ if !started {
+ started = true
+ sink.StreamStarted()
+ }
+ }
+ for {
+ line, err := readLineLimited(reader, maxMessageBytes)
+ if err != nil {
+ if timeoutErr := guard.timeoutErr(upstream.Name, "", nil); timeoutErr != nil {
+ return InvokeResult{}, timeoutErr
+ }
+ if stderr.Len() > 0 {
+ return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String()))
+ }
+ return InvokeResult{}, err
+ }
+ // Any line — including a blank one — is upstream liveness, matching
+ // the HTTP relay's per-line idle reset (SSE comments count there).
+ guard.resetIdle()
+ line = bytes.TrimSpace(line)
+ if len(line) == 0 {
+ continue
+ }
+
+ switch classifyRPCMessage(line, expectedID) {
+ case rpcMessageTerminalResponse:
+ var resp rpcResponse
+ if err := json.Unmarshal(line, &resp); err != nil {
+ continue
+ }
+ if len(resp.Error) > 0 && string(resp.Error) != "null" {
+ return InvokeResult{StatusCode: http.StatusBadGateway, Body: resp.Error, Failed: true}, nil
+ }
+ body := resp.Result
+ if len(body) == 0 {
+ body = []byte(`{"ok":true}`)
+ }
+ return InvokeResult{StatusCode: http.StatusOK, Body: body, Failed: looksLikeToolError(body)}, nil
+ case rpcMessageServerRequest:
+ ensureStarted()
+ if err := sink.Event(StreamEvent{Data: line, ServerRequest: true}); err != nil {
+ return InvokeResult{}, err
+ }
+ case rpcMessageNotification:
+ ensureStarted()
+ if err := sink.Event(StreamEvent{Data: line}); err != nil {
+ return InvokeResult{}, err
+ }
+ default:
+ // Unparseable line, or a response to some other id. Not ours
+ // to interpret; ignore and keep reading.
+ }
+ }
+}
diff --git a/internal/mcp/stdio_stream_test.go b/internal/mcp/stdio_stream_test.go
new file mode 100644
index 00000000..0a7e0999
--- /dev/null
+++ b/internal/mcp/stdio_stream_test.go
@@ -0,0 +1,277 @@
+//go:build !windows
+
+package mcp
+
+import (
+ "context"
+ "errors"
+ "os"
+ "path/filepath"
+ "strings"
+ "testing"
+ "time"
+)
+
+// writeFakeStdioServer creates a one-call MCP subprocess for the stdio tests.
+func writeFakeStdioServer(t *testing.T, toolsCallScript string) string {
+ t.Helper()
+ path := filepath.Join(t.TempDir(), "fake-mcp.sh")
+ content := "#!/usr/bin/env bash\n" +
+ "set -euo pipefail\n" +
+ "while IFS= read -r line; do\n" +
+ " if [[ -z \"$line\" ]]; then continue; fi\n" +
+ " if [[ \"$line\" == *'\"method\":\"initialize\"'* ]]; then\n" +
+ " id=$(echo \"$line\" | grep -o '\"id\":[0-9]*' | head -1 | cut -d: -f2)\n" +
+ " printf '%s\\n' \"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"id\\\":${id},\\\"result\\\":{\\\"serverInfo\\\":{\\\"name\\\":\\\"fake\\\",\\\"version\\\":\\\"0.1.0\\\"},\\\"capabilities\\\":{}}}\"\n" +
+ " elif [[ \"$line\" == *'\"method\":\"notifications/initialized\"'* ]]; then\n" +
+ " continue\n" +
+ " elif [[ \"$line\" == *'\"method\":\"tools/call\"'* ]]; then\n" +
+ toolsCallScript +
+ " exit 0\n" +
+ " fi\n" +
+ "done\n"
+ if err := os.WriteFile(path, []byte(content), 0o755); err != nil {
+ t.Fatal(err)
+ }
+ return path
+}
+
+func TestInvokeStreamStdioRelaysEventsBeforeTerminalResponseExists(t *testing.T) {
+ releaseFile := filepath.Join(t.TempDir(), "release")
+ script := writeFakeStdioServer(t, ""+
+ " id=$(echo \"$line\" | grep -o '\"id\":[0-9]*' | head -1 | cut -d: -f2)\n"+
+ " printf '%s\\n' '{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}'\n"+
+ " while [ ! -f \"$RELEASE_FILE\" ]; do sleep 0.02; done\n"+
+ " printf '%s\\n' \"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"id\\\":${id},\\\"result\\\":{\\\"content\\\":[{\\\"type\\\":\\\"text\\\",\\\"text\\\":\\\"done\\\"}]}}\"\n",
+ )
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: script, Env: map[string]string{"RELEASE_FILE": releaseFile}}
+ sink := &fakeStreamSink{onEvent: func(StreamEvent) error {
+ // The subprocess is blocked in its own `while [ ! -f ... ]` loop and
+ // cannot write the terminal response until this file exists — it
+ // only gets created here, inside the callback fired once the client
+ // has actually delivered the notification to the sink.
+ return os.WriteFile(releaseFile, []byte("go"), 0o644)
+ }}
+
+ result, err := client.InvokeStream(context.Background(), upstream, "demo", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !sink.started {
+ t.Fatal("expected StreamStarted to fire")
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected exactly one relayed event, got %d", len(sink.events))
+ }
+ if !strings.Contains(string(sink.events[0].Data), "notifications/progress") {
+ t.Fatalf("expected progress notification relayed, got %s", sink.events[0].Data)
+ }
+ if !strings.Contains(string(result.Body), "done") {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+}
+
+func TestInvokeStreamStdioTerminalOnlyResponseNeverTouchesSink(t *testing.T) {
+ script := writeFakeStdioServer(t, ""+
+ " id=$(echo \"$line\" | grep -o '\"id\":[0-9]*' | head -1 | cut -d: -f2)\n"+
+ " printf '%s\\n' \"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"id\\\":${id},\\\"result\\\":{\\\"content\\\":[{\\\"type\\\":\\\"text\\\",\\\"text\\\":\\\"ok\\\"}]}}\"\n",
+ )
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: script}
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(context.Background(), upstream, "demo", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if sink.started || len(sink.events) != 0 {
+ t.Fatalf("expected sink to never be touched when the upstream emits nothing but its terminal response, got started=%t events=%#v", sink.started, sink.events)
+ }
+ if !strings.Contains(string(result.Body), `"text":"ok"`) {
+ t.Fatalf("expected terminal result body, got %s", result.Body)
+ }
+}
+
+func TestInvokeStreamStdioRejectsOversizedMessage(t *testing.T) {
+ script := writeFakeStdioServer(t, ""+
+ " printf '%0200d\\n' 0\n",
+ )
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: script}
+ _, err := client.InvokeStream(
+ context.Background(),
+ upstream,
+ "demo",
+ map[string]any{},
+ nil,
+ nil,
+ &fakeStreamSink{},
+ StreamOptions{MaxMessageBytes: 128},
+ )
+ if !errors.Is(err, ErrStreamMessageTooLarge) {
+ t.Fatalf("InvokeStream error = %v, want ErrStreamMessageTooLarge", err)
+ }
+}
+
+func TestInvokeStreamStdioTerminalErrorAfterNotification(t *testing.T) {
+ script := writeFakeStdioServer(t, ""+
+ " printf '%s\\n' '{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}'\n"+
+ " id=$(echo \"$line\" | grep -o '\"id\":[0-9]*' | head -1 | cut -d: -f2)\n"+
+ " printf '%s\\n' \"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"id\\\":${id},\\\"error\\\":{\\\"code\\\":-32000,\\\"message\\\":\\\"tool exploded\\\"}}\"\n",
+ )
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: script}
+ sink := &fakeStreamSink{}
+
+ result, err := client.InvokeStream(context.Background(), upstream, "demo", map[string]any{}, nil, nil, sink, StreamOptions{})
+ if err != nil {
+ t.Fatalf("InvokeStream returned error: %v", err)
+ }
+ if !result.Failed {
+ t.Fatalf("expected Failed result, got %#v", result)
+ }
+ if !strings.Contains(string(result.Body), "tool exploded") {
+ t.Fatalf("expected error body, got %s", result.Body)
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected the progress notification to have been relayed before the terminal error, got %d", len(sink.events))
+ }
+}
+
+func TestInvokeStreamStdioIdleTimeoutAbortsRead(t *testing.T) {
+ script := writeFakeStdioServer(t, ""+
+ " printf '%s\\n' '{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}'\n"+
+ " sleep 30\n", // never sends the terminal event; killed by the idle timeout well before this returns
+ )
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: script}
+ sink := &fakeStreamSink{}
+
+ start := time.Now()
+ _, err := client.InvokeStream(context.Background(), upstream, "demo", map[string]any{}, nil, nil, sink, StreamOptions{IdleTimeout: 50 * time.Millisecond})
+ elapsed := time.Since(start)
+
+ if err == nil {
+ t.Fatal("expected an idle timeout error")
+ }
+ if !errors.Is(err, ErrStreamTimeout) {
+ t.Fatalf("expected errors.Is(err, ErrStreamTimeout) to hold, got %v", err)
+ }
+ if elapsed > 5*time.Second {
+ t.Fatalf("idle timeout took too long to abort: %s", elapsed)
+ }
+ if len(sink.events) != 1 {
+ t.Fatalf("expected exactly one relayed event before the timeout, got %d", len(sink.events))
+ }
+}
+
+// TestInvokeStreamStdioSinkErrorDoesNotDeadlockOnPersistentServer is a
+// regression test for the cleanup-defer ordering: when the sink aborts the
+// relay (agent disconnected) with no timeout having fired, the cleanup
+// defer runs cmd.Wait() — and a stdio server that keeps running (its outer
+// read loop is stuck inside an inner emit loop, so it never notices
+// stdin-close) would block Wait forever unless guard.stop() cancels the
+// context (killing the process group) BEFORE the Wait. With separate
+// defers in the natural order, LIFO ran Wait first — a deadlock this test
+// would catch by hanging.
+func TestInvokeStreamStdioSinkErrorDoesNotDeadlockOnPersistentServer(t *testing.T) {
+ script := writeFakeStdioServer(t, ""+
+ " while true; do\n"+
+ " printf '%s\\n' '{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":1}}'\n"+
+ " sleep 0.05\n"+
+ " done\n",
+ )
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: script}
+ sink := &fakeStreamSink{onEvent: func(StreamEvent) error {
+ return errors.New("downstream connection closed")
+ }}
+
+ start := time.Now()
+ _, err := client.InvokeStream(context.Background(), upstream, "demo", map[string]any{}, nil, nil, sink, StreamOptions{})
+ elapsed := time.Since(start)
+
+ if err == nil {
+ t.Fatal("expected the sink's error to be returned")
+ }
+ if !strings.Contains(err.Error(), "downstream connection closed") {
+ t.Fatalf("expected the sink's error, got %v", err)
+ }
+ if elapsed > 5*time.Second {
+ t.Fatalf("expected a prompt return after the sink aborted (process-group kill before Wait), took %s", elapsed)
+ }
+}
+
+// TestInvokeStreamStdioHandshakeHangBoundedByHeaderTimeout is a regression
+// test: the stdio initialize handshake happens before body timeouts are
+// armed, so it needs the header-phase bound — without it, a subprocess
+// that starts but never answers initialize blocks the call with no bound
+// of its own.
+func TestInvokeStreamStdioHandshakeHangBoundedByHeaderTimeout(t *testing.T) {
+ path := filepath.Join(t.TempDir(), "hang-mcp.sh")
+ content := "#!/usr/bin/env bash\n" +
+ "set -euo pipefail\n" +
+ "while IFS= read -r line; do\n" +
+ " if [[ \"$line\" == *'\"method\":\"initialize\"'* ]]; then\n" +
+ " sleep 30\n" + // never answers initialize; killed by the header timeout
+ " fi\n" +
+ "done\n"
+ if err := os.WriteFile(path, []byte(content), 0o755); err != nil {
+ t.Fatal(err)
+ }
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: path}
+ sink := &fakeStreamSink{}
+
+ start := time.Now()
+ _, err := client.InvokeStream(context.Background(), upstream, "demo", map[string]any{}, nil, nil, sink, StreamOptions{HeaderTimeout: 50 * time.Millisecond})
+ elapsed := time.Since(start)
+
+ if err == nil {
+ t.Fatal("expected a handshake timeout error")
+ }
+ if !errors.Is(err, ErrStreamTimeout) {
+ t.Fatalf("expected errors.Is(err, ErrStreamTimeout) for the hung handshake, got %v", err)
+ }
+ if elapsed > 5*time.Second {
+ t.Fatalf("handshake timeout took too long to abort: %s", elapsed)
+ }
+ if sink.started || len(sink.events) != 0 {
+ t.Fatalf("expected the sink never to be touched during a failed handshake, got started=%t events=%d", sink.started, len(sink.events))
+ }
+}
+
+// TestInvokeStdioSkipsStrayServerRequestBeforeTerminalResponse is a
+// regression test for the readRPC correctness fix: a "first message with
+// any id/result/error wins" reader would misinterpret a stray incoming
+// server-to-client request (it has an id, but no result/error — readRPC's
+// old check only looked for "any of id/result/error present") as the
+// answer to our own call, before the real response ever arrives. This uses
+// the plain buffered Invoke (not InvokeStream) since the fix applies
+// there too.
+func TestInvokeStdioSkipsStrayServerRequestBeforeTerminalResponse(t *testing.T) {
+ script := writeFakeStdioServer(t, ""+
+ " printf '%s\\n' '{\"jsonrpc\":\"2.0\",\"id\":\"srv-1\",\"method\":\"sampling/createMessage\",\"params\":{}}'\n"+
+ " id=$(echo \"$line\" | grep -o '\"id\":[0-9]*' | head -1 | cut -d: -f2)\n"+
+ " printf '%s\\n' \"{\\\"jsonrpc\\\":\\\"2.0\\\",\\\"id\\\":${id},\\\"result\\\":{\\\"content\\\":[{\\\"type\\\":\\\"text\\\",\\\"text\\\":\\\"real answer\\\"}]}}\"\n",
+ )
+
+ client := NewHTTPClient()
+ upstream := Upstream{Name: "fake-stdio", Mode: UpstreamModeStdio, Command: script}
+
+ result, err := client.Invoke(context.Background(), upstream, "demo", map[string]any{}, nil, nil)
+ if err != nil {
+ t.Fatalf("Invoke returned error: %v", err)
+ }
+ if !strings.Contains(string(result.Body), "real answer") {
+ t.Fatalf("expected the real terminal response (readRPC must skip the stray server-to-client request), got %s", result.Body)
+ }
+}
diff --git a/internal/mcp/stream.go b/internal/mcp/stream.go
new file mode 100644
index 00000000..7fd145e1
--- /dev/null
+++ b/internal/mcp/stream.go
@@ -0,0 +1,167 @@
+package mcp
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "time"
+)
+
+const defaultStreamMaxMessageBytes = 4 * 1024 * 1024
+
+// ErrStreamMessageTooLarge marks an upstream response whose encoded JSON-RPC
+// message exceeds StreamOptions.MaxMessageBytes.
+var ErrStreamMessageTooLarge = errors.New("stream message exceeds configured byte limit")
+
+// StreamEvent is one intermediate upstream JSON-RPC message, independent of
+// whether HTTP SSE or stdio carried it. It is either a notification (progress,
+// logging, or another server-to-client notification) or, more rarely, a
+// server-to-client request.
+type StreamEvent struct {
+ // Data is one raw JSON-RPC message. HTTP SSE joins the event's data lines
+ // with newlines; stdio removes its newline framing.
+ Data []byte
+ // ServerRequest is true when Data is a JSON-RPC request from the
+ // upstream (has both id and method) rather than a notification. Atryum
+ // does not broker server-initiated requests (sampling, elicitation,
+ // roots); these are surfaced to the sink for audit only, never relayed
+ // to the agent.
+ ServerRequest bool
+}
+
+// StreamStats contains transport observations that are not themselves
+// relayed JSON-RPC messages.
+type StreamStats struct {
+ StandaloneEventsDropped int64
+}
+
+// StreamStatsSink is an optional extension implemented by sinks that record
+// transport-level statistics. HTTP calls using the standalone progress path
+// call it once before return.
+type StreamStatsSink interface {
+ StreamStats(stats StreamStats)
+}
+
+// StreamSink receives intermediate upstream messages live, as InvokeStream
+// reads them, so a caller can relay them onward (or just audit them) before
+// the terminal response exists. Its methods run synchronously on the same
+// goroutine as the InvokeStream call — there is no concurrent access to the
+// sink, and no need for the sink to synchronize internally on that account.
+type StreamSink interface {
+ // StreamStarted fires at most once. HTTP SSE calls it before the first
+ // event or terminal response; stdio calls it only before the first
+ // intermediate event. A silently retried attempt never calls it.
+ StreamStarted()
+ // Event delivers one intermediate (non-terminal) message. A returned
+ // error aborts the stream: InvokeStream stops reading and returns that
+ // error to its caller.
+ Event(evt StreamEvent) error
+}
+
+// StreamOptions bounds InvokeStream's setup and response-reading phases. A
+// zero-valued field disables that particular bound.
+type StreamOptions struct {
+ // HeaderTimeout bounds setup before tool response reading begins: HTTP
+ // session initialization and response headers, or the stdio initialize
+ // handshake. Zero leaves setup bounded only by ctx's deadline, if any.
+ HeaderTimeout time.Duration
+ // IdleTimeout bounds response-reading inactivity. Streaming transports
+ // reset it when upstream activity arrives: events routed over the shared
+ // standalone HTTP stream, stdio lines, and any SSE line on the call's
+ // own response — including comment (":keepalive") lines, so an upstream
+ // heartbeating through a long tool run is not treated as idle. For a
+ // plain HTTP JSON response it bounds the complete body read. Zero
+ // disables the check.
+ IdleTimeout time.Duration
+ // MaxDuration bounds the complete response-reading phase after HTTP
+ // headers or the stdio handshake. Zero disables the check.
+ MaxDuration time.Duration
+ // MaxMessageBytes bounds one decoded SSE event, one stdio JSON-RPC line,
+ // or one plain JSON response body. Zero uses the safe 4 MiB default.
+ MaxMessageBytes int
+}
+
+func (o StreamOptions) maxMessageBytes() int {
+ if o.MaxMessageBytes > 0 {
+ return o.MaxMessageBytes
+ }
+ return defaultStreamMaxMessageBytes
+}
+
+// InvokeStream behaves like Invoke while also relaying intermediate JSON-RPC
+// messages to sink as they arrive. HTTP upstreams select streaming with an
+// SSE response, so StreamStarted fires even when that response contains only
+// its terminal message. Stdio starts the sink only when an intermediate
+// message arrives. A nil sink always uses Invoke's buffered path.
+func (c *Client) InvokeStream(ctx context.Context, upstream Upstream, tool string, input map[string]any, requestID *string, meta map[string]any, sink StreamSink, opts StreamOptions) (InvokeResult, error) {
+ switch upstream.Mode {
+ case UpstreamModeStdio:
+ if sink == nil {
+ return c.Invoke(ctx, upstream, tool, input, requestID, meta)
+ }
+ started := time.Now()
+ defer func() {
+ c.debugf("upstream invoke-stream transport=%s server=%s tool=%s duration_ms=%d", upstream.Mode, upstream.Name, tool, time.Since(started).Milliseconds())
+ }()
+ return c.invokeStdioStream(ctx, upstream, tool, input, requestID, meta, sink, opts)
+ case UpstreamModeHTTP, "":
+ if sink == nil {
+ return c.Invoke(ctx, upstream, tool, input, requestID, meta)
+ }
+ started := time.Now()
+ defer func() {
+ c.debugf("upstream invoke-stream transport=%s server=%s tool=%s duration_ms=%d", upstream.Mode, upstream.Name, tool, time.Since(started).Milliseconds())
+ }()
+ return c.invokeHTTPStream(ctx, upstream, tool, input, requestID, meta, sink, opts)
+ default:
+ return InvokeResult{}, fmt.Errorf("unsupported upstream mode %q", upstream.Mode)
+ }
+}
+
+type rpcMessageKind int
+
+const (
+ rpcMessageUnknown rpcMessageKind = iota
+ rpcMessageTerminalResponse
+ rpcMessageNotification
+ rpcMessageServerRequest
+)
+
+// classifyRPCMessage identifies one already-parsed JSON-RPC message for the
+// streaming relay: the terminal response to our request (matches
+// expectedID, or is the null-id error the JSON-RPC spec uses when a server
+// can't identify which request an error belongs to), a notification (no id),
+// a server-to-client request (id and method, no result/error), or unknown
+// (e.g. a response to some other id — not ours to interpret). Transport-
+// neutral: used for both the HTTP SSE relay (relaySSEToolCall) and the
+// stdio relay (relayStdioToolCall), and for stdio's non-streaming readRPC,
+// since a JSON-RPC message's shape doesn't depend on how it was framed on
+// the wire.
+func classifyRPCMessage(payload []byte, expectedID json.RawMessage) rpcMessageKind {
+ var message map[string]json.RawMessage
+ if err := json.Unmarshal(payload, &message); err != nil {
+ return rpcMessageUnknown
+ }
+ id, hasID := message["id"]
+ _, hasMethod := message["method"]
+ _, hasResult := message["result"]
+ _, hasError := message["error"]
+
+ if hasID && (hasResult || hasError) {
+ if hasError && jsonRawIsNull(id) {
+ return rpcMessageTerminalResponse
+ }
+ if jsonRPCIDsMatch(id, expectedID) {
+ return rpcMessageTerminalResponse
+ }
+ return rpcMessageUnknown
+ }
+ if hasID && hasMethod {
+ return rpcMessageServerRequest
+ }
+ if !hasID && hasMethod {
+ return rpcMessageNotification
+ }
+ return rpcMessageUnknown
+}
diff --git a/internal/mcp/stream_timeout.go b/internal/mcp/stream_timeout.go
new file mode 100644
index 00000000..3eef1fec
--- /dev/null
+++ b/internal/mcp/stream_timeout.go
@@ -0,0 +1,167 @@
+package mcp
+
+import (
+ "context"
+ "fmt"
+ "sync"
+ "sync/atomic"
+ "time"
+)
+
+// callTimeoutGuard implements InvokeStream's setup/idle/max-duration timeout
+// scheme by canceling one shared context. The setup timer covers HTTP response
+// headers or the stdio initialize handshake. After setup, callers replace it
+// with idle and maximum-duration timers for response reading.
+type callTimeoutGuard struct {
+ ctx context.Context
+ cancel context.CancelFunc
+
+ // mu guards trippedWhy, stopped, the timer fields, and idleTimeout.
+ // The timer fields need it because a time.AfterFunc callback starts
+ // its clock before the assignment of the returned *Timer completes:
+ // checkIdle (running on the timer's goroutine) could otherwise read
+ // g.idleTimer before/while armBodyTimeouts writes it — a data race by
+ // the memory model even if the window is nanoseconds in practice.
+ mu sync.Mutex
+ trippedWhy string
+ stopped bool
+ setupTimer *time.Timer
+ idleTimer *time.Timer
+ maxTimer *time.Timer
+ idleTimeout time.Duration
+
+ // lastActivity (unix nanoseconds) is updated by resetIdle and read by
+ // checkIdle. It exists so the idle timer's firing can be verified
+ // rather than trusted outright — see checkIdle. Atomic, not mu-guarded:
+ // resetIdle runs once per relayed event on the hot path and must not
+ // contend with the timer goroutine.
+ lastActivity atomic.Int64
+}
+
+func newCallTimeoutGuard(parent context.Context) *callTimeoutGuard {
+ ctx, cancel := context.WithCancel(parent)
+ return &callTimeoutGuard{ctx: ctx, cancel: cancel}
+}
+
+func (g *callTimeoutGuard) trip(why string) {
+ g.mu.Lock()
+ if g.trippedWhy == "" {
+ g.trippedWhy = why
+ }
+ g.mu.Unlock()
+ g.cancel()
+}
+
+func (g *callTimeoutGuard) armSetupTimeout(d time.Duration) {
+ if d <= 0 {
+ return
+ }
+ g.mu.Lock()
+ defer g.mu.Unlock()
+ if g.stopped {
+ return
+ }
+ g.setupTimer = time.AfterFunc(d, func() { g.trip("stream setup timeout exceeded") })
+}
+
+func (g *callTimeoutGuard) disarmSetupTimeout() {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+ if g.setupTimer != nil {
+ g.setupTimer.Stop()
+ }
+}
+
+func (g *callTimeoutGuard) armBodyTimeouts(idle, max time.Duration) {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+ if g.stopped {
+ return
+ }
+ g.idleTimeout = idle
+ if idle > 0 {
+ g.lastActivity.Store(time.Now().UnixNano())
+ g.idleTimer = time.AfterFunc(idle, g.checkIdle)
+ }
+ if max > 0 {
+ g.maxTimer = time.AfterFunc(max, func() { g.trip("max stream duration exceeded") })
+ }
+}
+
+// checkIdle is the idle timer's callback. It does not trust "the timer
+// fired" to mean "genuinely idle": time.Timer.Reset called concurrently
+// with a timer's own firing is explicitly documented as racy (the AfterFunc
+// callback may already be running by the time Reset takes effect), so
+// resetIdle deliberately never calls Reset at all — it only records the
+// latest activity timestamp. checkIdle re-derives the real elapsed time
+// from that timestamp and either trips (elapsed genuinely exceeds the
+// bound) or reschedules for the remaining time (an event arrived
+// concurrently with this firing). This makes the idle bound correct
+// regardless of how resetIdle and the timer callback interleave. The
+// stopped check makes a firing that lost the race with stop() a no-op
+// instead of re-arming a timer the guard's owner believes is dead.
+func (g *callTimeoutGuard) checkIdle() {
+ g.mu.Lock()
+ if g.stopped {
+ g.mu.Unlock()
+ return
+ }
+ idleTimeout := g.idleTimeout
+ elapsed := time.Duration(time.Now().UnixNano() - g.lastActivity.Load())
+ if elapsed < idleTimeout {
+ if g.idleTimer != nil {
+ g.idleTimer.Reset(idleTimeout - elapsed)
+ }
+ g.mu.Unlock()
+ return
+ }
+ g.mu.Unlock()
+ // trip acquires g.mu itself; called outside the lock.
+ g.trip("idle timeout waiting for the next stream event")
+}
+
+func (g *callTimeoutGuard) resetIdle() {
+ g.lastActivity.Store(time.Now().UnixNano())
+}
+
+// stop is idempotent: the guard's owners defer it both at guard creation
+// (covering early-error returns) and inside subprocess-cleanup defers that
+// must cancel the context before waiting on the process.
+func (g *callTimeoutGuard) stop() {
+ g.mu.Lock()
+ g.stopped = true
+ if g.setupTimer != nil {
+ g.setupTimer.Stop()
+ }
+ if g.idleTimer != nil {
+ g.idleTimer.Stop()
+ }
+ if g.maxTimer != nil {
+ g.maxTimer.Stop()
+ }
+ g.mu.Unlock()
+ g.cancel()
+}
+
+func (g *callTimeoutGuard) reason() string {
+ g.mu.Lock()
+ defer g.mu.Unlock()
+ return g.trippedWhy
+}
+
+// timeoutErr distinguishes "the guard's own bound fired" from "the transport
+// failed on its own", a check every streaming read path needs after a read
+// error: if the guard tripped, it wraps upstreamName and the trip reason
+// (plus situation, a short phrase like "while resuming" — or "" for none) as
+// ErrStreamTimeout; otherwise it returns fallback unchanged, letting the
+// caller pass nil when it still has its own fallback logic to run.
+func (g *callTimeoutGuard) timeoutErr(upstreamName, situation string, fallback error) error {
+ reason := g.reason()
+ if reason == "" {
+ return fallback
+ }
+ if situation == "" {
+ return fmt.Errorf("upstream %q: %s: %w", upstreamName, reason, ErrStreamTimeout)
+ }
+ return fmt.Errorf("upstream %q %s %s: %w", upstreamName, reason, situation, ErrStreamTimeout)
+}
diff --git a/internal/mcp/stream_timeout_test.go b/internal/mcp/stream_timeout_test.go
new file mode 100644
index 00000000..b7e1cc06
--- /dev/null
+++ b/internal/mcp/stream_timeout_test.go
@@ -0,0 +1,112 @@
+package mcp
+
+import (
+ "context"
+ "strings"
+ "testing"
+ "time"
+)
+
+func TestCallTimeoutGuardIdleAndMaxDuration(t *testing.T) {
+ idle := newCallTimeoutGuard(context.Background())
+ defer idle.stop()
+ idle.armBodyTimeouts(10*time.Millisecond, 0)
+ <-idle.ctx.Done()
+ if reason := idle.reason(); !strings.Contains(reason, "idle timeout") {
+ t.Fatalf("expected idle timeout reason, got %q", reason)
+ }
+
+ max := newCallTimeoutGuard(context.Background())
+ defer max.stop()
+ max.armBodyTimeouts(0, 10*time.Millisecond)
+ <-max.ctx.Done()
+ if reason := max.reason(); !strings.Contains(reason, "max stream duration") {
+ t.Fatalf("expected max duration reason, got %q", reason)
+ }
+}
+
+func TestCallTimeoutGuardSetupTimeoutCanBeDisarmed(t *testing.T) {
+ g := newCallTimeoutGuard(context.Background())
+ defer g.stop()
+ g.armSetupTimeout(10 * time.Millisecond)
+ g.disarmSetupTimeout()
+ time.Sleep(30 * time.Millisecond)
+ if reason := g.reason(); reason != "" {
+ t.Fatalf("expected a disarmed setup timeout not to fire, got %q", reason)
+ }
+}
+
+// TestCallTimeoutGuardCheckIdleReschedulesOnRecentActivity is a
+// deterministic regression test for the idle-timer reset race: time.Timer's
+// docs explicitly warn that Reset racing with the timer's own firing is
+// unsafe to reason about naively (the AfterFunc callback may already be
+// running by the time Reset takes effect). checkIdle closes that race by
+// re-deriving real elapsed time from lastActivity instead of trusting that
+// "the timer fired" means "genuinely idle". This calls checkIdle directly
+// with a lastActivity timestamp from a moment ago — simulating the timer
+// firing at the exact instant resetIdle recorded fresh activity — and
+// verifies it reschedules rather than tripping.
+func TestCallTimeoutGuardCheckIdleReschedulesOnRecentActivity(t *testing.T) {
+ g := newCallTimeoutGuard(context.Background())
+ defer g.stop()
+ g.idleTimeout = 100 * time.Millisecond
+ g.lastActivity.Store(time.Now().UnixNano())
+ g.idleTimer = time.NewTimer(time.Hour) // dummy target for checkIdle's Reset call
+
+ g.checkIdle()
+
+ if reason := g.reason(); reason != "" {
+ t.Fatalf("expected checkIdle to reschedule (not trip) when real elapsed time is well under idleTimeout, got %q", reason)
+ }
+}
+
+func TestCallTimeoutGuardCheckIdleTripsWhenElapsedExceedsTimeout(t *testing.T) {
+ g := newCallTimeoutGuard(context.Background())
+ defer g.stop()
+ g.idleTimeout = 10 * time.Millisecond
+ g.lastActivity.Store(time.Now().Add(-time.Hour).UnixNano())
+
+ g.checkIdle()
+
+ if reason := g.reason(); !strings.Contains(reason, "idle timeout") {
+ t.Fatalf("expected checkIdle to trip when elapsed time genuinely exceeds idleTimeout, got %q", reason)
+ }
+}
+
+// TestCallTimeoutGuardCheckIdleIsNoOpAfterStop is a regression test for the
+// stop/checkIdle race: a checkIdle firing that loses the race with stop()
+// must neither trip the guard nor re-arm the timer. Simulated directly by
+// calling checkIdle after stop() with an ancient lastActivity — without
+// the stopped check it would trip.
+func TestCallTimeoutGuardCheckIdleIsNoOpAfterStop(t *testing.T) {
+ g := newCallTimeoutGuard(context.Background())
+ g.armBodyTimeouts(time.Hour, 0)
+ g.lastActivity.Store(time.Now().Add(-2 * time.Hour).UnixNano())
+ g.stop()
+
+ g.checkIdle()
+
+ if reason := g.reason(); reason != "" {
+ t.Fatalf("expected checkIdle after stop to be a no-op, got %q", reason)
+ }
+}
+
+// TestCallTimeoutGuardSurvivesContinuousResetIdlePressure is a stress test:
+// hammering resetIdle from a tight loop must never spuriously trip the
+// idle timer, even though the timer's own firing schedule and the reset
+// calls are running on different goroutines with no shared lock between
+// them (by design — resetIdle only writes an atomic timestamp).
+func TestCallTimeoutGuardSurvivesContinuousResetIdlePressure(t *testing.T) {
+ g := newCallTimeoutGuard(context.Background())
+ defer g.stop()
+ g.armBodyTimeouts(5*time.Millisecond, 0)
+
+ deadline := time.Now().Add(200 * time.Millisecond)
+ for time.Now().Before(deadline) {
+ g.resetIdle()
+ }
+
+ if reason := g.reason(); reason != "" {
+ t.Fatalf("expected the idle timer never to trip while resetIdle is called continuously, got %q", reason)
+ }
+}
diff --git a/pkg/atryum/atryum.go b/pkg/atryum/atryum.go
index d49bb81e..5bdb4051 100644
--- a/pkg/atryum/atryum.go
+++ b/pkg/atryum/atryum.go
@@ -268,6 +268,18 @@ func runServer(args []string, o options) error {
service.SetInvocationSummarizer(&summaryAdapter{client: backendClient})
}
service.SetSessionStore(store.NewExternalSessionRepoWithDialect(db, dialect))
+ service.SetStreamOptions(
+ mcp.StreamOptions{
+ HeaderTimeout: time.Duration(cfg.Defaults.EffectiveStreamHeaderTimeoutSeconds()) * time.Second,
+ IdleTimeout: time.Duration(cfg.Defaults.StreamIdleTimeoutSeconds) * time.Second,
+ MaxDuration: time.Duration(cfg.Defaults.StreamMaxDurationSeconds) * time.Second,
+ MaxMessageBytes: cfg.Defaults.StreamMaxMessageBytes,
+ },
+ invocation.StreamAuditLimits{
+ MaxEvents: cfg.Defaults.StreamAuditMaxEvents,
+ MaxEventBytes: cfg.Defaults.StreamAuditMaxEventBytes,
+ },
+ )
serverAdmin := api.NewServerAdminService(serverRepo, oauthRepo, client, 5*time.Second, cfg.Server.PublicBaseURL)
if *initServers {
if err := initEnabledServerStatuses(context.Background(), serverRepo, serverAdmin); err != nil {
@@ -282,6 +294,7 @@ func runServer(args []string, o options) error {
}
handler := api.NewHandler(service, serverAdmin, policyRegistry, rulesRepo, agentsRepo, agentSyncSettingsRepo, llmConfigsRepo, syncAgentsFn, backendClient, localEvaluator)
handler.SetManagedAgentBindings(managedAgentBindingRepo)
+ handler.SetStreamRelayEnabled(cfg.Defaults.StreamRelayEnabled)
for _, register := range o.extraRoutes {
handler.AddExtraRoutes(register)
}
diff --git a/scripts/fake_agent.py b/scripts/fake_agent.py
index 1f01eff0..044face9 100644
--- a/scripts/fake_agent.py
+++ b/scripts/fake_agent.py
@@ -37,6 +37,14 @@
python fake_agent.py mcp --list-tools
python fake_agent.py mcp --tool add --arguments '{"a":2,"b":3}'
+ # Prove Atryum relays live progress instead of buffering the whole
+ # response: sends Accept: text/event-stream and fails unless progress
+ # notifications actually arrive measurably before the terminal result.
+ python fake_agent.py mcp everything-streaming \\
+ --tool trigger-long-running-operation \\
+ --arguments '{"duration":3,"steps":3}' \\
+ --stream --min-progress-events 3
+
# Pretend to be a specific harness:
python fake_agent.py mcp --client-name cursor --client-version 0.45.7 --list-tools
@@ -602,6 +610,106 @@ def mcp_call(
return payload
+def mcp_call_stream(
+ base: str,
+ server: str | None,
+ tool: str,
+ arguments: dict[str, Any],
+ req_id: int,
+ protocol_version: str = "2025-06-18",
+ extra_headers: dict[str, str] | None = None,
+ timeout: float = 60.0,
+ progress_token: str = "fake-agent-stream",
+) -> dict[str, Any]:
+ """Like mcp_call, but for a tools/call that requests an SSE response and
+ reads it incrementally instead of buffering with resp.read() — this is
+ what lets the caller record each frame's real arrival time. That's the
+ only way to prove Atryum relayed progress live rather than replaying
+ everything at once when the tool finally finished: the same technique
+ internal/api/mcp_everything_test.go uses in-process, exercised here over
+ a real HTTP connection to a real running atryum binary.
+
+ Deliberately simpler than _request_once: no 401-retry-with-forced-refresh,
+ since that would mean replaying a request whose response we've already
+ started reading. Proactive token refresh (_authorization_headers) still
+ applies, so this works for the same auth protocols as long as the token
+ isn't already expired when the call starts.
+ """
+ path = "/mcp/" + (server or "")
+ body = {
+ "jsonrpc": "2.0",
+ "id": req_id,
+ "method": "tools/call",
+ "params": {
+ "name": tool,
+ "arguments": arguments,
+ # Progress notifications are tied to a caller-supplied token —
+ # an upstream has nothing to report progress against without
+ # one, and would just answer with a plain buffered result.
+ "_meta": {"progressToken": progress_token},
+ },
+ }
+ headers = {
+ "MCP-Protocol-Version": protocol_version,
+ "Accept": "application/json, text/event-stream",
+ }
+ if extra_headers:
+ headers.update(extra_headers)
+ if "Authorization" not in headers:
+ headers.update(_authorization_headers())
+ headers.setdefault("Content-Type", "application/json")
+
+ req = urlrequest.Request(
+ base + path,
+ data=json.dumps(body).encode("utf-8"),
+ method="POST",
+ headers=headers,
+ )
+ events: list[dict[str, Any]] = []
+ started = time.monotonic()
+ with urlrequest.urlopen(req, timeout=timeout) as resp:
+ content_type = resp.headers.get("Content-Type", "")
+ if "text/event-stream" not in content_type:
+ raw = resp.read().decode("utf-8")
+ return {
+ "streamed": False,
+ "content_type": content_type,
+ "events": [],
+ "result": json.loads(raw) if raw else {},
+ }
+
+ # SSE framing: one or more "data:" lines per event, terminated by a
+ # blank line. Reading resp line-by-line (rather than resp.read()) is
+ # what makes each event's arrival timestamp real instead of "whenever
+ # the whole response finally finished".
+ data_lines: list[str] = []
+
+ def flush() -> None:
+ if not data_lines:
+ return
+ payload = "\n".join(data_lines)
+ try:
+ parsed = json.loads(payload)
+ except json.JSONDecodeError:
+ parsed = None
+ events.append({"t": time.monotonic() - started, "data": payload, "parsed": parsed})
+ data_lines.clear()
+
+ for raw_line in resp:
+ line = raw_line.decode("utf-8", errors="replace").rstrip("\r\n")
+ if line == "":
+ flush()
+ continue
+ if line.startswith("data:"):
+ data_lines.append(line[len("data:"):].lstrip(" "))
+ # Other SSE fields (event:, id:, retry:, ":" comments) don't
+ # matter here — Atryum's relay only ever sends data lines (see
+ # internal/api/sse_relay.go's writeSSEEvent).
+ flush()
+
+ return {"streamed": True, "content_type": content_type, "events": events}
+
+
def run_mcp(
base: str,
server: str | None,
@@ -611,6 +719,9 @@ def run_mcp(
bearer: str | None,
client_name: str,
client_version: str,
+ stream: bool = False,
+ min_progress_events: int = 1,
+ min_live_gap_seconds: float = 1.0,
) -> None:
headers: dict[str, str] = {}
if bearer:
@@ -651,15 +762,50 @@ def run_mcp(
return
# 4. tools/call
- call = mcp_call(
- base,
- server,
- "tools/call",
- {"name": tool, "arguments": arguments or {}},
- req_id=4,
- extra_headers=headers,
+ if not stream:
+ call = mcp_call(
+ base,
+ server,
+ "tools/call",
+ {"name": tool, "arguments": arguments or {}},
+ req_id=4,
+ extra_headers=headers,
+ )
+ _print_json("tools/call", call)
+ return
+
+ outcome = mcp_call_stream(base, server, tool, arguments or {}, req_id=4, extra_headers=headers)
+ if not outcome["streamed"]:
+ raise SystemExit(
+ "expected a streamed (text/event-stream) tools/call response, got "
+ f"content-type {outcome['content_type']!r} instead — the upstream "
+ "may not have opened a stream, or stream_relay_enabled may be off"
+ )
+
+ events = outcome["events"]
+ progress = [e for e in events if (e["parsed"] or {}).get("method") == "notifications/progress"]
+ terminal = events[-1] if events else None
+ if terminal is None or not isinstance(terminal["parsed"], dict) or (
+ "result" not in terminal["parsed"] and "error" not in terminal["parsed"]
+ ):
+ raise SystemExit(f"stream ended without a recognizable terminal JSON-RPC response: {events!r}")
+ if len(progress) < min_progress_events:
+ raise SystemExit(
+ f"expected at least {min_progress_events} live progress notification(s), "
+ f"got {len(progress)} — the relay may be buffering instead of streaming live"
+ )
+ gap = terminal["t"] - progress[0]["t"]
+ if gap < min_live_gap_seconds:
+ raise SystemExit(
+ f"progress and the terminal result arrived {gap:.3f}s apart — too "
+ f"close together to prove live delivery rather than a buffered "
+ f"replay (want >= {min_live_gap_seconds}s)"
+ )
+ print(
+ f"tools/call streamed: {len(progress)} progress notification(s) over "
+ f"{gap:.2f}s before the terminal result"
)
- _print_json("tools/call", call)
+ _print_json("tools/call (terminal)", terminal["parsed"])
# ─── argparse ───────────────────────────────────────────────────────────────
@@ -746,6 +892,31 @@ def main(argv: list[str] | None = None) -> int:
default=default_version,
help=f"clientInfo.version sent in initialize (default this run: {default_version})",
)
+ pm.add_argument(
+ "--stream",
+ action="store_true",
+ help=(
+ "send Accept: text/event-stream on tools/call and verify progress "
+ "actually arrives live before the terminal result, instead of just "
+ "checking the final answer"
+ ),
+ )
+ pm.add_argument(
+ "--min-progress-events",
+ type=int,
+ default=1,
+ help="with --stream, minimum live progress notifications required (default 1)",
+ )
+ pm.add_argument(
+ "--min-live-gap-seconds",
+ type=float,
+ default=1.0,
+ help=(
+ "with --stream, minimum seconds required between the first progress "
+ "notification and the terminal result, proving live delivery rather "
+ "than a buffered replay (default 1.0)"
+ ),
+ )
args = p.parse_args(argv)
_set_token_cache_key(args.base)
@@ -780,6 +951,9 @@ def main(argv: list[str] | None = None) -> int:
bearer=args.bearer,
client_name=args.client_name,
client_version=args.client_version,
+ stream=args.stream,
+ min_progress_events=args.min_progress_events,
+ min_live_gap_seconds=args.min_live_gap_seconds,
)
else:
p.error(f"unknown mode {args.mode!r}")