From 94a6f960e83f3ac9ccce63eaaa2b633b7eea6ab6 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Wed, 22 Jul 2026 17:36:01 -0400 Subject: [PATCH 01/18] Support streaming with mcp --- CHANGELOG.md | 20 + atryum.example.toml | 29 + cmd/atryum/main.go | 16 + docs/architecture.md | 169 ++++ internal/api/handlers.go | 289 +++++- internal/api/handlers_test.go | 552 +++++++++++ internal/config/config.go | 38 +- internal/config/config_test.go | 45 + internal/invocation/model.go | 5 + internal/invocation/service.go | 160 +++- internal/invocation/service_test.go | 567 ++++++++++- internal/invocation/stream_sink.go | 188 ++++ internal/mcp/client.go | 1162 +++++++++++++++++++++-- internal/mcp/client_test.go | 1042 +++++++++++++++++++- internal/mcp/stdio_process_unix.go | 58 ++ internal/mcp/stdio_process_unix_test.go | 51 + internal/mcp/stdio_process_windows.go | 12 + 17 files changed, 4286 insertions(+), 117 deletions(-) create mode 100644 internal/invocation/stream_sink.go create mode 100644 internal/mcp/stdio_process_unix.go create mode 100644 internal/mcp/stdio_process_unix_test.go create mode 100644 internal/mcp/stdio_process_windows.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 33b72912..a9b8467a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,26 @@ 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 are audited into `invocation_events` + (`invocation.stream_event` / `invocation.stream_completed`). New + `[defaults]` config knobs: `stream_relay_enabled` (kill-switch, default + on), `stream_header_timeout_seconds`, `stream_idle_timeout_seconds`, + `stream_max_duration_seconds`, `stream_audit_max_events`, + `stream_audit_max_event_bytes`. See `docs/architecture.md` for the full + design. + ## [0.2.0] - 2026-07-14 ### Added diff --git a/atryum.example.toml b/atryum.example.toml index d1e87dbd..ea1a0181 100644 --- a/atryum.example.toml +++ b/atryum.example.toml @@ -49,6 +49,35 @@ connection_timeout_seconds = 5 [defaults] request_timeout_seconds = 30 +# Live SSE relay for tools/call: when an upstream MCP server answers a +# tools/call with a Server-Sent Events stream (per the Streamable HTTP +# transport, MCP spec 2025-11-25), Atryum relays intermediate messages +# (progress, logging, other notifications) to the agent as they arrive, +# instead of buffering the whole response and returning only the terminal +# result. This only ever activates when the agent's own request also sends +# Accept: text/event-stream and the upstream responds with one; non-streaming +# tool calls are completely unaffected. stream_relay_enabled is a kill-switch +# to disable the relay globally without a rollback. +stream_relay_enabled = true +# Bounds waiting for the upstream's response headers on a streaming +# tools/call (the connect phase, before Atryum knows whether the response +# will stream). 0 falls back to request_timeout_seconds above. +stream_header_timeout_seconds = 0 +# Bounds the gap between successive relayed events once a stream has +# started; resets on every event. Does not bound the call's total duration. +stream_idle_timeout_seconds = 60 +# Bounds the whole call once a stream has started. 0 = unlimited. +stream_max_duration_seconds = 600 +# 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 +# audit queue can still drop events if storage cannot keep up, and reports +# those drops in invocation.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/cmd/atryum/main.go b/cmd/atryum/main.go index fa39b42e..9c5b4d65 100644 --- a/cmd/atryum/main.go +++ b/cmd/atryum/main.go @@ -244,6 +244,21 @@ func runServer(args []string) error { service.SetInvocationSummarizer(&summaryAdapter{client: backendClient}) } service.SetSessionStore(store.NewExternalSessionRepoWithDialect(db, dialect)) + streamHeaderTimeoutSeconds := cfg.Defaults.StreamHeaderTimeoutSeconds + if streamHeaderTimeoutSeconds <= 0 { + streamHeaderTimeoutSeconds = cfg.Defaults.RequestTimeoutSeconds + } + service.SetStreamOptions( + mcp.StreamOptions{ + HeaderTimeout: time.Duration(streamHeaderTimeoutSeconds) * time.Second, + IdleTimeout: time.Duration(cfg.Defaults.StreamIdleTimeoutSeconds) * time.Second, + MaxDuration: time.Duration(cfg.Defaults.StreamMaxDurationSeconds) * time.Second, + }, + 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 { @@ -258,6 +273,7 @@ func runServer(args []string) error { } handler := api.NewHandler(service, serverAdmin, policyRegistry, rulesRepo, agentsRepo, agentSyncSettingsRepo, llmConfigsRepo, syncAgentsFn, backendClient, localEvaluator) handler.SetManagedAgentBindings(managedAgentBindingRepo) + handler.SetStreamRelayEnabled(cfg.Defaults.StreamRelayEnabled) authValidator, err := auth.NewValidator(cfg.Auth, nil) if err != nil { diff --git a/docs/architecture.md b/docs/architecture.md index f3ced24d..aae84bdf 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -186,6 +186,175 @@ 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 + +This section is a special case of the execution flow above, scoped to just one of the +two `Invoke`-backed entry points from the Runtime ingress table: the MCP proxy's +`tools/call`. Direct invocation (`POST /api/v1/invocations`) is a plain REST endpoint — +it does not negotiate `Accept: text/event-stream` and never streams. + +**Background, read this first.** Every message in this protocol is JSON-RPC: a +*request* asks for something and carries an `id`; a *response* answers one specific +request by carrying that same `id` back; a *notification* is a one-way heads-up with no +`id` and no reply expected. Server-Sent Events (SSE) is just a way to send several of +these messages down one HTTP connection over time — as `data:` lines, one message each, +separated by blank lines — instead of sending one message all at once. A tool call +answered this way might look like: + +``` +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 line is a notification — a progress update nobody has to reply to. The second +is the real answer (it has an `id` and a `result`). This doc calls that second one the +**terminal response**: it's the one and only thing a non-streaming call would have +returned. + +**What Atryum does with this.** If the upstream tool streams progress notifications +while working on a call, Atryum forwards each one to the agent as it happens — the agent +sees them live instead of waiting in silence until the call finishes. + +If the upstream tool answers with a single plain response and no streaming, the agent +simply receives that response. There are no live updates to relay in that case. + +**The three layers involved**, in order: the part of Atryum facing the agent decides +whether to relay live and writes the response back; the part in the middle runs +approval rules and keeps the audit trail; the part facing the upstream tool speaks the +actual wire protocol and hands messages up as they arrive. (If you want to find the +code: `internal/api` → `internal/invocation` → `internal/mcp`, in that order.) + +```mermaid +sequenceDiagram + autonumber + participant Agent + participant Facing as Atryum: agent-facing layer + participant Middle as Atryum: rules & audit layer + participant Client as Atryum: upstream-facing layer + participant Upstream as Upstream tool server + + Agent->>Facing: Call a tool (willing to receive a live stream) + Facing->>Middle: Run the call + Middle->>Client: Send the call upstream + Client->>Upstream: Call the tool + alt Upstream streams progress before answering + Upstream-->>Client: Starts streaming + Client->>Middle: A progress update arrived + Middle->>Middle: Record it for the audit trail (in the background) + Middle->>Facing: Forward the update + Facing-->>Agent: Relay it live + Note over Client,Upstream: repeats for every update sent + Upstream-->>Client: Final answer + Client-->>Middle: Done + Middle->>Middle: Save the result, close out the audit trail + Middle-->>Facing: Done + Facing-->>Agent: Send the final answer + else Upstream just answers directly, no streaming + Upstream-->>Client: Single reply + Note over Client,Facing: nothing has been sent to the agent yet + Client-->>Middle: Done + Middle-->>Facing: Done + Facing-->>Agent: Send the one reply + end +``` + +**Approval gating applies before any streaming can start.** If a rule says a tool call +needs a human to approve it first, Atryum pauses *before ever contacting the upstream +tool* — which is before any streaming could even start. Nothing is sent to the agent +while a call is waiting on approval, whether or not the agent asked for a live stream. + +```mermaid +sequenceDiagram + participant Agent + participant Atryum + participant Reviewer as Human reviewer + participant Upstream as Upstream tool server + + Agent->>Atryum: Call a tool + Atryum->>Atryum: Rule says this needs approval — pause and wait + Note over Agent,Atryum: Connection stays open. Nothing sent yet. + Reviewer->>Atryum: Approve + Atryum->>Upstream: Now make the call + Note over Agent,Upstream: From here it's the same as the diagram above + Upstream-->>Agent: Progress updates, then the final answer +``` + +**Every relayed update is also recorded for the audit trail**, in the background, so a +slow database write can never delay a live update reaching the agent or eat into the +time budget described below. + +**Why there are three separate time limits, not one.** A normal (non-streaming) call has +one clock: if it takes too long overall, Atryum gives up. A streaming call can't work +that way, because "this is taking a while" is expected and fine as long as *something* +keeps happening. So there are three limits instead of one: + +- How long to wait for the upstream tool to even start responding. +- How long to wait between one update and the next, once it has started — this one + resets every time something new arrives, so a stream that's still actively sending + updates is never punished just for running long overall. +- A hard ceiling on the whole call, just in case, so nothing runs forever. + +Whichever limit is hit first ends the call, and Atryum records *which one* (a stalled +tool, an agent that disconnected, and an ordinary network failure all get a different, +searchable label in the audit trail) — so whoever's debugging later doesn't have to +guess which of the three actually happened. + +**The agent's own connection has a time limit too.** If the agent stops reading — its +connection dies, or it just stops paying attention — Atryum needs to notice, otherwise +it would sit forever trying to hand off data nobody is receiving, which would also +stall the upstream tool call waiting behind it. So every write to the agent has its own +short time limit, and while the upstream tool is quiet, Atryum sends small "still here" +pings on the open connection, so ordinary network equipment sitting in between doesn't +mistake a slow-but-healthy call for a dead one and close it. + +**If the upstream tool's own connection drops mid-stream, Atryum reconnects and picks up +where it left off** — this is a feature of the underlying protocol, not something Atryum +invented — so the agent doesn't see a glitch. The connection *from* Atryum *to* the +agent intentionally does not support this kind of resume: promising it would mean +promising to remember exactly where every agent left off, even across an Atryum +restart, which isn't a promise Atryum makes today. + +Atryum does not forward messages that the *upstream tool* might try to send back toward +the agent as if it were the agent's own message (a few advanced, rarely-used parts of +the protocol allow this) — those get recorded for the audit trail and dropped, since +Atryum has no way to get a reply back to where it came from. The always-on keepalive +connection agents can poll separately is a distinct mechanism from this relay. + +A single on/off switch lets an operator disable the relay entirely, so every `tools/call` +gets a single buffered response regardless of what the agent or upstream would otherwise +support. + +#### Edge cases and failure paths + +- **A tool call that failed to connect properly must not get relayed twice.** If Atryum + needs to retry the setup for a call, it only does so before anything has been sent to + the agent — retrying after the agent has already seen live updates would relay + everything a second time, so Atryum refuses to retry once that's happened. +- **Something going wrong mid-stream must still leave the agent with a real answer.** + Even if a problem is discovered after Atryum has already started relaying updates, the + agent still gets a proper closing message explaining what happened — never a + connection that just quietly closes with no explanation. +- **A duplicate update on reconnect must not reach the agent twice.** When Atryum + reconnects to an upstream tool mid-stream, some tools resend the very last update as + part of "catching up." Atryum recognizes and skips that one repeated update, so the + agent only ever sees it once. +- **An update that spans multiple lines must not get corrupted.** The wire format allows + a single update to be written across more than one line. Atryum preserves that shape + correctly when relaying it, instead of accidentally squashing it into something that + looks like a different, malformed message. +- **A stuck external tool process must be fully stopped, not half-stopped, when it times + out.** Some upstream tools run as an external program that can itself launch further + child programs. If Atryum only stopped the program it directly started, a leftover + child process could keep running and keep things open, and the timeout would never + actually free anything up. Atryum stops the whole group of processes together. +- **Resetting a "how long has it been quiet" timer must not accidentally end a call + that's actually fine.** Naively restarting a countdown every time an update arrives + can, in rare bad timing, let the old countdown finish a split second before the + restart takes effect — ending a call that was actually still healthy. Atryum + double-checks how much time has *really* passed before deciding to end a call, so a + well-timed update can never be wrongly punished by bad luck in the timing. + ## Decision-only calls `Submit` persists and returns a decision without contacting an MCP server. An external diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 3c0122e6..a99f12ad 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -1,6 +1,7 @@ package api import ( + "bytes" "context" "crypto/rand" "database/sql" @@ -42,6 +43,7 @@ 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) 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) @@ -151,6 +153,13 @@ type Handler struct { authValidator *auth.Validator apiKeyAuth auth.APIKeyConfig + // streamRelayEnabled is the kill-switch for the tools/call SSE relay + // (see handleMCPProxy). The relay only ever activates when the agent's + // POST also sends Accept: text/event-stream and the upstream answers + // with an SSE body, so leaving this on by default is safe; it exists so + // the feature 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 @@ -707,7 +716,15 @@ 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. +func (h *Handler) SetStreamRelayEnabled(enabled bool) { + h.streamRelayEnabled = enabled } // SetAuthValidator installs the inbound auth validator. When non-nil, the @@ -1003,6 +1020,198 @@ func writeSSEComment(w io.Writer, comment string) { fmt.Fprint(w, "\n") } +// writeSSEEvent writes and flushes one SSE frame. Per the SSE spec, a +// multi-line payload must be sent as one "data:" line per line of content — +// a raw newline embedded in a single "data:" line breaks framing, since any +// continuation line lacking its own field prefix is ignored by a compliant +// parser. This matters here: a relayed notification's data (evt.Data, +// reconstructed by mcp.sseEventReader) can genuinely be multi-line if the +// upstream sent it that way (see mcp.TestListToolsDecodesMultilineSSEData +// for a real example) — only the terminal frame (always compact +// json.Marshal output) is guaranteed single-line. No "id:" field is ever +// emitted: doing so implies Last-Event-ID resumability, which Atryum does +// not support and must not advertise. Returns whatever error the write +// itself produced, which is how a broken downstream connection (the agent +// disconnected) is detected. +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 bytes.Split(data, []byte("\n")) { + 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 implements mcp.StreamSink for one agent-facing tools/call +// request. It relays every intermediate upstream event to the agent as an +// SSE frame, live, as InvokeStreaming reads it from the upstream, and keeps +// the downstream connection alive with heartbeat comments while the +// upstream is silent. +// +// The downstream response only switches to SSE mode when StreamStarted +// fires — until then, handleMCPProxy has written nothing, so a JSON +// (non-streaming) upstream response still produces today's exact buffered +// reply. Once started is true, headers have already been sent: the caller +// must never call writeRPCResult/writeRPCError/WriteHeader again for this +// request; the terminal response is written via finishStream instead. +// +// Concurrency: StreamStarted/Event run synchronously on the handler +// goroutine (inside InvokeStreaming's call stack); the heartbeat runs on +// its own goroutine. mu serializes every write to w so a heartbeat can +// never interleave with (and corrupt) an event or terminal frame. +// finishStream stops the heartbeat before writing the terminal frame, so +// no write can occur after the handler returns. +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.eventCount++ + s.setWriteDeadlineLocked() + if err := writeSSEEvent(s.w, s.flusher, "", evt.Data); err != nil { + s.writeErr = err + return err + } + return nil +} + +// 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.stopOnce.Do(func() { close(s.heartbeatStop) }) + <-s.heartbeatDone + 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 +} + func isJSONRPCRequest(r *http.Request) bool { if strings.Contains(strings.ToLower(r.Header.Get("Content-Type")), "application/json") { return true @@ -1286,6 +1495,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") @@ -1295,7 +1505,7 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server h.handleAtryumRulesToolCall(w, r, req.ID, server, 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) } @@ -1306,22 +1516,63 @@ 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 response mode switch + // happens lazily, inside svc.InvokeStreaming, only if the upstream + // actually answers with an SSE stream (sink.StreamStarted). Until + // then nothing has been written, so a JSON upstream response still + // produces exactly today's buffered reply below. + var sink *sseRelaySink + if flusher, ok := w.(http.Flusher); ok && h.streamRelayEnabled && acceptsEventStream(r) { + sink = newSSERelaySink(w, flusher) + } + 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 { + h.debugf("mcp tools.call error after stream started server=%s tool=%s err=%v", server, params.Name, err) + errBody, _ := json.Marshal(map[string]any{"code": -32000, "message": err.Error()}) + terminal, _ := json.Marshal(jsonRPCResponse{JSONRPC: "2.0", ID: req.ID, Error: errBody}) + _ = sink.finishStream(terminal) + 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}) + _ = sink.finishStream(terminal) 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 { @@ -1639,6 +1890,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 { diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index c4e1adf0..441ec028 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -1,6 +1,7 @@ package api import ( + "bufio" "bytes" "context" "database/sql" @@ -11,6 +12,7 @@ import ( "net/http/httptest" "net/url" "strings" + "sync/atomic" "testing" "time" @@ -18,6 +20,7 @@ import ( backendclient "atryum/internal/backend" "atryum/internal/config" "atryum/internal/invocation" + "atryum/internal/invocation/policy" "atryum/internal/managedagents" "atryum/internal/mcp" "atryum/internal/store" @@ -49,6 +52,11 @@ type stubService struct { createSessionReq *invocation.CreateSessionRequest createSessionAgentID 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) } func (s *stubService) Invoke(ctx context.Context, req invocation.CreateInvocationRequest) (invocation.InvocationResponse, error) { @@ -56,6 +64,20 @@ 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) ListTools(context.Context, string) ([]mcp.Tool, error) { return s.tools, s.listErr } @@ -1210,6 +1232,536 @@ 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 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. +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 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, "persisting result failed") { + t.Fatalf("expected a terminal JSON-RPC error frame, got %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. +func TestMCPToolsCallEndToEndRelaysLiveBeforeTerminalResponseExists(t *testing.T) { + releaseTerminal := 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 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() + + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("open db: %v", err) + } + defer 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: "demo", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}, + }) + 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{}, 5*time.Second, nil, nil, nil, nil, + ) + + h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil) + agentServer := httptest.NewServer(h.Routes()) + defer agentServer.Close() + + 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) + } +} + 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/config/config.go b/internal/config/config.go index d7f21b55..31e689c1 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -115,6 +115,37 @@ 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 + // relay only ever activates when the agent's own request also sends + // Accept: text/event-stream and the upstream answers with an SSE + // body, so leaving this on by default is safe. + StreamRelayEnabled bool `toml:"stream_relay_enabled"` + // StreamHeaderTimeoutSeconds bounds waiting for the upstream's + // response headers on a streaming tools/call — the connect phase, + // before Atryum knows whether the response will be an SSE stream. + // Zero falls back to RequestTimeoutSeconds, preserving today's + // connect-phase behavior. + StreamHeaderTimeoutSeconds int `toml:"stream_header_timeout_seconds"` + // StreamIdleTimeoutSeconds bounds the gap between successive relayed + // events once a stream has started; it resets on every event. Unlike + // RequestTimeoutSeconds, this does not bound the call's total + // duration — only how long it may go without producing anything. + StreamIdleTimeoutSeconds int `toml:"stream_idle_timeout_seconds"` + // StreamMaxDurationSeconds bounds the whole call once a stream has + // started. Zero disables the bound (unlimited). + StreamMaxDurationSeconds int `toml:"stream_max_duration_seconds"` + // 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"` } type UpstreamConfig struct { @@ -150,7 +181,12 @@ func Load(path string) (Config, error) { ConnectionTimeoutSecs: 5, }, Defaults: DefaultsConfig{ - RequestTimeoutSeconds: 30, + RequestTimeoutSeconds: 30, + StreamRelayEnabled: true, + StreamIdleTimeoutSeconds: 60, + StreamMaxDurationSeconds: 600, + StreamAuditMaxEvents: 100, + StreamAuditMaxEventBytes: 4096, }, } _, err := toml.DecodeFile(path, &cfg) diff --git a/internal/config/config_test.go b/internal/config/config_test.go index f4b53379..7becaa1b 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -72,6 +72,51 @@ 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.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) { diff --git a/internal/invocation/model.go b/internal/invocation/model.go index 2b1dc948..aa5e4444 100644 --- a/internal/invocation/model.go +++ b/internal/invocation/model.go @@ -100,6 +100,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 6dd3a4e1..31f09237 100644 --- a/internal/invocation/service.go +++ b/internal/invocation/service.go @@ -171,13 +171,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 audit sink's +// bounded queue remains a final backpressure guard even with MaxEvents zero, +// and reports any queue 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 @@ -199,6 +217,12 @@ type Service struct { mu sync.Mutex pendingApprovals map[string]chan approvalDecision + // streamOptions and streamAuditLimits govern InvokeStreaming's execution + // once a sink is present (see finishExecution). The zero value of each + // disables its bounds; SetStreamOptions installs real values. + streamOptions mcp.StreamOptions + streamAuditLimits StreamAuditLimits + toolCatalogMu sync.Mutex toolCatalog map[string]toolCatalogEntry } @@ -237,6 +261,14 @@ 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 finishExecution). Calls with a nil sink are unaffected. +func (s *Service) SetStreamOptions(opts mcp.StreamOptions, auditLimits StreamAuditLimits) { + s.streamOptions = opts + s.streamAuditLimits = auditLimits +} + // 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. @@ -291,7 +323,21 @@ 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 when +// sink is non-nil and the upstream answers the tools/call with an SSE +// stream, intermediate JSON-RPC messages (progress, logging, other +// notifications) are relayed to sink as they arrive — see finishExecution. +// 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") } @@ -430,12 +476,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) } } @@ -977,7 +1023,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 { @@ -993,11 +1039,11 @@ 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) { +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 @@ -1100,14 +1146,26 @@ 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 if err != nil { @@ -1140,6 +1198,88 @@ func (s *Service) finishExecution(ctx context.Context, inv Invocation, upstream return s.toResponse(inv), nil } +// finishExecutionStreaming is finishExecution's live-relay path: it wraps +// the caller's sink in an auditing decorator (so every relayed event and +// the call's outcome are recorded as invocation_events rows regardless of +// whether the downstream write later fails) and calls InvokeStream with +// s.streamOptions instead of the fixed s.defaultTimeout. +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.streamAuditLimits) + result, err := s.client.InvokeStream(ctx, upstream, req.Tool, req.Input, req.RequestID, req.Meta, audited, s.streamOptions) + completed := time.Now().UTC() + inv.CompletedAt = &completed + + if err != nil { + inv.Status = StatusFailed + reason, message := classifyStreamError(audited, err) + inv.Error = mustJSON(map[string]any{"message": message}) + audited.finish(completed, "failed") + persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout) + defer cancelPersist() + if updateErr := s.invocations.UpdateResult(persistCtx, inv); updateErr != nil { + return InvocationResponse{}, fmt.Errorf("persist streaming invocation failure: %w", updateErr) + } + // 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_relayed": audited.seq}), + CreatedAt: completed, + }) + return s.toResponse(inv), nil + } + var terminalEvent Event + if result.Failed { + inv.Status = StatusFailed + inv.Error = result.Body + audited.finish(completed, "failed") + 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 + audited.finish(completed, "succeeded") + 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 { + return InvocationResponse{}, err + } + _ = 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(audited *auditingSink, err error) (reason string, message string) { + if audited.downstreamErr != nil { + return "stream_aborted_downstream", audited.downstreamErr.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() + } + return "transport_error", err.Error() +} + func (s *Service) Approve(ctx context.Context, invocationID string, actorID string) error { s.mu.Lock() ch, ok := s.pendingApprovals[invocationID] diff --git a/internal/invocation/service_test.go b/internal/invocation/service_test.go index a00bfefa..67e905cb 100644 --- a/internal/invocation/service_test.go +++ b/internal/invocation/service_test.go @@ -5,6 +5,7 @@ import ( "database/sql" "encoding/json" "errors" + "fmt" "net/http" "net/http/httptest" "os" @@ -70,12 +71,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 @@ -234,6 +244,559 @@ func TestSubmitLogsAndAuditsRuleLoadFailure(t *testing.T) { } } +// 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 +} + +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.Fatalf("decode request: %v", err) + } + 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.Fatalf("unexpected method %q", body["method"]) + } + })) +} + +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) + } +} + +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 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) { + var toolsCallCount int + upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) { + toolsCallCount++ + 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 toolsCallCount != 1 { + t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed mid-stream)", toolsCallCount) + } + + 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") + } +} + +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() { + time.Sleep(50 * time.Millisecond) + if sink.touched() { + t.Errorf("sink touched before approval — approval gating must precede any relay") + } + list, err := service.List(context.Background(), invocation.InvocationListFilter{Limit: 10}) + if err != nil || len(list.Items) == 0 { + t.Errorf("expected a pending invocation to approve") + return + } + if err := service.Approve(context.Background(), list.Items[0].InvocationID, ""); 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) + } +} + func TestResolverBootstrapsServersFromConfigWhenDBEmpty(t *testing.T) { db := newSQLiteTestDB(t) repo := store.NewServerRepo(db) diff --git a/internal/invocation/stream_sink.go b/internal/invocation/stream_sink.go new file mode 100644 index 00000000..6a61ab7e --- /dev/null +++ b/internal/invocation/stream_sink.go @@ -0,0 +1,188 @@ +package invocation + +import ( + "context" + "sync" + "sync/atomic" + "time" + + "atryum/internal/mcp" +) + +const ( + streamAuditQueueCapacity = 128 + streamAuditWriteTimeout = 500 * time.Millisecond + streamAuditFlushTimeout = 2 * time.Second +) + +// 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 a bounded background queue rather than on the +// relay's hot path. A stalled audit repository must neither delay an event +// reaching the agent nor defeat the stream's idle/max-duration bounds. Each +// write and the final queue drain are time-bounded; failures and queue drops +// are summarized in invocation.stream_completed. +type auditingSink struct { + inner mcp.StreamSink // may be nil: audit-only, no relay + events eventRepo + invocationID string + requestID *string + upstreamName string + limits StreamAuditLimits + auditCtx context.Context + cancelAudit context.CancelFunc + auditQueue chan Event + auditDone chan struct{} + closeOnce sync.Once + persisted atomic.Int64 + failed atomic.Int64 + dropped atomic.Int64 + + 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, + } + if events != nil { + a.auditCtx, a.cancelAudit = context.WithCancel(context.Background()) + a.auditQueue = make(chan Event, streamAuditQueueCapacity) + a.auditDone = make(chan struct{}) + go a.runAuditWriter() + } + 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) recordEvent(evt mcp.StreamEvent) { + if a.auditQueue == 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(), + } + select { + case a.auditQueue <- record: + default: + a.dropped.Add(1) + } +} + +func (a *auditingSink) runAuditWriter() { + defer close(a.auditDone) + for evt := range a.auditQueue { + writeCtx, cancel := context.WithTimeout(a.auditCtx, streamAuditWriteTimeout) + err := a.events.Create(writeCtx, evt) + cancel() + if err != nil { + a.failed.Add(1) + continue + } + a.persisted.Add(1) + } +} + +func (a *auditingSink) stopAuditWriter() { + if a.auditQueue == nil { + return + } + a.closeOnce.Do(func() { close(a.auditQueue) }) + timer := time.NewTimer(streamAuditFlushTimeout) + defer timer.Stop() + select { + case <-a.auditDone: + a.cancelAudit() + case <-timer.C: + // Cancel the in-flight write and make every queued write fail fast. + // Do not wait indefinitely if an eventRepo violates context + // cancellation; terminal invocation persistence must remain bounded. + a.cancelAudit() + select { + case <-a.auditDone: + case <-time.After(streamAuditWriteTimeout): + } + } +} + +// finish records the invocation.stream_completed totals row. terminal is +// "succeeded", "failed", or "aborted". +func (a *auditingSink) finish(completed time.Time, terminal string) { + if a.events == nil { + return + } + a.stopAuditWriter() + 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(), + "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/mcp/client.go b/internal/mcp/client.go index 46a7e3c5..fcb761b7 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" @@ -260,6 +261,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 @@ -408,21 +423,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 +627,51 @@ 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) { +// marshalToolCallEnvelope builds the JSON-RPC tools/call request body Atryum +// sends upstream. The envelope always uses id "1" — 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) { + params := map[string]any{"name": tool, "arguments": input} + if merged := mergeRequestMeta(meta, requestID); merged != nil { + params["_meta"] = merged + } + return json.Marshal(Envelope{JSONRPC: "2.0", ID: json.RawMessage([]byte("1")), 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, json.RawMessage([]byte("1"))) + 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 +679,609 @@ 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 + 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 + } + } + c.debugf("upstream http tools.call server=%s status=%d failed=%t", upstream.Name, result.StatusCode, invoke.Failed) + return invoke, nil +} + +// StreamEvent is one upstream SSE event carrying a JSON-RPC message that is +// not the terminal response to the call: either a notification (progress, +// logging, or any other server-to-client notification) or, more rarely, a +// server-to-client request. +type StreamEvent struct { + // Data is the joined "data:" payload for this event: one JSON-RPC message. + 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 +} + +// StreamSink receives upstream SSE events 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 exactly once, before the first event is + // delivered or the terminal response is returned — never for an + // attempt that gets silently retried (see relaySSEToolCall). + 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 how long InvokeStream may take once a stream has +// started. A zero-valued field disables that particular bound. +type StreamOptions struct { + // HeaderTimeout bounds waiting for the upstream's initial response + // headers, i.e. before we know whether the response is streaming. Zero + // leaves this phase bounded only by ctx's own deadline, if any. + HeaderTimeout time.Duration + // IdleTimeout bounds the gap between successive events once the stream + // has started; it resets after every event. Zero disables the check. + IdleTimeout time.Duration + // MaxDuration bounds the whole call once the stream has started. Zero + // disables the check. + MaxDuration time.Duration +} + +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 +} + +// callTimeoutGuard implements InvokeStream's header/idle/max-duration +// timeout scheme by canceling one shared context — the same context the +// HTTP request and its body reads run under. Header timing bounds only the +// wait for response headers; once headers arrive the caller disarms it and +// arms the idle/max-duration timers for the body-read phase instead, so a +// slow-to-start upstream and a slow-once-started upstream are judged against +// the right bound for each phase, rather than one fixed wall-clock budget +// covering both (which is what the plain per-call http.Client timeout does). +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 + headerTimer *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) armHeaderTimeout(d time.Duration) { + if d <= 0 { + return + } + g.mu.Lock() + defer g.mu.Unlock() + if g.stopped { + return + } + g.headerTimer = time.AfterFunc(d, func() { g.trip("timed out waiting for upstream response headers") }) +} + +func (g *callTimeoutGuard) disarmHeaderTimeout() { + g.mu.Lock() + defer g.mu.Unlock() + if g.headerTimer != nil { + g.headerTimer.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.headerTimer != nil { + g.headerTimer.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 +} + +// 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, opts StreamOptions) (streamCallOutcome, error) { + guard := newCallTimeoutGuard(ctx) + defer guard.stop() + guard.armHeaderTimeout(opts.HeaderTimeout) + + h, err := c.doHTTPEnvelopeHeaders(guard.ctx, upstream, body, DefaultMCPProtocolVersion, true) + guard.disarmHeaderTimeout() + if err != nil { + if reason := guard.reason(); reason != "" { + return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) + } + return streamCallOutcome{}, 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() + _, _ = io.Copy(io.Discard, resp.Body) + return streamCallOutcome{missingSession: true, sessionID: h.sessionID}, nil + } + + if !strings.Contains(strings.ToLower(h.contentType), "text/event-stream") { + defer resp.Body.Close() + bodyBytes, err := io.ReadAll(resp.Body) + if err != nil { + if reason := guard.reason(); reason != "" { + return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) } - 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 + return streamCallOutcome{}, 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, guard, upstream, h.sessionID) +} + +// relaySSEToolCall reads resp's SSE body incrementally via an +// sseEventReader, relaying every intermediate (non-terminal) message to +// sink as it arrives, and returns once the terminal JSON-RPC response for +// id "1" is read. resp.Body is not closed here — the caller does that. +// sessionID is the session this attempt was sent under; it is always +// stamped onto the returned outcome (even a missing-session terminal +// response) so a caller retry can identify and clear the right session — +// mirroring doHTTPEnvelope's ForwardResult.SessionID contract. +// +// sink.StreamStarted fires lazily, right before the first thing is actually +// delivered — not simply because the response's Content-Type was SSE. This +// matters for the missing-session retry: if the very first (and only) +// message is a missing-session terminal error, the whole attempt is +// discarded and silently retried (see invokeHTTPStream), so the sink must +// never have been told a stream started for it. Once a real event has been +// relayed, or the terminal response is anything other than a +// zero-events missing-session error, the attempt is the one that counts and +// StreamStarted fires exactly once for it. +func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, guard *callTimeoutGuard, upstream Upstream, sessionID string) (streamCallOutcome, error) { + expectedID := json.RawMessage([]byte("1")) + currentResp := resp + reader := newSSEEventReader(currentResp.Body) + relayed := 0 + started := false + lastEventID := "" + retryDelay := time.Duration(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 := "" + ensureStarted := func() { + if !started { + started = true + sink.StreamStarted() + } + } + for { + evt, err := reader.NextEvent() + if err != nil { + if reason := guard.reason(); reason != "" { + _ = currentResp.Body.Close() + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q %s after %d relayed event(s): %w", upstream.Name, reason, relayed, ErrStreamTimeout) } - rpcResp, err = decodeRPCResponse(result, json.RawMessage([]byte("1"))) + if err != io.EOF { + _ = currentResp.Body.Close() + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + _ = currentResp.Body.Close() + if lastEventID == "" { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q closed the stream without a JSON-RPC response or resumable event id", upstream.Name) + } + if err := waitForSSEReconnect(guard.ctx, retryDelay); err != nil { + if reason := guard.reason(); reason != "" { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q %s while waiting to resume after %d relayed event(s): %w", upstream.Name, reason, relayed, ErrStreamTimeout) + } + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + currentResp, err = c.resumeSSEStream(guard.ctx, upstream, lastEventID) 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 reason := guard.reason(); reason != "" { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q %s while resuming after %d relayed event(s): %w", upstream.Name, reason, relayed, ErrStreamTimeout) + } + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + reader = newSSEEventReader(currentResp.Body) + resumedFrom = lastEventID + continue + } + guard.resetIdle() + 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 + } + payload := evt.Data + + switch classifyRPCMessage(payload, expectedID) { + case rpcMessageTerminalResponse: + var rpcResp rpcResponse + if err := json.Unmarshal(payload, &rpcResp); err != nil { + _ = currentResp.Body.Close() + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, currentResp.StatusCode) + if !(missingSession && relayed == 0) { + ensureStarted() + } + _ = currentResp.Body.Close() + return streamCallOutcome{invoke: invoke, missingSession: missingSession, eventsRelayed: relayed, sessionID: sessionID}, nil + case rpcMessageServerRequest: + ensureStarted() + relayed++ + if err := sink.Event(StreamEvent{Data: payload, ServerRequest: true}); err != nil { + _ = currentResp.Body.Close() + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + case rpcMessageNotification: + ensureStarted() + relayed++ + if err := sink.Event(StreamEvent{Data: payload}); err != nil { + _ = currentResp.Body.Close() + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + default: + // Unrecognized payload shape (e.g. a response to some other id). + // Not ours to interpret; ignore and keep reading. } } - 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 +} + +func waitForSSEReconnect(ctx context.Context, delay time.Duration) error { + if delay <= 0 { + return nil } - bodyBytes := rpcResp.Result - if len(bodyBytes) == 0 || string(bodyBytes) == "null" { - bodyBytes = []byte(`{"content":[{"type":"text","text":"ok"}]}`) + 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 + } + body, err := marshalToolCallEnvelope(tool, input, requestID, meta) + if err != nil { + return InvokeResult{}, err + } + + outcome, err := c.doHTTPToolCallStream(ctx, upstream, body, sink, 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 + } + outcome, err = c.doHTTPToolCallStream(ctx, upstream, body, sink, 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 +} + +// InvokeStream behaves like Invoke, except that if the upstream emits +// intermediate JSON-RPC messages (progress, logging, other notifications) +// for the call — as an SSE stream over HTTP, or as extra newline-delimited +// messages before the response over stdio — they are relayed to sink as +// they arrive, before the terminal response exists. When sink is nil, or +// the upstream never emits anything beyond its terminal response, sink is +// never called and the returned InvokeResult is identical to what Invoke +// would return. +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) } - 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 } func (c *Client) listToolsHTTP(ctx context.Context, upstream Upstream) ([]Tool, error) { @@ -692,10 +1336,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 +1364,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 +1386,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 +1399,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 +1648,48 @@ 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. +type boundedBuffer struct { + buf bytes.Buffer + limit int +} + +func newBoundedBuffer(limit int) *boundedBuffer { + return &boundedBuffer{limit: limit} +} + +func (b *boundedBuffer) Write(p []byte) (int, error) { + 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 { return b.buf.Len() } +func (b *boundedBuffer) String() string { 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 +1698,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 +1707,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 +1718,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())) @@ -1034,6 +1755,155 @@ func (c *Client) invokeStdio(ctx context.Context, upstream Upstream, tool string return InvokeResult{StatusCode: http.StatusOK, Body: body, Failed: looksLikeToolError(body)}, nil } +// invokeStdioStream is invokeStdio's live-relay counterpart. Unlike HTTP, +// stdio has no header phase or Content-Type to signal in advance whether +// the upstream will emit anything beyond its terminal response — every +// stdio reply is the same newline-delimited JSON-RPC framing regardless. +// So instead of a transport-level signal, StreamStarted fires lazily, the +// same way the HTTP path already does for its own edge case (see +// relaySSEToolCall): only right before the first actual notification or +// server-request is relayed. A call that produces nothing but its terminal +// response never touches sink at all, staying identical to invokeStdio. +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.armHeaderTimeout(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 := readRPC(reader, rpcIDMessage(initID)); err != nil { + if reason := guard.reason(); reason != "" { + return InvokeResult{}, fmt.Errorf("upstream %q: %s during stdio initialize: %w", upstream.Name, reason, ErrStreamTimeout) + } + if stderr.Len() > 0 { + return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String())) + } + return InvokeResult{}, err + } + guard.disarmHeaderTimeout() + _ = 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) +} + +// 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) (InvokeResult, error) { + expectedID := rpcIDMessage(callID) + started := false + ensureStarted := func() { + if !started { + started = true + sink.StreamStarted() + } + } + for { + line, err := reader.ReadBytes('\n') + if err != nil { + if reason := guard.reason(); reason != "" { + return InvokeResult{}, fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout) + } + if stderr.Len() > 0 { + return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String())) + } + return InvokeResult{}, err + } + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + guard.resetIdle() + + 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. + } + } +} + func (c *Client) listToolsStdio(ctx context.Context, upstream Upstream) ([]Tool, error) { if upstream.Command == "" { return nil, fmt.Errorf("stdio upstream %q missing command", upstream.Name) @@ -1043,6 +1913,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 +1922,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 +1932,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 +2123,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 +2134,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 +2145,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 +2154,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,29 +2166,46 @@ 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) { +// 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 + dataLines []string + eventID string + retry time.Duration + hasData bool + hasID bool + hasRetry bool +} + +type sseWireEvent struct { + Data []byte + ID string + Retry time.Duration + HasData bool + HasID bool + HasRetry bool +} + +func newSSEEventReader(r io.Reader) *sseEventReader { 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() + return &sseEventReader{scanner: scanner} +} + +// 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() { + line := r.scanner.Text() if line == "" { - if payload, ok := flush(); ok { - return payload, nil + if !r.hasData && !r.hasID && !r.hasRetry { + continue } - continue + return r.takeEvent(), nil } if strings.HasPrefix(line, ":") { continue @@ -1325,17 +2217,90 @@ func extractSSEJSONRPCResponse(r io.Reader, expectedID json.RawMessage) ([]byte, } else if strings.HasPrefix(value, " ") { value = strings.TrimPrefix(value, " ") } - if field == "data" { - dataLines = append(dataLines, 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 := scanner.Err(); err != nil { - return nil, err + if err := r.scanner.Err(); err != nil { + return sseWireEvent{}, err } - if payload, ok := flush(); ok { - return payload, nil + 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 + 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 + } } - return nil, fmt.Errorf("no JSON-RPC response in SSE stream") } type jsonRPCResponseMatch int @@ -1544,7 +2509,14 @@ 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) { for { line, err := reader.ReadBytes('\n') if err != nil { @@ -1554,11 +2526,11 @@ 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 @@ -1608,6 +2580,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 367cc6b4..63661477 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -5,11 +5,14 @@ import ( "context" "database/sql" "encoding/json" + "errors" "io" "log" "net/http" "net/http/httptest" "net/url" + "os" + "path/filepath" "strings" "sync" "sync/atomic" @@ -467,7 +470,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,6 +479,998 @@ func TestInvokeSkipsSSENotificationBeforeResponse(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 { + t.Fatalf("decode request: %v", err) + } + switch req.Method { + case "initialize": + 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/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) + } + })) + defer server.Close() + + client := NewHTTPClient() + client.httpClient = server.Client() + 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 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) + } +} + +// invokeStreamTestServer builds the initialize/notifications.initialized +// scaffolding shared by the InvokeStream tests below, dispatching tools/call +// to callHandler. +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) { + 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", 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.Fatalf("unexpected method %q", req.Method) + } + })) +} + +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) { + var resumeRequests int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + resumeRequests++ + 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 resumeRequests != 1 { + t.Fatalf("resume request count = %d, want 1", resumeRequests) + } + 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) + } +} + +// 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) + } +} + +// 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) { + 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": + sessionID := "sid-1" + if len(sessions) > 0 { + sessionID = "sid-2" + } + sessions = append(sessions, sessionID) + 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": + toolsCallCount++ + w.Header().Set("Content-Type", "text/event-stream") + flusher := w.(http.Flusher) + if toolsCallCount == 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) + } + if toolsCallCount != 2 { + t.Fatalf("tools/call count = %d, want 2", toolsCallCount) + } + if len(sessions) != 2 { + t.Fatalf("initialize sessions = %#v, want two sessions", sessions) + } + 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) { + var initializeCount int + 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": + initializeCount++ + 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++ + 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 toolsCallCount != 1 { + t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed)", toolsCallCount) + } + if initializeCount != 1 { + t.Fatalf("initialize count = %d, want 1 (no reinitialize attempt)", initializeCount) + } + 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)) + } +} + +// writeFakeStdioServer writes an executable bash script implementing the +// initialize/notifications.initialized handshake and dispatching tools/call +// to script (a bash fragment appended verbatim, given $line as the raw +// incoming JSON and able to compute its id via `id=$(echo "$line" | grep -o +// '"id":[0-9]*' | head -1 | cut -d: -f2)`). +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 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) + } +} + +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 TestCallTimeoutGuardHeaderTimeoutDisarmedAfterHeadersArrive(t *testing.T) { + g := newCallTimeoutGuard(context.Background()) + defer g.stop() + g.armHeaderTimeout(10 * time.Millisecond) + g.disarmHeaderTimeout() + time.Sleep(30 * time.Millisecond) + if reason := g.reason(); reason != "" { + t.Fatalf("expected a disarmed header 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) + } +} + func TestListToolsDecodesMultilineSSEData(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var req Envelope @@ -511,6 +1506,20 @@ func TestListToolsDecodesMultilineSSEData(t *testing.T) { } } +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 TestMissingSessionRPCErrorDetection(t *testing.T) { if !isMissingSessionRPCError(json.RawMessage(`{"code":-32000,"message":"No session ID provided for non-initialization request"}`)) { t.Fatal("expected missing session error to be detected") @@ -821,3 +1830,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/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) {} From fbe3f6ef57184e51531ec4025c1bcd51f84aa380 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 11:40:10 -0400 Subject: [PATCH 02/18] more testing --- CHANGELOG.md | 8 + Justfile | 18 + docs/architecture.md | 75 ++ internal/api/handlers_test.go | 9 + internal/api/mcp_everything_test.go | 170 +++++ .../api/mcp_external_fixture_helpers_test.go | 22 + .../api/mcp_external_process_unix_test.go | 33 + .../api/mcp_external_process_windows_test.go | 15 + internal/api/mcp_standalone_stream_test.go | 172 +++++ .../testdata/mcp_standalone_fixture/server.py | 37 + internal/mcp/client.go | 706 ++++++++++++++++-- internal/mcp/client_test.go | 514 +++++++++++++ 12 files changed, 1704 insertions(+), 75 deletions(-) create mode 100644 internal/api/mcp_everything_test.go create mode 100644 internal/api/mcp_external_fixture_helpers_test.go create mode 100644 internal/api/mcp_external_process_unix_test.go create mode 100644 internal/api/mcp_external_process_windows_test.go create mode 100644 internal/api/mcp_standalone_stream_test.go create mode 100644 internal/api/testdata/mcp_standalone_fixture/server.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a9b8467a..384cb1a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `stream_max_duration_seconds`, `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. ## [0.2.0] - 2026-07-14 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/docs/architecture.md b/docs/architecture.md index aae84bdf..17c7375f 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -219,6 +219,31 @@ sees them live instead of waiting in silence until the call finishes. If the upstream tool answers with a single plain response and no streaming, the agent simply receives that response. There are no live updates to relay in that case. +**Some upstream tools use a second, separate connection for progress updates instead of +the one above.** The wire protocol actually allows two different channels for anything a +tool sends before its final answer: updates riding along on the same connection as the +call itself (the channel described above), and a second, independent connection carrying +updates that aren't tied to any one specific call. Some widely-used tool-building +frameworks always use this second channel for progress updates rather than the first. +Atryum listens on both, so it makes no difference which channel a given upstream tool +happens to prefer — the agent sees the same live updates either way. + +**Because that second channel isn't tied to any one call, Atryum has to work out whose +update belongs to whom.** Several agents can have calls in flight against the same +upstream tool server at once, all sharing that one second channel. Each update on it +carries a tracking number the calling agent chose — but two unrelated agents could easily +pick the same tracking number, since neither knows about the other. Atryum doesn't know in +advance which of the two channels a given upstream will actually use to answer, so for +every call that asks for update tracking it swaps in its own guaranteed-unique tracking +number in place of whatever the agent supplied, before the call ever goes upstream — not +only for calls that end up using the second channel. It restores the agent's own original +number before handing an update back, the same way no matter which of the two channels +that update actually arrives on. That's what makes a coincidental match between two +unrelated agents' tracking numbers harmless either way. An update on the second channel +with no tracking number at all (a plain log-style message, say) is only ever handed to an +agent when exactly one call is currently sharing that channel; with more than one, there's +no way to know whose it is, and Atryum drops it rather than guess wrong. + **The three layers involved**, in order: the part of Atryum facing the agent decides whether to relay live and writes the response back; the part in the middle runs approval rules and keeps the audit trail; the part facing the upstream tool speaks the @@ -240,6 +265,7 @@ sequenceDiagram Client->>Upstream: Call the tool alt Upstream streams progress before answering Upstream-->>Client: Starts streaming + Note over Client,Upstream: some tools send updates on a second, separate
connection instead — Atryum listens on that one too Client->>Middle: A progress update arrived Middle->>Middle: Record it for the audit trail (in the background) Middle->>Facing: Forward the update @@ -259,6 +285,41 @@ sequenceDiagram end ``` +The diagram above shows updates arriving on the same connection as the call itself — +the most common case, and the only one some upstream tools use at all. Here's what +happens instead when an upstream tool sends its updates on the second, standalone +connection described earlier: + +```mermaid +sequenceDiagram + autonumber + participant Agent + participant Facing as Atryum: agent-facing layer + participant Middle as Atryum: rules & audit layer + participant Client as Atryum: upstream-facing layer + participant Upstream as Upstream tool server + + Agent->>Facing: Call a tool, willing to receive live updates + Facing->>Middle: Run the call + Middle->>Client: Send the call upstream + Client->>Client: Swap the agent's tracking number for a unique one + Client->>Upstream: Call the tool (this connection will carry the final answer) + Client->>Upstream: Open the second, standalone connection
(shared with any other call in flight against this upstream) + Note over Upstream: This upstream tool sends its updates on the standalone
connection, never on the call's own connection + Upstream-->>Client: Update arrives on the standalone connection + Client->>Client: Match it, by tracking number, back to this call + Client->>Middle: A progress update arrived (tracking number restored) + Middle->>Middle: Record it for the audit trail (in the background) + Middle->>Facing: Forward the update + Facing-->>Agent: Relay it live, with the agent's own original tracking number + Note over Upstream,Client: repeats for every update sent this way + Upstream-->>Client: Final answer, on the call's own connection + Client-->>Middle: Done + Middle->>Middle: Save the result, close out the audit trail + Middle-->>Facing: Done + Facing-->>Agent: Send the final answer +``` + **Approval gating applies before any streaming can start.** If a rule says a tool call needs a human to approve it first, Atryum pauses *before ever contacting the upstream tool* — which is before any streaming could even start. Nothing is sent to the agent @@ -354,6 +415,20 @@ support. restart takes effect — ending a call that was actually still healthy. Atryum double-checks how much time has *really* passed before deciding to end a call, so a well-timed update can never be wrongly punished by bad luck in the timing. +- **Two agents that happen to pick the same tracking number for their updates must + never get mixed up.** The second, standalone update channel described above is shared + across every call currently in flight against a given upstream tool server, and + agents don't know about each other's choices. Since Atryum can't tell in advance which + channel a given upstream will actually use to answer, it assigns its own + guaranteed-unique tracking number to *every* call that asks for update tracking, not + only ones that end up using the second channel, and restores the agent's original + number the same way regardless of which channel the update comes back on — so a + coincidental match can never cross-deliver one agent's update to another either way. +- **An upstream tool that doesn't support the second, standalone channel at all doesn't + affect the call itself.** Some tool servers simply don't offer it. Atryum notices on + the first attempt and stops trying again for the rest of that session, but the actual + tool call still completes and answers normally either way — the only thing lost is any + update that server would have sent exclusively on that unsupported channel. ## Decision-only calls diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index 441ec028..e0c77c8e 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -1656,6 +1656,15 @@ func readNextSSEFrame(t *testing.T, reader *bufio.Reader) sseEventFrame { 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) diff --git a/internal/api/mcp_everything_test.go b/internal/api/mcp_everything_test.go new file mode 100644 index 00000000..0b638ebc --- /dev/null +++ b/internal/api/mcp_everything_test.go @@ -0,0 +1,170 @@ +//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" + "context" + "database/sql" + "fmt" + "net" + "net/http" + "net/http/httptest" + "os/exec" + "strings" + "testing" + "time" + + "atryum/internal/config" + "atryum/internal/invocation" + "atryum/internal/invocation/policy" + "atryum/internal/mcp" + "atryum/internal/store" +) + +// 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) + deadline := time.Now().Add(30 * time.Second) + for time.Now().Before(deadline) { + conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 200*time.Millisecond) + if err == nil { + _ = conn.Close() + return baseURL + } + time.Sleep(200 * time.Millisecond) + } + 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 "" +} + +func TestMCPToolsCallAgainstRealEverythingServer(t *testing.T) { + baseURL := startEverythingServer(t) + + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("open db: %v", err) + } + defer 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: "everything", Mode: "http", BaseURL: baseURL, Enabled: true, TimeoutSeconds: 30}}, + }) + 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{}, 30*time.Second, nil, nil, nil, nil, + ) + 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()) + defer agentServer.Close() + + 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..0390ae6f --- /dev/null +++ b/internal/api/mcp_standalone_stream_test.go @@ -0,0 +1,172 @@ +//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" + "context" + "database/sql" + "fmt" + "net" + "net/http" + "net/http/httptest" + "os/exec" + "strings" + "testing" + "time" + + "atryum/internal/config" + "atryum/internal/invocation" + "atryum/internal/invocation/policy" + "atryum/internal/mcp" + "atryum/internal/store" +) + +// 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) + deadline := time.Now().Add(30 * time.Second) + for time.Now().Before(deadline) { + conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 200*time.Millisecond) + if err == nil { + _ = conn.Close() + return baseURL + } + time.Sleep(200 * time.Millisecond) + } + 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 "" +} + +func TestMCPToolsCallAgainstRealStandaloneStreamServer(t *testing.T) { + baseURL := startStandaloneFixtureServer(t) + + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("open db: %v", err) + } + defer 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: "standalone-fixture", Mode: "http", BaseURL: baseURL, Enabled: true, TimeoutSeconds: 30}}, + }) + 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{}, 30*time.Second, nil, nil, nil, nil, + ) + 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()) + defer agentServer.Close() + + 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/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/mcp/client.go b/internal/mcp/client.go index fcb761b7..dceabb99 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -253,6 +253,14 @@ type Client struct { sessionInitLocks map[string]*sync.Mutex sessions map[string]string sessionProtocols map[string]string + + // standaloneStreams holds, per upstream name, the shared "standalone" + // SSE GET connection used to receive server-initiated messages that + // aren't tied to any specific request — notably progress notifications + // from servers (e.g. the reference Python SDK) that don't attribute + // them to the request that triggered them. See standaloneStream. + standaloneMu sync.Mutex + standaloneStreams map[string]*standaloneStream } type InvokeResult struct { @@ -325,7 +333,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[string]*standaloneStream)} } func (r *Resolver) Resolve(name string) (Upstream, error) { @@ -632,9 +640,17 @@ func (c *Client) TestConnection(ctx context.Context, upstream Upstream) Connecti // 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 merged := mergeRequestMeta(meta, requestID); merged != nil { - params["_meta"] = merged + if meta != nil { + params["_meta"] = meta } return json.Marshal(Envelope{JSONRPC: "2.0", ID: json.RawMessage([]byte("1")), Method: "tools/call", Params: mustRawJSON(params)}) } @@ -960,7 +976,7 @@ type streamCallOutcome struct { // 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, opts StreamOptions) (streamCallOutcome, error) { +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.armHeaderTimeout(opts.HeaderTimeout) @@ -1010,33 +1026,88 @@ func (c *Client) doHTTPToolCallStream(ctx context.Context, upstream Upstream, bo // 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, guard, upstream, h.sessionID) + return c.relaySSEToolCall(resp, sink, progressCh, guard, upstream, h.sessionID) +} + +// 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) *postStreamPump { + p := &postStreamPump{msgs: make(chan postStreamMsg), current: resp, done: make(chan struct{})} + go p.run(c, guard, upstream) + 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() + } + }) } -// relaySSEToolCall reads resp's SSE body incrementally via an -// sseEventReader, relaying every intermediate (non-terminal) message to -// sink as it arrives, and returns once the terminal JSON-RPC response for -// id "1" is read. resp.Body is not closed here — the caller does that. -// sessionID is the session this attempt was sent under; it is always -// stamped onto the returned outcome (even a missing-session terminal -// response) so a caller retry can identify and clear the right session — -// mirroring doHTTPEnvelope's ForwardResult.SessionID contract. -// -// sink.StreamStarted fires lazily, right before the first thing is actually -// delivered — not simply because the response's Content-Type was SSE. This -// matters for the missing-session retry: if the very first (and only) -// message is a missing-session terminal error, the whole attempt is -// discarded and silently retried (see invokeHTTPStream), so the sink must -// never have been told a stream started for it. Once a real event has been -// relayed, or the terminal response is anything other than a -// zero-events missing-session error, the attempt is the one that counts and -// StreamStarted fires exactly once for it. -func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, guard *callTimeoutGuard, upstream Upstream, sessionID string) (streamCallOutcome, error) { - expectedID := json.RawMessage([]byte("1")) - currentResp := resp - reader := newSSEEventReader(currentResp.Body) - relayed := 0 - started := false +// setCurrent installs resp as the response the pump is currently reading +// from (after a resume). 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() + defer p.mu.Unlock() + if p.stopped { + return false + } + p.current = resp + 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: + } +} + +func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstream) { + defer close(p.msgs) + reader := newSSEEventReader(p.current.Body) lastEventID := "" retryDelay := time.Duration(0) // resumedFrom holds, after a resume, the cursor id the Last-Event-ID @@ -1046,41 +1117,49 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, guard *c // The guard window closes at the first event bearing any other id, so a // server legitimately reusing the id much later is unaffected. resumedFrom := "" - ensureStarted := func() { - if !started { - started = true - sink.StreamStarted() - } - } for { evt, err := reader.NextEvent() if err != nil { + p.mu.Lock() + stopped := p.stopped + p.mu.Unlock() + if stopped { + return + } if reason := guard.reason(); reason != "" { - _ = currentResp.Body.Close() - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q %s after %d relayed event(s): %w", upstream.Name, reason, relayed, ErrStreamTimeout) + p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout)}) + return } if err != io.EOF { - _ = currentResp.Body.Close() - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + p.send(postStreamMsg{err: err}) + return } - _ = currentResp.Body.Close() if lastEventID == "" { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q closed the stream without a JSON-RPC response or resumable event id", upstream.Name) + p.send(postStreamMsg{err: fmt.Errorf("upstream %q closed the stream without a JSON-RPC response or resumable event id", upstream.Name)}) + return } if err := waitForSSEReconnect(guard.ctx, retryDelay); err != nil { if reason := guard.reason(); reason != "" { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q %s while waiting to resume after %d relayed event(s): %w", upstream.Name, reason, relayed, ErrStreamTimeout) + p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while waiting to resume: %w", upstream.Name, reason, ErrStreamTimeout)}) + return } - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + p.send(postStreamMsg{err: err}) + return } - currentResp, err = c.resumeSSEStream(guard.ctx, upstream, lastEventID) + resumed, err := c.resumeSSEStream(guard.ctx, upstream, lastEventID) if err != nil { if reason := guard.reason(); reason != "" { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q %s while resuming after %d relayed event(s): %w", upstream.Name, reason, relayed, ErrStreamTimeout) + p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while resuming: %w", upstream.Name, reason, ErrStreamTimeout)}) + return } - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + p.send(postStreamMsg{err: err}) + return } - reader = newSSEEventReader(currentResp.Body) + if !p.setCurrent(resumed) { + _ = resumed.Body.Close() + return + } + reader = newSSEEventReader(resumed.Body) resumedFrom = lastEventID continue } @@ -1101,38 +1180,124 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, guard *c if !evt.HasData { continue } - payload := evt.Data + p.send(postStreamMsg{data: evt.Data}) + } +} - switch classifyRPCMessage(payload, expectedID) { - case rpcMessageTerminalResponse: - var rpcResp rpcResponse - if err := json.Unmarshal(payload, &rpcResp); err != nil { - _ = currentResp.Body.Close() - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err - } - invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, currentResp.StatusCode) - if !(missingSession && relayed == 0) { - ensureStarted() +// relaySSEToolCall reads resp's SSE body incrementally (via postStreamPump, +// on a dedicated goroutine), relaying every intermediate (non-terminal) +// message to sink as it arrives, and returns once the terminal JSON-RPC +// response for id "1" is read. It also drains progressCh — standalone- +// stream notifications matched to this call (see progressWaiter) — via the +// same select loop, so exactly one goroutine ever calls sink.Event for a +// given call. resp.Body is not closed here directly; postStreamPump owns +// that (including across resumes, which replace it with a new response). +// sessionID is the session this attempt was sent under; it is always +// stamped onto the returned outcome (even a missing-session terminal +// response) so a caller retry can identify and clear the right session — +// mirroring doHTTPEnvelope's ForwardResult.SessionID contract. +// +// sink.StreamStarted fires lazily, right before the first thing is actually +// delivered — not simply because the response's Content-Type was SSE. This +// matters for the missing-session retry: if the very first (and only) +// message is a missing-session terminal error, the whole attempt is +// discarded and silently retried (see invokeHTTPStream), so the sink must +// never have been told a stream started for it. Once a real event has been +// relayed, or the terminal response is anything other than a +// zero-events missing-session error, the attempt is the one that counts and +// StreamStarted fires exactly once for it. +func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progressCh <-chan StreamEvent, guard *callTimeoutGuard, upstream Upstream, sessionID string) (streamCallOutcome, error) { + expectedID := json.RawMessage([]byte("1")) + statusCode := resp.StatusCode + relayed := 0 + started := false + ensureStarted := func() { + if !started { + started = true + sink.StreamStarted() + } + } + + pump := newPostStreamPump(c, guard, upstream, resp) + 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 } - _ = currentResp.Body.Close() - return streamCallOutcome{invoke: invoke, missingSession: missingSession, eventsRelayed: relayed, sessionID: sessionID}, nil - case rpcMessageServerRequest: ensureStarted() relayed++ - if err := sink.Event(StreamEvent{Data: payload, ServerRequest: true}); err != nil { - _ = currentResp.Body.Close() + if err := sink.Event(evt); err != nil { return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err } - case rpcMessageNotification: - ensureStarted() - relayed++ - if err := sink.Event(StreamEvent{Data: payload}); err != nil { - _ = currentResp.Body.Close() - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + case msg, ok := <-pump.msgs: + if !ok { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q: stream ended unexpectedly", upstream.Name) + } + if msg.err != nil { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, msg.err + } + payload := msg.data + switch classifyRPCMessage(payload, expectedID) { + case rpcMessageTerminalResponse: + var rpcResp rpcResponse + if err := json.Unmarshal(payload, &rpcResp); err != nil { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, statusCode) + if progressCh != nil { + // See terminalSettleWindow: give a notification already in + // flight on the standalone stream a brief, bounded chance + // to arrive before finalizing. + settle := time.NewTimer(terminalSettleWindow) + settleLoop: + for { + select { + case evt, ok := <-progressCh: + if !ok { + break settleLoop + } + ensureStarted() + relayed++ + if err := sink.Event(evt); err != nil { + settle.Stop() + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + if !settle.Stop() { + <-settle.C + } + settle.Reset(terminalSettleWindow) + case <-settle.C: + break settleLoop + } + } + } + if !(missingSession && relayed == 0) { + ensureStarted() + } + return streamCallOutcome{invoke: invoke, missingSession: missingSession, eventsRelayed: relayed, sessionID: sessionID}, nil + case rpcMessageServerRequest: + ensureStarted() + relayed++ + if err := sink.Event(StreamEvent{Data: payload, ServerRequest: true}); err != nil { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + case rpcMessageNotification: + ensureStarted() + relayed++ + if err := sink.Event(StreamEvent{Data: payload}); err != nil { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + default: + // Unrecognized payload shape (e.g. a response to some other id). + // Not ours to interpret; ignore and keep reading. } - default: - // Unrecognized payload shape (e.g. a response to some other id). - // Not ours to interpret; ignore and keep reading. } } } @@ -1151,6 +1316,360 @@ func waitForSSEReconnect(ctx context.Context, delay time.Duration) error { } } +// progressWaiter is one streaming call's registration with a +// standaloneStream: a notification matching wireToken has its raw payload +// sent to events. events is buffered and drained only by +// relaySSEToolCall's own goroutine (via select, alongside that call's own +// POST-response reads — see postStreamPump) rather than delivered directly +// to the sink from the shared reader goroutine that owns routeStandaloneEvent. +// That indirection is what makes it safe: without it, a notification for +// this call and this call's own terminal response race across two +// independent goroutines with no ordering guarantee, and a delivery +// attempt could land after the call has already returned to its caller — +// which, at the HTTP handler layer, may already have written the terminal +// SSE frame and returned, making a later write to the same +// http.ResponseWriter unsafe. +type progressWaiter struct { + events chan StreamEvent +} + +// standaloneWaiterEventBuffer bounds progressWaiter.events. Sized generously +// relative to realistic progress-update rates: a full buffer means the +// receiving call's own goroutine isn't draining it (already finished, or +// deep in its own resume/retry handling), in which case routeStandaloneEvent +// drops the event rather than blocking — a shared reader goroutine also +// serving other concurrent calls must never block on one slow receiver. +const standaloneWaiterEventBuffer = 32 + +// standaloneStream manages one shared "standalone" SSE GET connection per +// upstream — the channel the MCP Streamable HTTP transport defines for +// server-initiated messages that aren't tied to any specific request. +// +// This exists because the reference MCP Python SDK's Context.report_progress +// does not attribute its notification to the request that triggered it (it +// calls send_progress_notification without related_request_id), so the +// server's message router sends it to this standalone stream, never to the +// tools/call POST response body that relaySSEToolCall reads. Without this, +// Atryum cannot see those notifications at all. +// +// Atryum multiplexes every downstream caller of a given upstream onto one +// shared session, so this stream is refcounted across concurrent streaming +// calls rather than opened per call: acquireStandaloneStream starts the +// connection for the first waiter and releaseStandaloneStream tears it down +// once the last waiter is gone. It deliberately does not implement +// Last-Event-ID resumption (unlike relaySSEToolCall's per-call stream): if +// the connection drops mid-flight, any calls still waiting on it simply stop +// receiving standalone-routed notifications until the next acquire cycle +// reopens it — an accepted limitation, not a correctness hazard, since the +// call's own terminal response still arrives normally on its POST stream. +// +// standaloneWaiterGracePeriod bounds how long a call's progressWaiter +// lingers in the waiters map after the call itself has completed, before +// invokeHTTPStream's deferred cleanup actually removes it. See that cleanup +// for why immediate removal is unsafe. +const standaloneWaiterGracePeriod = 2 * time.Second + +// terminalSettleWindow bounds how long relaySSEToolCall waits, once it has +// read this call's terminal response, for anything further to arrive on +// progressCh before finalizing — reset each time something does arrive, so +// a burst of trailing notifications is fully drained rather than cut off +// after one. This call's own POST response and the shared standalone +// stream are independent connections read by independent goroutines: even +// with progressWaiter's channel already holding a pending notification by +// the time the terminal is read, Go's select has no rule preferring one +// ready case over another, so without this window a notification that +// arrived at essentially the same instant as the terminal could be skipped +// — not because it never arrived, but because select happened not to pick +// it up first. +const terminalSettleWindow = 25 * time.Millisecond + +type standaloneStream struct { + 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 +} + +// 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 { + c.standaloneMu.Lock() + s := c.standaloneStreams[upstream.Name] + if s == nil { + s = &standaloneStream{waiters: make(map[string]progressWaiter)} + c.standaloneStreams[upstream.Name] = s + } + c.standaloneMu.Unlock() + + s.mu.Lock() + s.refCount++ + start := s.refCount == 1 && !s.unsupported + if start { + streamCtx, cancel := context.WithCancel(context.Background()) + s.cancel = cancel + s.done = make(chan struct{}) + go c.runStandaloneStream(streamCtx, upstream, s) + } + s.mu.Unlock() + return s +} + +// releaseStandaloneStream drops one reference acquired via +// acquireStandaloneStream. Once the last reference is gone, it cancels the +// reader goroutine, waits for it to fully exit, and evicts the entry so a +// future acquire opens a fresh connection (picking up, e.g., a session that +// was reinitialized in the meantime). +func (c *Client) releaseStandaloneStream(upstream Upstream, s *standaloneStream) { + 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 + } + s.mu.Unlock() + if cancel != nil { + cancel() + <-done + } + if last { + c.standaloneMu.Lock() + if c.standaloneStreams[upstream.Name] == s { + delete(c.standaloneStreams, upstream.Name) + } + c.standaloneMu.Unlock() + } +} + +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) (*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 := 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 standalone stream returned 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 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) { + defer close(s.done) + resp, err := c.openStandaloneGET(ctx, upstream) + if err != nil { + c.debugf("standalone stream unavailable server=%s err=%v", upstream.Name, err) + s.mu.Lock() + s.unsupported = true + s.mu.Unlock() + return + } + defer resp.Body.Close() + reader := newSSEEventReader(resp.Body) + for { + evt, err := reader.NextEvent() + if err != nil { + return + } + if !evt.HasData { + continue + } + c.routeStandaloneEvent(s, evt.Data) + } +} + +// 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 + } + 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}: + 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. + } +} + +// 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 wraps the caller's sink for one streaming call that requested +// progress tracking, restoring the caller's original progressToken in +// place of the wire-level token Atryum minted (see rewriteProgressToken) on +// every Event call — regardless of whether relaySSEToolCall read the +// message from the call's own POST response or from the standalone +// stream's per-call channel (see progressWaiter). Some upstreams echo a +// call's progress notifications on the tools/call POST response itself +// rather than the standalone stream — that's the more spec-typical case, +// in fact — so the restore can't live only in the standalone-delivery +// path, or the agent would see Atryum's internal token leak through there. +// +// relaySSEToolCall drains both sources from a single goroutine (see +// postStreamPump), so, unlike an earlier version of this type, Event and +// StreamStarted need no guard against concurrent calls. +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) +} + // 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 @@ -1220,12 +1739,49 @@ func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool s }); err != nil { return InvokeResult{}, err } - body, err := marshalToolCallEnvelope(tool, input, requestID, meta) + + merged := mergeRequestMeta(meta, requestID) + // effectiveSink is what actually gets passed to doHTTPToolCallStream. + // When this call requested progress tracking, it becomes a *callSink + // that restores the caller's original progressToken on every Event + // call, regardless of which of the two upstream channels (this call's + // own POST response, or the standalone stream via progressCh) the + // underlying message arrived on. + effectiveSink := sink + var progressCh chan StreamEvent + if rewritten, wireToken, original, ok := c.rewriteProgressToken(merged); ok { + merged = rewritten + effectiveSink = newCallSink(sink, wireToken, original) + progressCh = make(chan StreamEvent, standaloneWaiterEventBuffer) + standalone := c.acquireStandaloneStream(upstream) + standalone.registerWaiter(wireToken, progressWaiter{events: progressCh}) + defer func() { + c.releaseStandaloneStream(upstream, standalone) + // Deliberately not unregistered synchronously here: this call's + // own POST-response stream and the shared standalone connection + // are two independent connections read by two independent + // goroutines, with no ordering guarantee between them. A + // notification for this exact call can still be in flight on the + // standalone connection at the moment this call's own terminal + // response arrives — removing the waiter immediately risks the + // reader goroutine finding nothing for a notification that was + // legitimately on its way, silently dropping it. Wire tokens are + // never reused (always a fresh atomic counter value), so nothing + // is unsafe about the waiter lingering a little longer; delaying + // the removal trades a small, bounded amount of memory for + // closing that window. + time.AfterFunc(standaloneWaiterGracePeriod, func() { + standalone.unregisterWaiter(wireToken) + }) + }() + } + + body, err := marshalToolCallEnvelopeWithMeta(tool, input, merged) if err != nil { return InvokeResult{}, err } - outcome, err := c.doHTTPToolCallStream(ctx, upstream, body, sink, opts) + outcome, err := c.doHTTPToolCallStream(ctx, upstream, body, effectiveSink, progressCh, opts) if err != nil { return InvokeResult{}, err } @@ -1239,7 +1795,7 @@ func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool s }); retryErr != nil { return InvokeResult{}, retryErr } - outcome, err = c.doHTTPToolCallStream(ctx, upstream, body, sink, opts) + outcome, err = c.doHTTPToolCallStream(ctx, upstream, body, effectiveSink, progressCh, opts) if err != nil { return InvokeResult{}, err } diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 63661477..84de46a1 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -6,6 +6,7 @@ import ( "database/sql" "encoding/json" "errors" + "fmt" "io" "log" "net/http" @@ -616,6 +617,15 @@ func TestMergeRequestMeta(t *testing.T) { 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.Fatalf("decode request: %v", err) @@ -1121,6 +1131,510 @@ func TestInvokeStreamIdleTimeoutAbortsRead(t *testing.T) { } } +// syncFakeStreamSink is fakeStreamSink's mutex-protected counterpart. It's +// needed wherever a test can have both relaySSEToolCall's own read loop and +// routeStandaloneEvent deliver to the same sink concurrently — the plain +// fakeStreamSink above assumes single-goroutine delivery and would race. +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: + } +} + +// 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) + } +} + +// 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. +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(upstream, 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(upstream, 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(upstream, s3) +} + +// 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()) + } +} + // writeFakeStdioServer writes an executable bash script implementing the // initialize/notifications.initialized handshake and dispatching tools/call // to script (a bash fragment appended verbatim, given $line as the raw From 05e2b967928a19b670c61767431b33178a379fe3 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 12:46:35 -0400 Subject: [PATCH 03/18] fix: keep active standalone streams alive --- internal/mcp/client.go | 22 +++++------ internal/mcp/client_test.go | 78 +++++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 12 deletions(-) diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 84567a98..f3a7a3f9 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -1217,6 +1217,12 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress sink.StreamStarted() } } + deliver := func(evt StreamEvent) error { + guard.resetIdle() + ensureStarted() + relayed++ + return sink.Event(evt) + } pump := newPostStreamPump(c, guard, upstream, resp) defer pump.stop() @@ -1231,9 +1237,7 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress progressCh = nil continue } - ensureStarted() - relayed++ - if err := sink.Event(evt); err != nil { + if err := deliver(evt); err != nil { return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err } case msg, ok := <-pump.msgs: @@ -1263,9 +1267,7 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress if !ok { break settleLoop } - ensureStarted() - relayed++ - if err := sink.Event(evt); err != nil { + if err := deliver(evt); err != nil { settle.Stop() return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err } @@ -1283,15 +1285,11 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress } return streamCallOutcome{invoke: invoke, missingSession: missingSession, eventsRelayed: relayed, sessionID: sessionID}, nil case rpcMessageServerRequest: - ensureStarted() - relayed++ - if err := sink.Event(StreamEvent{Data: payload, ServerRequest: true}); err != nil { + if err := deliver(StreamEvent{Data: payload, ServerRequest: true}); err != nil { return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err } case rpcMessageNotification: - ensureStarted() - relayed++ - if err := sink.Event(StreamEvent{Data: payload}); err != nil { + if err := deliver(StreamEvent{Data: payload}); err != nil { return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err } default: diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 803b5afb..109013f3 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -1338,6 +1338,84 @@ func TestInvokeStreamStandaloneStreamRelaysProgressNotification(t *testing.T) { } } +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)) + } +} + // 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 From eb95cea61a1f16c2d35c55b365466c107f1cba77 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 12:47:42 -0400 Subject: [PATCH 04/18] docs: clarify streaming transport contracts --- internal/api/handlers.go | 18 ++++----- internal/config/config.go | 26 ++++++------- internal/invocation/service.go | 13 +++---- internal/mcp/client.go | 67 ++++++++++++++++------------------ 4 files changed, 59 insertions(+), 65 deletions(-) diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 8e0bf690..34c6bdb8 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -154,10 +154,10 @@ type Handler struct { apiKeyAuth auth.APIKeyConfig // streamRelayEnabled is the kill-switch for the tools/call SSE relay - // (see handleMCPProxy). The relay only ever activates when the agent's - // POST also sends Accept: text/event-stream and the upstream answers - // with an SSE body, so leaving this on by default is safe; it exists so - // the feature can be disabled globally without a rollback. + // (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` @@ -1534,11 +1534,11 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server toolReq.ClientName = snap.Name toolReq.ClientVersion = snap.Version } - // A stream-capable agent gets a relay sink; the response mode switch - // happens lazily, inside svc.InvokeStreaming, only if the upstream - // actually answers with an SSE stream (sink.StreamStarted). Until - // then nothing has been written, so a JSON upstream response still - // produces exactly today's buffered reply below. + // 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) diff --git a/internal/config/config.go b/internal/config/config.go index 011f703f..9b2689ea 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -118,23 +118,21 @@ type DefaultsConfig struct { // StreamRelayEnabled is the kill-switch for the tools/call SSE relay // (see api.Handler.SetStreamRelayEnabled / docs/architecture.md). The - // relay only ever activates when the agent's own request also sends - // Accept: text/event-stream and the upstream answers with an SSE - // body, so leaving this on by default is safe. + // 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 waiting for the upstream's - // response headers on a streaming tools/call — the connect phase, - // before Atryum knows whether the response will be an SSE stream. - // Zero falls back to RequestTimeoutSeconds, preserving today's - // connect-phase behavior. + // 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 the gap between successive relayed - // events once a stream has started; it resets on every event. Unlike - // RequestTimeoutSeconds, this does not bound the call's total - // duration — only how long it may go without producing anything. + // 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 the whole call once a stream has - // started. Zero disables the bound (unlimited). + // StreamMaxDurationSeconds bounds response reading after setup completes. + // Zero disables the bound (unlimited). StreamMaxDurationSeconds int `toml:"stream_max_duration_seconds"` // StreamAuditMaxEvents caps how many invocation.stream_event audit // rows get persisted per call; beyond the cap, events are still diff --git a/internal/invocation/service.go b/internal/invocation/service.go index 19511618..8680228a 100644 --- a/internal/invocation/service.go +++ b/internal/invocation/service.go @@ -218,8 +218,8 @@ type Service struct { pendingApprovals map[string]chan approvalDecision // streamOptions and streamAuditLimits govern InvokeStreaming's execution - // once a sink is present (see finishExecution). The zero value of each - // disables its bounds; SetStreamOptions installs real values. + // once a sink is present (see finishExecutionStreaming). The zero value + // of each disables its bounds; SetStreamOptions installs real values. streamOptions mcp.StreamOptions streamAuditLimits StreamAuditLimits @@ -263,7 +263,7 @@ func (s *Service) SetInvocationSummarizer(client SummaryClient) { // SetStreamOptions configures the header/idle/max-duration timeout scheme // and per-event audit caps used by InvokeStreaming once a caller supplies a -// sink (see finishExecution). Calls with a nil sink are unaffected. +// sink (see finishExecutionStreaming). Calls with a nil sink are unaffected. func (s *Service) SetStreamOptions(opts mcp.StreamOptions, auditLimits StreamAuditLimits) { s.streamOptions = opts s.streamAuditLimits = auditLimits @@ -330,10 +330,9 @@ func (s *Service) Invoke(ctx context.Context, req CreateInvocationRequest) (Invo return s.InvokeStreaming(ctx, req, nil) } -// InvokeStreaming runs one tool call exactly like Invoke, except that when -// sink is non-nil and the upstream answers the tools/call with an SSE -// stream, intermediate JSON-RPC messages (progress, logging, other -// notifications) are relayed to sink as they arrive — see finishExecution. +// 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. diff --git a/internal/mcp/client.go b/internal/mcp/client.go index f3a7a3f9..7ad64336 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -720,12 +720,13 @@ func (c *Client) invokeHTTP(ctx context.Context, upstream Upstream, tool string, return invoke, nil } -// StreamEvent is one upstream SSE event carrying a JSON-RPC message that is -// not the terminal response to the call: either a notification (progress, -// logging, or any other server-to-client notification) or, more rarely, a +// 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 the joined "data:" payload for this event: one JSON-RPC message. + // 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 @@ -735,15 +736,15 @@ type StreamEvent struct { ServerRequest bool } -// StreamSink receives upstream SSE events 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 +// 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 exactly once, before the first event is - // delivered or the terminal response is returned — never for an - // attempt that gets silently retried (see relaySSEToolCall). + // 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 @@ -751,18 +752,20 @@ type StreamSink interface { Event(evt StreamEvent) error } -// StreamOptions bounds how long InvokeStream may take once a stream has -// started. A zero-valued field disables that particular bound. +// StreamOptions bounds InvokeStream's setup and response-reading phases. A +// zero-valued field disables that particular bound. type StreamOptions struct { - // HeaderTimeout bounds waiting for the upstream's initial response - // headers, i.e. before we know whether the response is streaming. Zero - // leaves this phase bounded only by ctx's own deadline, if any. + // 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 the gap between successive events once the stream - // has started; it resets after every event. Zero disables the check. + // IdleTimeout bounds response-reading inactivity. Streaming transports + // reset it when upstream activity arrives, including events routed over + // the shared standalone HTTP stream. For a plain HTTP JSON response it + // bounds the complete body read. Zero disables the check. IdleTimeout time.Duration - // MaxDuration bounds the whole call once the stream has started. Zero - // disables the check. + // MaxDuration bounds the complete response-reading phase after HTTP + // headers or the stdio handshake. Zero disables the check. MaxDuration time.Duration } @@ -813,14 +816,10 @@ func classifyRPCMessage(payload []byte, expectedID json.RawMessage) rpcMessageKi return rpcMessageUnknown } -// callTimeoutGuard implements InvokeStream's header/idle/max-duration -// timeout scheme by canceling one shared context — the same context the -// HTTP request and its body reads run under. Header timing bounds only the -// wait for response headers; once headers arrive the caller disarms it and -// arms the idle/max-duration timers for the body-read phase instead, so a -// slow-to-start upstream and a slow-once-started upstream are judged against -// the right bound for each phase, rather than one fixed wall-clock budget -// covering both (which is what the plain per-call http.Client timeout does). +// 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 @@ -1805,14 +1804,12 @@ func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool s return outcome.invoke, nil } -// InvokeStream behaves like Invoke, except that if the upstream emits -// intermediate JSON-RPC messages (progress, logging, other notifications) -// for the call — as an SSE stream over HTTP, or as extra newline-delimited -// messages before the response over stdio — they are relayed to sink as -// they arrive, before the terminal response exists. When sink is nil, or -// the upstream never emits anything beyond its terminal response, sink is -// never called and the returned InvokeResult is identical to what Invoke -// would return. +// 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 SSE response contains +// only its terminal message. Stdio has no equivalent transport signal, so its +// sink starts 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: From 75567da6cd5fefe3d704dd45fd354f06300c20ac Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 12:54:54 -0400 Subject: [PATCH 05/18] refactor: separate streaming responsibilities Move transport, timeout, parsing, relay, and audit logic into focused files so each path can be understood and tested independently. --- internal/api/handlers.go | 193 --- internal/api/handlers_test.go | 525 ------ internal/api/sse_relay.go | 179 +++ internal/api/sse_relay_test.go | 543 +++++++ internal/invocation/service.go | 82 - internal/invocation/service_test.go | 554 ------- internal/invocation/stream_execution.go | 93 ++ internal/invocation/stream_execution_test.go | 572 +++++++ internal/mcp/client.go | 1401 ---------------- internal/mcp/client_test.go | 1505 ------------------ internal/mcp/http_stream.go | 473 ++++++ internal/mcp/http_stream_test.go | 572 +++++++ internal/mcp/sse_reader.go | 148 ++ internal/mcp/sse_reader_test.go | 59 + internal/mcp/standalone_stream.go | 314 ++++ internal/mcp/standalone_stream_test.go | 596 +++++++ internal/mcp/stdio_stream.go | 159 ++ internal/mcp/stdio_stream_test.go | 253 +++ internal/mcp/stream.go | 134 ++ internal/mcp/stream_timeout.go | 149 ++ internal/mcp/stream_timeout_test.go | 112 ++ 21 files changed, 4356 insertions(+), 4260 deletions(-) create mode 100644 internal/api/sse_relay.go create mode 100644 internal/api/sse_relay_test.go create mode 100644 internal/invocation/stream_execution.go create mode 100644 internal/invocation/stream_execution_test.go create mode 100644 internal/mcp/http_stream.go create mode 100644 internal/mcp/http_stream_test.go create mode 100644 internal/mcp/sse_reader.go create mode 100644 internal/mcp/sse_reader_test.go create mode 100644 internal/mcp/standalone_stream.go create mode 100644 internal/mcp/standalone_stream_test.go create mode 100644 internal/mcp/stdio_stream.go create mode 100644 internal/mcp/stdio_stream_test.go create mode 100644 internal/mcp/stream.go create mode 100644 internal/mcp/stream_timeout.go create mode 100644 internal/mcp/stream_timeout_test.go diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 34c6bdb8..686b93ae 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -1,7 +1,6 @@ package api import ( - "bytes" "context" "crypto/rand" "database/sql" @@ -1038,198 +1037,6 @@ func writeSSEComment(w io.Writer, comment string) { fmt.Fprint(w, "\n") } -// writeSSEEvent writes and flushes one SSE frame. Per the SSE spec, a -// multi-line payload must be sent as one "data:" line per line of content — -// a raw newline embedded in a single "data:" line breaks framing, since any -// continuation line lacking its own field prefix is ignored by a compliant -// parser. This matters here: a relayed notification's data (evt.Data, -// reconstructed by mcp.sseEventReader) can genuinely be multi-line if the -// upstream sent it that way (see mcp.TestListToolsDecodesMultilineSSEData -// for a real example) — only the terminal frame (always compact -// json.Marshal output) is guaranteed single-line. No "id:" field is ever -// emitted: doing so implies Last-Event-ID resumability, which Atryum does -// not support and must not advertise. Returns whatever error the write -// itself produced, which is how a broken downstream connection (the agent -// disconnected) is detected. -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 bytes.Split(data, []byte("\n")) { - 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 implements mcp.StreamSink for one agent-facing tools/call -// request. It relays every intermediate upstream event to the agent as an -// SSE frame, live, as InvokeStreaming reads it from the upstream, and keeps -// the downstream connection alive with heartbeat comments while the -// upstream is silent. -// -// The downstream response only switches to SSE mode when StreamStarted -// fires — until then, handleMCPProxy has written nothing, so a JSON -// (non-streaming) upstream response still produces today's exact buffered -// reply. Once started is true, headers have already been sent: the caller -// must never call writeRPCResult/writeRPCError/WriteHeader again for this -// request; the terminal response is written via finishStream instead. -// -// Concurrency: StreamStarted/Event run synchronously on the handler -// goroutine (inside InvokeStreaming's call stack); the heartbeat runs on -// its own goroutine. mu serializes every write to w so a heartbeat can -// never interleave with (and corrupt) an event or terminal frame. -// finishStream stops the heartbeat before writing the terminal frame, so -// no write can occur after the handler returns. -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.eventCount++ - s.setWriteDeadlineLocked() - if err := writeSSEEvent(s.w, s.flusher, "", evt.Data); err != nil { - s.writeErr = err - return err - } - return nil -} - -// 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.stopOnce.Do(func() { close(s.heartbeatStop) }) - <-s.heartbeatDone - 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 -} - func isJSONRPCRequest(r *http.Request) bool { if strings.Contains(strings.ToLower(r.Header.Get("Content-Type")), "application/json") { return true diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index 03d5b6c5..1ef49793 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -1,7 +1,6 @@ package api import ( - "bufio" "bytes" "context" "database/sql" @@ -12,7 +11,6 @@ import ( "net/http/httptest" "net/url" "strings" - "sync/atomic" "testing" "time" @@ -20,7 +18,6 @@ import ( backendclient "github.com/validmind/atryum/internal/backend" "github.com/validmind/atryum/internal/config" "github.com/validmind/atryum/internal/invocation" - "github.com/validmind/atryum/internal/invocation/policy" "github.com/validmind/atryum/internal/managedagents" "github.com/validmind/atryum/internal/mcp" "github.com/validmind/atryum/internal/store" @@ -1249,528 +1246,6 @@ func TestMCPToolsCallForwardsMetaToInvocation(t *testing.T) { } } -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. -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 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, "persisting result failed") { - t.Fatalf("expected a terminal JSON-RPC error frame, got %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. -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() - - db, err := sql.Open("sqlite", ":memory:") - if err != nil { - t.Fatalf("open db: %v", err) - } - defer 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: "demo", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}, - }) - 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{}, 5*time.Second, nil, nil, nil, nil, - ) - - h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil) - agentServer := httptest.NewServer(h.Routes()) - defer agentServer.Close() - - 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) - } -} - 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/sse_relay.go b/internal/api/sse_relay.go new file mode 100644 index 00000000..983ce668 --- /dev/null +++ b/internal/api/sse_relay.go @@ -0,0 +1,179 @@ +package api + +import ( + "bytes" + "fmt" + "io" + "net/http" + "sync" + "time" + + "github.com/validmind/atryum/internal/mcp" +) + +// 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 bytes.Split(data, []byte("\n")) { + 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.eventCount++ + s.setWriteDeadlineLocked() + if err := writeSSEEvent(s.w, s.flusher, "", evt.Data); err != nil { + s.writeErr = err + return err + } + return nil +} + +// 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.stopOnce.Do(func() { close(s.heartbeatStop) }) + <-s.heartbeatDone + 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..7f54b9b6 --- /dev/null +++ b/internal/api/sse_relay_test.go @@ -0,0 +1,543 @@ +package api + +import ( + "bufio" + "context" + "database/sql" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "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" +) + +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. +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 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, "persisting result failed") { + t.Fatalf("expected a terminal JSON-RPC error frame, got %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. +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() + + db, err := sql.Open("sqlite", ":memory:") + if err != nil { + t.Fatalf("open db: %v", err) + } + defer 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: "demo", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}, + }) + 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{}, 5*time.Second, nil, nil, nil, nil, + ) + + h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil) + agentServer := httptest.NewServer(h.Routes()) + defer agentServer.Close() + + 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) + } +} diff --git a/internal/invocation/service.go b/internal/invocation/service.go index 8680228a..8794ed88 100644 --- a/internal/invocation/service.go +++ b/internal/invocation/service.go @@ -1197,88 +1197,6 @@ func (s *Service) finishExecutionBuffered(ctx context.Context, inv Invocation, u return s.toResponse(inv), nil } -// finishExecutionStreaming is finishExecution's live-relay path: it wraps -// the caller's sink in an auditing decorator (so every relayed event and -// the call's outcome are recorded as invocation_events rows regardless of -// whether the downstream write later fails) and calls InvokeStream with -// s.streamOptions instead of the fixed s.defaultTimeout. -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.streamAuditLimits) - result, err := s.client.InvokeStream(ctx, upstream, req.Tool, req.Input, req.RequestID, req.Meta, audited, s.streamOptions) - completed := time.Now().UTC() - inv.CompletedAt = &completed - - if err != nil { - inv.Status = StatusFailed - reason, message := classifyStreamError(audited, err) - inv.Error = mustJSON(map[string]any{"message": message}) - audited.finish(completed, "failed") - persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout) - defer cancelPersist() - if updateErr := s.invocations.UpdateResult(persistCtx, inv); updateErr != nil { - return InvocationResponse{}, fmt.Errorf("persist streaming invocation failure: %w", updateErr) - } - // 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_relayed": audited.seq}), - CreatedAt: completed, - }) - return s.toResponse(inv), nil - } - var terminalEvent Event - if result.Failed { - inv.Status = StatusFailed - inv.Error = result.Body - audited.finish(completed, "failed") - 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 - audited.finish(completed, "succeeded") - 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 { - return InvocationResponse{}, err - } - _ = 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(audited *auditingSink, err error) (reason string, message string) { - if audited.downstreamErr != nil { - return "stream_aborted_downstream", audited.downstreamErr.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() - } - return "transport_error", err.Error() -} - func (s *Service) Approve(ctx context.Context, invocationID string, actorID string) error { s.mu.Lock() ch, ok := s.pendingApprovals[invocationID] diff --git a/internal/invocation/service_test.go b/internal/invocation/service_test.go index 1727c531..629f25be 100644 --- a/internal/invocation/service_test.go +++ b/internal/invocation/service_test.go @@ -5,7 +5,6 @@ import ( "database/sql" "encoding/json" "errors" - "fmt" "net/http" "net/http/httptest" "os" @@ -244,559 +243,6 @@ func TestSubmitLogsAndAuditsRuleLoadFailure(t *testing.T) { } } -// 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 -} - -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.Fatalf("decode request: %v", err) - } - 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.Fatalf("unexpected method %q", body["method"]) - } - })) -} - -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) - } -} - -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 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) { - var toolsCallCount int - upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) { - toolsCallCount++ - 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 toolsCallCount != 1 { - t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed mid-stream)", toolsCallCount) - } - - 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") - } -} - -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() { - time.Sleep(50 * time.Millisecond) - if sink.touched() { - t.Errorf("sink touched before approval — approval gating must precede any relay") - } - list, err := service.List(context.Background(), invocation.InvocationListFilter{Limit: 10}) - if err != nil || len(list.Items) == 0 { - t.Errorf("expected a pending invocation to approve") - return - } - if err := service.Approve(context.Background(), list.Items[0].InvocationID, ""); 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) - } -} - func TestResolverBootstrapsServersFromConfigWhenDBEmpty(t *testing.T) { db := newSQLiteTestDB(t) repo := store.NewServerRepo(db) diff --git a/internal/invocation/stream_execution.go b/internal/invocation/stream_execution.go new file mode 100644 index 00000000..3e79babc --- /dev/null +++ b/internal/invocation/stream_execution.go @@ -0,0 +1,93 @@ +package invocation + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/validmind/atryum/internal/mcp" +) + +// finishExecutionStreaming is finishExecution's live-relay path: it wraps +// the caller's sink in an auditing decorator (so every relayed event and +// the call's outcome are recorded as invocation_events rows regardless of +// whether the downstream write later fails) and calls InvokeStream with +// s.streamOptions instead of the fixed s.defaultTimeout. +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.streamAuditLimits) + result, err := s.client.InvokeStream(ctx, upstream, req.Tool, req.Input, req.RequestID, req.Meta, audited, s.streamOptions) + completed := time.Now().UTC() + inv.CompletedAt = &completed + + if err != nil { + inv.Status = StatusFailed + reason, message := classifyStreamError(audited, err) + inv.Error = mustJSON(map[string]any{"message": message}) + audited.finish(completed, "failed") + persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout) + defer cancelPersist() + if updateErr := s.invocations.UpdateResult(persistCtx, inv); updateErr != nil { + return InvocationResponse{}, fmt.Errorf("persist streaming invocation failure: %w", updateErr) + } + // 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_relayed": audited.seq}), + CreatedAt: completed, + }) + return s.toResponse(inv), nil + } + var terminalEvent Event + if result.Failed { + inv.Status = StatusFailed + inv.Error = result.Body + audited.finish(completed, "failed") + 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 + audited.finish(completed, "succeeded") + 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 { + return InvocationResponse{}, err + } + _ = 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(audited *auditingSink, err error) (reason string, message string) { + if audited.downstreamErr != nil { + return "stream_aborted_downstream", audited.downstreamErr.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() + } + 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..6641dc32 --- /dev/null +++ b/internal/invocation/stream_execution_test.go @@ -0,0 +1,572 @@ +package invocation_test + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "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 +} + +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.Fatalf("decode request: %v", err) + } + 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.Fatalf("unexpected method %q", body["method"]) + } + })) +} + +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) + } +} + +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 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) { + var toolsCallCount int + upstream := sseToolCallUpstream(t, func(w http.ResponseWriter, r *http.Request, body map[string]any) { + toolsCallCount++ + 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 toolsCallCount != 1 { + t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed mid-stream)", toolsCallCount) + } + + 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") + } +} + +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() { + time.Sleep(50 * time.Millisecond) + if sink.touched() { + t.Errorf("sink touched before approval — approval gating must precede any relay") + } + list, err := service.List(context.Background(), invocation.InvocationListFilter{Limit: 10}) + if err != nil || len(list.Items) == 0 { + t.Errorf("expected a pending invocation to approve") + return + } + if err := service.Approve(context.Background(), list.Items[0].InvocationID, ""); 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) + } +} diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 7ad64336..5aee47df 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -720,1121 +720,6 @@ func (c *Client) invokeHTTP(ctx context.Context, upstream Upstream, tool string, return invoke, nil } -// 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 -} - -// 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, including events routed over - // the shared standalone HTTP stream. 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 -} - -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 -} - -// 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 - headerTimer *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) armHeaderTimeout(d time.Duration) { - if d <= 0 { - return - } - g.mu.Lock() - defer g.mu.Unlock() - if g.stopped { - return - } - g.headerTimer = time.AfterFunc(d, func() { g.trip("timed out waiting for upstream response headers") }) -} - -func (g *callTimeoutGuard) disarmHeaderTimeout() { - g.mu.Lock() - defer g.mu.Unlock() - if g.headerTimer != nil { - g.headerTimer.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.headerTimer != nil { - g.headerTimer.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 -} - -// 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.armHeaderTimeout(opts.HeaderTimeout) - - h, err := c.doHTTPEnvelopeHeaders(guard.ctx, upstream, body, DefaultMCPProtocolVersion, true) - guard.disarmHeaderTimeout() - if err != nil { - if reason := guard.reason(); reason != "" { - return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) - } - return streamCallOutcome{}, 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() - _, _ = io.Copy(io.Discard, resp.Body) - return streamCallOutcome{missingSession: true, sessionID: h.sessionID}, nil - } - - if !strings.Contains(strings.ToLower(h.contentType), "text/event-stream") { - defer resp.Body.Close() - bodyBytes, err := io.ReadAll(resp.Body) - if err != nil { - if reason := guard.reason(); reason != "" { - return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) - } - return streamCallOutcome{}, 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) -} - -// 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) *postStreamPump { - p := &postStreamPump{msgs: make(chan postStreamMsg), current: resp, done: make(chan struct{})} - go p.run(c, guard, upstream) - 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). 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() - defer p.mu.Unlock() - if p.stopped { - return false - } - p.current = resp - 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: - } -} - -func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstream) { - defer close(p.msgs) - reader := newSSEEventReader(p.current.Body) - lastEventID := "" - retryDelay := time.Duration(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 reason := guard.reason(); reason != "" { - p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout)}) - 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 - } - if err := waitForSSEReconnect(guard.ctx, retryDelay); err != nil { - if reason := guard.reason(); reason != "" { - p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while waiting to resume: %w", upstream.Name, reason, ErrStreamTimeout)}) - return - } - p.send(postStreamMsg{err: err}) - return - } - resumed, err := c.resumeSSEStream(guard.ctx, upstream, lastEventID) - if err != nil { - if reason := guard.reason(); reason != "" { - p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while resuming: %w", upstream.Name, reason, ErrStreamTimeout)}) - return - } - p.send(postStreamMsg{err: err}) - return - } - if !p.setCurrent(resumed) { - _ = resumed.Body.Close() - return - } - reader = newSSEEventReader(resumed.Body) - resumedFrom = lastEventID - continue - } - guard.resetIdle() - 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 reads resp's SSE body incrementally (via postStreamPump, -// on a dedicated goroutine), relaying every intermediate (non-terminal) -// message to sink as it arrives, and returns once the terminal JSON-RPC -// response for id "1" is read. It also drains progressCh — standalone- -// stream notifications matched to this call (see progressWaiter) — via the -// same select loop, so exactly one goroutine ever calls sink.Event for a -// given call. resp.Body is not closed here directly; postStreamPump owns -// that (including across resumes, which replace it with a new response). -// sessionID is the session this attempt was sent under; it is always -// stamped onto the returned outcome (even a missing-session terminal -// response) so a caller retry can identify and clear the right session — -// mirroring doHTTPEnvelope's ForwardResult.SessionID contract. -// -// sink.StreamStarted fires lazily, right before the first thing is actually -// delivered — not simply because the response's Content-Type was SSE. This -// matters for the missing-session retry: if the very first (and only) -// message is a missing-session terminal error, the whole attempt is -// discarded and silently retried (see invokeHTTPStream), so the sink must -// never have been told a stream started for it. Once a real event has been -// relayed, or the terminal response is anything other than a -// zero-events missing-session error, the attempt is the one that counts and -// StreamStarted fires exactly once for it. -func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progressCh <-chan StreamEvent, guard *callTimeoutGuard, upstream Upstream, sessionID string) (streamCallOutcome, error) { - expectedID := json.RawMessage([]byte("1")) - 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) - } - - pump := newPostStreamPump(c, guard, upstream, resp) - 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 streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err - } - case msg, ok := <-pump.msgs: - if !ok { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q: stream ended unexpectedly", upstream.Name) - } - if msg.err != nil { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, msg.err - } - payload := msg.data - switch classifyRPCMessage(payload, expectedID) { - case rpcMessageTerminalResponse: - var rpcResp rpcResponse - if err := json.Unmarshal(payload, &rpcResp); err != nil { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err - } - invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, statusCode) - if progressCh != nil { - // See terminalSettleWindow: give a notification already in - // flight on the standalone stream a brief, bounded chance - // to arrive before finalizing. - settle := time.NewTimer(terminalSettleWindow) - settleLoop: - for { - select { - case evt, ok := <-progressCh: - if !ok { - break settleLoop - } - if err := deliver(evt); err != nil { - settle.Stop() - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err - } - if !settle.Stop() { - <-settle.C - } - settle.Reset(terminalSettleWindow) - case <-settle.C: - break settleLoop - } - } - } - 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 streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err - } - case rpcMessageNotification: - if err := deliver(StreamEvent{Data: payload}); err != nil { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err - } - default: - // Unrecognized payload shape (e.g. a response to some other id). - // Not ours to interpret; ignore and keep reading. - } - } - } -} - -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() - } -} - -// progressWaiter is one streaming call's registration with a -// standaloneStream: a notification matching wireToken has its raw payload -// sent to events. events is buffered and drained only by -// relaySSEToolCall's own goroutine (via select, alongside that call's own -// POST-response reads — see postStreamPump) rather than delivered directly -// to the sink from the shared reader goroutine that owns routeStandaloneEvent. -// That indirection is what makes it safe: without it, a notification for -// this call and this call's own terminal response race across two -// independent goroutines with no ordering guarantee, and a delivery -// attempt could land after the call has already returned to its caller — -// which, at the HTTP handler layer, may already have written the terminal -// SSE frame and returned, making a later write to the same -// http.ResponseWriter unsafe. -type progressWaiter struct { - events chan StreamEvent -} - -// standaloneWaiterEventBuffer bounds progressWaiter.events. Sized generously -// relative to realistic progress-update rates: a full buffer means the -// receiving call's own goroutine isn't draining it (already finished, or -// deep in its own resume/retry handling), in which case routeStandaloneEvent -// drops the event rather than blocking — a shared reader goroutine also -// serving other concurrent calls must never block on one slow receiver. -const standaloneWaiterEventBuffer = 32 - -// standaloneStream manages one shared "standalone" SSE GET connection per -// upstream — the channel the MCP Streamable HTTP transport defines for -// server-initiated messages that aren't tied to any specific request. -// -// This exists because the reference MCP Python SDK's Context.report_progress -// does not attribute its notification to the request that triggered it (it -// calls send_progress_notification without related_request_id), so the -// server's message router sends it to this standalone stream, never to the -// tools/call POST response body that relaySSEToolCall reads. Without this, -// Atryum cannot see those notifications at all. -// -// Atryum multiplexes every downstream caller of a given upstream onto one -// shared session, so this stream is refcounted across concurrent streaming -// calls rather than opened per call: acquireStandaloneStream starts the -// connection for the first waiter and releaseStandaloneStream tears it down -// once the last waiter is gone. It deliberately does not implement -// Last-Event-ID resumption (unlike relaySSEToolCall's per-call stream): if -// the connection drops mid-flight, any calls still waiting on it simply stop -// receiving standalone-routed notifications until the next acquire cycle -// reopens it — an accepted limitation, not a correctness hazard, since the -// call's own terminal response still arrives normally on its POST stream. -// -// standaloneWaiterGracePeriod bounds how long a call's progressWaiter -// lingers in the waiters map after the call itself has completed, before -// invokeHTTPStream's deferred cleanup actually removes it. See that cleanup -// for why immediate removal is unsafe. -const standaloneWaiterGracePeriod = 2 * time.Second - -// terminalSettleWindow bounds how long relaySSEToolCall waits, once it has -// read this call's terminal response, for anything further to arrive on -// progressCh before finalizing — reset each time something does arrive, so -// a burst of trailing notifications is fully drained rather than cut off -// after one. This call's own POST response and the shared standalone -// stream are independent connections read by independent goroutines: even -// with progressWaiter's channel already holding a pending notification by -// the time the terminal is read, Go's select has no rule preferring one -// ready case over another, so without this window a notification that -// arrived at essentially the same instant as the terminal could be skipped -// — not because it never arrived, but because select happened not to pick -// it up first. -const terminalSettleWindow = 25 * time.Millisecond - -type standaloneStream struct { - 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 -} - -// 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 { - c.standaloneMu.Lock() - s := c.standaloneStreams[upstream.Name] - if s == nil { - s = &standaloneStream{waiters: make(map[string]progressWaiter)} - c.standaloneStreams[upstream.Name] = s - } - c.standaloneMu.Unlock() - - s.mu.Lock() - s.refCount++ - start := s.refCount == 1 && !s.unsupported - if start { - streamCtx, cancel := context.WithCancel(context.Background()) - s.cancel = cancel - s.done = make(chan struct{}) - go c.runStandaloneStream(streamCtx, upstream, s) - } - s.mu.Unlock() - return s -} - -// releaseStandaloneStream drops one reference acquired via -// acquireStandaloneStream. Once the last reference is gone, it cancels the -// reader goroutine, waits for it to fully exit, and evicts the entry so a -// future acquire opens a fresh connection (picking up, e.g., a session that -// was reinitialized in the meantime). -func (c *Client) releaseStandaloneStream(upstream Upstream, s *standaloneStream) { - 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 - } - s.mu.Unlock() - if cancel != nil { - cancel() - <-done - } - if last { - c.standaloneMu.Lock() - if c.standaloneStreams[upstream.Name] == s { - delete(c.standaloneStreams, upstream.Name) - } - c.standaloneMu.Unlock() - } -} - -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) (*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 := 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 standalone stream returned 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 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) { - defer close(s.done) - resp, err := c.openStandaloneGET(ctx, upstream) - if err != nil { - c.debugf("standalone stream unavailable server=%s err=%v", upstream.Name, err) - s.mu.Lock() - s.unsupported = true - s.mu.Unlock() - return - } - defer resp.Body.Close() - reader := newSSEEventReader(resp.Body) - for { - evt, err := reader.NextEvent() - if err != nil { - return - } - if !evt.HasData { - continue - } - c.routeStandaloneEvent(s, evt.Data) - } -} - -// 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 - } - 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}: - 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. - } -} - -// 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 wraps the caller's sink for one streaming call that requested -// progress tracking, restoring the caller's original progressToken in -// place of the wire-level token Atryum minted (see rewriteProgressToken) on -// every Event call — regardless of whether relaySSEToolCall read the -// message from the call's own POST response or from the standalone -// stream's per-call channel (see progressWaiter). Some upstreams echo a -// call's progress notifications on the tools/call POST response itself -// rather than the standalone stream — that's the more spec-typical case, -// in fact — so the restore can't live only in the standalone-delivery -// path, or the agent would see Atryum's internal token leak through there. -// -// relaySSEToolCall drains both sources from a single goroutine (see -// postStreamPump), so, unlike an earlier version of this type, Event and -// StreamStarted need no guard against concurrent calls. -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) -} - -// 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) - // effectiveSink is what actually gets passed to doHTTPToolCallStream. - // When this call requested progress tracking, it becomes a *callSink - // that restores the caller's original progressToken on every Event - // call, regardless of which of the two upstream channels (this call's - // own POST response, or the standalone stream via progressCh) the - // underlying message arrived on. - effectiveSink := sink - var progressCh chan StreamEvent - if rewritten, wireToken, original, ok := c.rewriteProgressToken(merged); ok { - merged = rewritten - effectiveSink = newCallSink(sink, wireToken, original) - progressCh = make(chan StreamEvent, standaloneWaiterEventBuffer) - standalone := c.acquireStandaloneStream(upstream) - standalone.registerWaiter(wireToken, progressWaiter{events: progressCh}) - defer func() { - c.releaseStandaloneStream(upstream, standalone) - // Deliberately not unregistered synchronously here: this call's - // own POST-response stream and the shared standalone connection - // are two independent connections read by two independent - // goroutines, with no ordering guarantee between them. A - // notification for this exact call can still be in flight on the - // standalone connection at the moment this call's own terminal - // response arrives — removing the waiter immediately risks the - // reader goroutine finding nothing for a notification that was - // legitimately on its way, silently dropping it. Wire tokens are - // never reused (always a fresh atomic counter value), so nothing - // is unsafe about the waiter lingering a little longer; delaying - // the removal trades a small, bounded amount of memory for - // closing that window. - time.AfterFunc(standaloneWaiterGracePeriod, func() { - standalone.unregisterWaiter(wireToken) - }) - }() - } - - 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 - } - 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 -} - -// 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 SSE response contains -// only its terminal message. Stdio has no equivalent transport signal, so its -// sink starts 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) - } -} - func (c *Client) listToolsHTTP(ctx context.Context, upstream Upstream) ([]Tool, error) { if err := c.ensureHTTPSession(ctx, upstream); err != nil { return nil, err @@ -2306,155 +1191,6 @@ func (c *Client) invokeStdio(ctx context.Context, upstream Upstream, tool string return InvokeResult{StatusCode: http.StatusOK, Body: body, Failed: looksLikeToolError(body)}, nil } -// invokeStdioStream is invokeStdio's live-relay counterpart. Unlike HTTP, -// stdio has no header phase or Content-Type to signal in advance whether -// the upstream will emit anything beyond its terminal response — every -// stdio reply is the same newline-delimited JSON-RPC framing regardless. -// So instead of a transport-level signal, StreamStarted fires lazily, the -// same way the HTTP path already does for its own edge case (see -// relaySSEToolCall): only right before the first actual notification or -// server-request is relayed. A call that produces nothing but its terminal -// response never touches sink at all, staying identical to invokeStdio. -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.armHeaderTimeout(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 := readRPC(reader, rpcIDMessage(initID)); err != nil { - if reason := guard.reason(); reason != "" { - return InvokeResult{}, fmt.Errorf("upstream %q: %s during stdio initialize: %w", upstream.Name, reason, ErrStreamTimeout) - } - if stderr.Len() > 0 { - return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String())) - } - return InvokeResult{}, err - } - guard.disarmHeaderTimeout() - _ = 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) -} - -// 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) (InvokeResult, error) { - expectedID := rpcIDMessage(callID) - started := false - ensureStarted := func() { - if !started { - started = true - sink.StreamStarted() - } - } - for { - line, err := reader.ReadBytes('\n') - if err != nil { - if reason := guard.reason(); reason != "" { - return InvokeResult{}, fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout) - } - if stderr.Len() > 0 { - return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String())) - } - return InvokeResult{}, err - } - line = bytes.TrimSpace(line) - if len(line) == 0 { - continue - } - guard.resetIdle() - - 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. - } - } -} - func (c *Client) listToolsStdio(ctx context.Context, upstream Upstream) ([]Tool, error) { if upstream.Command == "" { return nil, fmt.Errorf("stdio upstream %q missing command", upstream.Name) @@ -2717,143 +1453,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} } -// 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 - dataLines []string - eventID string - retry time.Duration - hasData bool - hasID bool - hasRetry bool -} - -type sseWireEvent struct { - Data []byte - ID string - Retry time.Duration - HasData bool - HasID bool - HasRetry bool -} - -func newSSEEventReader(r io.Reader) *sseEventReader { - scanner := bufio.NewScanner(r) - scanner.Buffer(make([]byte, 1024*1024), 4*1024*1024) - return &sseEventReader{scanner: scanner} -} - -// 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() { - line := r.scanner.Text() - if line == "" { - if !r.hasData && !r.hasID && !r.hasRetry { - continue - } - return r.takeEvent(), nil - } - 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 { - 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 - 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 - } - } -} - type jsonRPCResponseMatch int const ( diff --git a/internal/mcp/client_test.go b/internal/mcp/client_test.go index 109013f3..25df1af7 100644 --- a/internal/mcp/client_test.go +++ b/internal/mcp/client_test.go @@ -5,15 +5,11 @@ import ( "context" "database/sql" "encoding/json" - "errors" - "fmt" "io" "log" "net/http" "net/http/httptest" "net/url" - "os" - "path/filepath" "strings" "sync" "sync/atomic" @@ -611,1507 +607,6 @@ func TestMergeRequestMeta(t *testing.T) { } } -// invokeStreamTestServer builds the initialize/notifications.initialized -// scaffolding shared by the InvokeStream tests below, dispatching tools/call -// to callHandler. -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.Fatalf("decode request: %v", err) - } - 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.Fatalf("unexpected method %q", req.Method) - } - })) -} - -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) { - var resumeRequests int - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodGet { - resumeRequests++ - 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 resumeRequests != 1 { - t.Fatalf("resume request count = %d, want 1", resumeRequests) - } - 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) - } -} - -// 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) - } -} - -// 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) { - 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": - sessionID := "sid-1" - if len(sessions) > 0 { - sessionID = "sid-2" - } - sessions = append(sessions, sessionID) - 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": - toolsCallCount++ - w.Header().Set("Content-Type", "text/event-stream") - flusher := w.(http.Flusher) - if toolsCallCount == 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) - } - if toolsCallCount != 2 { - t.Fatalf("tools/call count = %d, want 2", toolsCallCount) - } - if len(sessions) != 2 { - t.Fatalf("initialize sessions = %#v, want two sessions", sessions) - } - 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) { - var initializeCount int - 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": - initializeCount++ - 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++ - 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 toolsCallCount != 1 { - t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed)", toolsCallCount) - } - if initializeCount != 1 { - t.Fatalf("initialize count = %d, want 1 (no reinitialize attempt)", initializeCount) - } - 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)) - } -} - -// syncFakeStreamSink is fakeStreamSink's mutex-protected counterpart. It's -// needed wherever a test can have both relaySSEToolCall's own read loop and -// routeStandaloneEvent deliver to the same sink concurrently — the plain -// fakeStreamSink above assumes single-goroutine delivery and would race. -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: - } -} - -// 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) - } -} - -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)) - } -} - -// 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. -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(upstream, 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(upstream, 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(upstream, s3) -} - -// 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()) - } -} - -// writeFakeStdioServer writes an executable bash script implementing the -// initialize/notifications.initialized handshake and dispatching tools/call -// to script (a bash fragment appended verbatim, given $line as the raw -// incoming JSON and able to compute its id via `id=$(echo "$line" | grep -o -// '"id":[0-9]*' | head -1 | cut -d: -f2)`). -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 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) - } -} - -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 TestCallTimeoutGuardHeaderTimeoutDisarmedAfterHeadersArrive(t *testing.T) { - g := newCallTimeoutGuard(context.Background()) - defer g.stop() - g.armHeaderTimeout(10 * time.Millisecond) - g.disarmHeaderTimeout() - time.Sleep(30 * time.Millisecond) - if reason := g.reason(); reason != "" { - t.Fatalf("expected a disarmed header 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) - } -} - -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 TestMissingSessionRPCErrorDetection(t *testing.T) { if !isMissingSessionRPCError(json.RawMessage(`{"code":-32000,"message":"No session ID provided for non-initialization request"}`)) { t.Fatal("expected missing session error to be detected") diff --git a/internal/mcp/http_stream.go b/internal/mcp/http_stream.go new file mode 100644 index 00000000..b264054c --- /dev/null +++ b/internal/mcp/http_stream.go @@ -0,0 +1,473 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync" + "time" +) + +// 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 { + if reason := guard.reason(); reason != "" { + return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) + } + return streamCallOutcome{}, 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() + _, _ = io.Copy(io.Discard, resp.Body) + return streamCallOutcome{missingSession: true, sessionID: h.sessionID}, nil + } + + if !strings.Contains(strings.ToLower(h.contentType), "text/event-stream") { + defer resp.Body.Close() + bodyBytes, err := io.ReadAll(resp.Body) + if err != nil { + if reason := guard.reason(); reason != "" { + return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) + } + return streamCallOutcome{}, 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) +} + +// 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) *postStreamPump { + p := &postStreamPump{msgs: make(chan postStreamMsg), current: resp, done: make(chan struct{})} + go p.run(c, guard, upstream) + 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). 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() + defer p.mu.Unlock() + if p.stopped { + return false + } + p.current = resp + 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: + } +} + +func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstream) { + defer close(p.msgs) + reader := newSSEEventReader(p.current.Body) + lastEventID := "" + retryDelay := time.Duration(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 reason := guard.reason(); reason != "" { + p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout)}) + 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 + } + if err := waitForSSEReconnect(guard.ctx, retryDelay); err != nil { + if reason := guard.reason(); reason != "" { + p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while waiting to resume: %w", upstream.Name, reason, ErrStreamTimeout)}) + return + } + p.send(postStreamMsg{err: err}) + return + } + resumed, err := c.resumeSSEStream(guard.ctx, upstream, lastEventID) + if err != nil { + if reason := guard.reason(); reason != "" { + p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while resuming: %w", upstream.Name, reason, ErrStreamTimeout)}) + return + } + p.send(postStreamMsg{err: err}) + return + } + if !p.setCurrent(resumed) { + _ = resumed.Body.Close() + return + } + reader = newSSEEventReader(resumed.Body) + resumedFrom = lastEventID + continue + } + guard.resetIdle() + 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) (streamCallOutcome, error) { + expectedID := json.RawMessage([]byte("1")) + 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) + } + + pump := newPostStreamPump(c, guard, upstream, resp) + 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 streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + case msg, ok := <-pump.msgs: + if !ok { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q: stream ended unexpectedly", upstream.Name) + } + if msg.err != nil { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, msg.err + } + payload := msg.data + switch classifyRPCMessage(payload, expectedID) { + case rpcMessageTerminalResponse: + var rpcResp rpcResponse + if err := json.Unmarshal(payload, &rpcResp); err != nil { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, statusCode) + if progressCh != nil { + // See terminalSettleWindow: give a notification already in + // flight on the standalone stream a brief, bounded chance + // to arrive before finalizing. + settle := time.NewTimer(terminalSettleWindow) + settleLoop: + for { + select { + case evt, ok := <-progressCh: + if !ok { + break settleLoop + } + if err := deliver(evt); err != nil { + settle.Stop() + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + if !settle.Stop() { + <-settle.C + } + settle.Reset(terminalSettleWindow) + case <-settle.C: + break settleLoop + } + } + } + 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 streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + case rpcMessageNotification: + if err := deliver(StreamEvent{Data: payload}); err != nil { + return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + } + default: + // Unrecognized payload shape (e.g. a response to some other id). + // Not ours to interpret; ignore and keep reading. + } + } + } +} + +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 + if rewritten, wireToken, original, ok := c.rewriteProgressToken(merged); ok { + merged = rewritten + effectiveSink = newCallSink(sink, wireToken, original) + progressCh = make(chan StreamEvent, standaloneWaiterEventBuffer) + standalone := c.acquireStandaloneStream(upstream) + standalone.registerWaiter(wireToken, progressWaiter{events: progressCh}) + defer func() { + c.releaseStandaloneStream(upstream, standalone) + // The POST and standalone connections can finish out of order. + // Keep the unique-token waiter briefly so an in-flight progress + // message is not dropped after the terminal response arrives. + time.AfterFunc(standaloneWaiterGracePeriod, func() { + standalone.unregisterWaiter(wireToken) + }) + }() + } + + 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 + } + 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..12a5c1fd --- /dev/null +++ b/internal/mcp/http_stream_test.go @@ -0,0 +1,572 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// invokeStreamTestServer builds the initialize/notifications.initialized +// scaffolding shared by the InvokeStream tests below, dispatching tools/call +// to callHandler. +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.Fatalf("decode request: %v", err) + } + 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.Fatalf("unexpected method %q", req.Method) + } + })) +} + +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) { + var resumeRequests int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet { + resumeRequests++ + 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 resumeRequests != 1 { + t.Fatalf("resume request count = %d, want 1", resumeRequests) + } + 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) + } +} + +// 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) { + 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": + sessionID := "sid-1" + if len(sessions) > 0 { + sessionID = "sid-2" + } + sessions = append(sessions, sessionID) + 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": + toolsCallCount++ + w.Header().Set("Content-Type", "text/event-stream") + flusher := w.(http.Flusher) + if toolsCallCount == 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) + } + if toolsCallCount != 2 { + t.Fatalf("tools/call count = %d, want 2", toolsCallCount) + } + if len(sessions) != 2 { + t.Fatalf("initialize sessions = %#v, want two sessions", sessions) + } + 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) { + var initializeCount int + 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": + initializeCount++ + 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++ + 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 toolsCallCount != 1 { + t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed)", toolsCallCount) + } + if initializeCount != 1 { + t.Fatalf("initialize count = %d, want 1 (no reinitialize attempt)", initializeCount) + } + 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)) + } +} diff --git a/internal/mcp/sse_reader.go b/internal/mcp/sse_reader.go new file mode 100644 index 00000000..b4b3292f --- /dev/null +++ b/internal/mcp/sse_reader.go @@ -0,0 +1,148 @@ +package mcp + +import ( + "bufio" + "encoding/json" + "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 + dataLines []string + eventID string + retry time.Duration + hasData bool + hasID bool + hasRetry bool +} + +type sseWireEvent struct { + Data []byte + ID string + Retry time.Duration + HasData bool + HasID bool + HasRetry bool +} + +func newSSEEventReader(r io.Reader) *sseEventReader { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 1024*1024), 4*1024*1024) + return &sseEventReader{scanner: scanner} +} + +// 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() { + line := r.scanner.Text() + if line == "" { + if !r.hasData && !r.hasID && !r.hasRetry { + continue + } + return r.takeEvent(), nil + } + 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 { + 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 + 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..ee67ef42 --- /dev/null +++ b/internal/mcp/sse_reader_test.go @@ -0,0 +1,59 @@ +package mcp + +import ( + "context" + "encoding/json" + "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) + } +} diff --git a/internal/mcp/standalone_stream.go b/internal/mcp/standalone_stream.go new file mode 100644 index 00000000..e04123af --- /dev/null +++ b/internal/mcp/standalone_stream.go @@ -0,0 +1,314 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync" + "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 +} + +// 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. +// One ref-counted stream serves all calls sharing an upstream session. It is +// not resumable; a later acquire opens a fresh connection. +// +// standaloneWaiterGracePeriod lets an in-flight progress message arrive after +// the POST terminal response without keeping its unique-token waiter forever. +const standaloneWaiterGracePeriod = 2 * time.Second + +// terminalSettleWindow briefly drains progress that races the terminal across +// the independent POST and standalone connections. Each arrival resets it so +// a trailing burst is drained completely. +const terminalSettleWindow = 25 * time.Millisecond + +type standaloneStream struct { + 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 +} + +// 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 { + c.standaloneMu.Lock() + s := c.standaloneStreams[upstream.Name] + if s == nil { + s = &standaloneStream{waiters: make(map[string]progressWaiter)} + c.standaloneStreams[upstream.Name] = s + } + c.standaloneMu.Unlock() + + s.mu.Lock() + s.refCount++ + start := s.refCount == 1 && !s.unsupported + if start { + streamCtx, cancel := context.WithCancel(context.Background()) + s.cancel = cancel + s.done = make(chan struct{}) + go c.runStandaloneStream(streamCtx, upstream, s) + } + s.mu.Unlock() + return s +} + +// releaseStandaloneStream drops one reference acquired via +// acquireStandaloneStream. Once the last reference is gone, it cancels the +// reader goroutine, waits for it to fully exit, and evicts the entry so a +// future acquire opens a fresh connection (picking up, e.g., a session that +// was reinitialized in the meantime). +func (c *Client) releaseStandaloneStream(upstream Upstream, s *standaloneStream) { + 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 + } + s.mu.Unlock() + if cancel != nil { + cancel() + <-done + } + if last { + c.standaloneMu.Lock() + if c.standaloneStreams[upstream.Name] == s { + delete(c.standaloneStreams, upstream.Name) + } + c.standaloneMu.Unlock() + } +} + +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) (*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 := 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 standalone stream returned 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 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) { + defer close(s.done) + resp, err := c.openStandaloneGET(ctx, upstream) + if err != nil { + c.debugf("standalone stream unavailable server=%s err=%v", upstream.Name, err) + s.mu.Lock() + s.unsupported = true + s.mu.Unlock() + return + } + defer resp.Body.Close() + reader := newSSEEventReader(resp.Body) + for { + evt, err := reader.NextEvent() + if err != nil { + return + } + if !evt.HasData { + continue + } + c.routeStandaloneEvent(s, evt.Data) + } +} + +// 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 + } + 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}: + 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. + } +} + +// 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) +} diff --git a/internal/mcp/standalone_stream_test.go b/internal/mcp/standalone_stream_test.go new file mode 100644 index 00000000..7fd42ec6 --- /dev/null +++ b/internal/mcp/standalone_stream_test.go @@ -0,0 +1,596 @@ +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. It's +// needed wherever a test can have both relaySSEToolCall's own read loop and +// routeStandaloneEvent deliver to the same sink concurrently — the plain +// fakeStreamSink above assumes single-goroutine delivery and would race. +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: + } +} + +// 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) + } +} + +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)) + } +} + +// 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. +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(upstream, 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(upstream, 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(upstream, s3) +} + +// 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()) + } +} diff --git a/internal/mcp/stdio_stream.go b/internal/mcp/stdio_stream.go new file mode 100644 index 00000000..f6a25ca5 --- /dev/null +++ b/internal/mcp/stdio_stream.go @@ -0,0 +1,159 @@ +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 := readRPC(reader, rpcIDMessage(initID)); err != nil { + if reason := guard.reason(); reason != "" { + return InvokeResult{}, fmt.Errorf("upstream %q: %s during stdio initialize: %w", upstream.Name, reason, ErrStreamTimeout) + } + 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) +} + +// 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) (InvokeResult, error) { + expectedID := rpcIDMessage(callID) + started := false + ensureStarted := func() { + if !started { + started = true + sink.StreamStarted() + } + } + for { + line, err := reader.ReadBytes('\n') + if err != nil { + if reason := guard.reason(); reason != "" { + return InvokeResult{}, fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout) + } + if stderr.Len() > 0 { + return InvokeResult{}, fmt.Errorf("stdio upstream error: %s", strings.TrimSpace(stderr.String())) + } + return InvokeResult{}, err + } + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + guard.resetIdle() + + 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..d3968e14 --- /dev/null +++ b/internal/mcp/stdio_stream_test.go @@ -0,0 +1,253 @@ +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 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..2c680a5b --- /dev/null +++ b/internal/mcp/stream.go @@ -0,0 +1,134 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "time" +) + +// 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 +} + +// 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, including events routed over + // the shared standalone HTTP stream. 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 +} + +// 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..98a67216 --- /dev/null +++ b/internal/mcp/stream_timeout.go @@ -0,0 +1,149 @@ +package mcp + +import ( + "context" + "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 +} 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) + } +} From 0a3ae30bcd9eaa824fc366be3a2dfdc669817847 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 14:25:53 -0400 Subject: [PATCH 06/18] fix: harden upstream stream recovery Bound upstream message memory, add reconnect backoff, and keep standalone progress streams aligned with MCP session renewal. Retry transient standalone failures and expose buffer drops to stream observers. --- atryum.example.toml | 32 ++-- internal/config/config.go | 4 + internal/config/config_test.go | 3 + internal/mcp/client.go | 38 +++- internal/mcp/http_stream.go | 104 ++++++++--- internal/mcp/http_stream_test.go | 46 +++++ internal/mcp/sse_reader.go | 24 ++- internal/mcp/sse_reader_test.go | 37 ++++ internal/mcp/standalone_stream.go | 162 ++++++++++++----- internal/mcp/standalone_stream_test.go | 236 ++++++++++++++++++++++++- internal/mcp/stdio_stream.go | 8 +- internal/mcp/stdio_stream_test.go | 22 +++ internal/mcp/stream.go | 30 ++++ pkg/atryum/atryum.go | 7 +- 14 files changed, 646 insertions(+), 107 deletions(-) diff --git a/atryum.example.toml b/atryum.example.toml index ea1a0181..4ffa8fd4 100644 --- a/atryum.example.toml +++ b/atryum.example.toml @@ -49,30 +49,26 @@ connection_timeout_seconds = 5 [defaults] request_timeout_seconds = 30 -# Live SSE relay for tools/call: when an upstream MCP server answers a -# tools/call with a Server-Sent Events stream (per the Streamable HTTP -# transport, MCP spec 2025-11-25), Atryum relays intermediate messages -# (progress, logging, other notifications) to the agent as they arrive, -# instead of buffering the whole response and returning only the terminal -# result. This only ever activates when the agent's own request also sends -# Accept: text/event-stream and the upstream responds with one; non-streaming -# tool calls are completely unaffected. stream_relay_enabled is a kill-switch -# to disable the relay globally without a rollback. +# 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 waiting for the upstream's response headers on a streaming -# tools/call (the connect phase, before Atryum knows whether the response -# will stream). 0 falls back to request_timeout_seconds above. +# 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 between successive relayed events once a stream has -# started; resets on every event. Does not bound the call's total duration. +# Bounds the gap between upstream messages while reading the response. +# Resets on every message. Does not bound the call's total duration. stream_idle_timeout_seconds = 60 -# Bounds the whole call once a stream has started. 0 = unlimited. +# 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 -# audit queue can still drop events if storage cannot keep up, and reports -# those drops in invocation.stream_completed. +# 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. diff --git a/internal/config/config.go b/internal/config/config.go index 9b2689ea..e77ec3a3 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -134,6 +134,9 @@ type DefaultsConfig struct { // 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 @@ -183,6 +186,7 @@ func Load(path string) (Config, error) { StreamRelayEnabled: true, StreamIdleTimeoutSeconds: 60, StreamMaxDurationSeconds: 600, + StreamMaxMessageBytes: 4 * 1024 * 1024, StreamAuditMaxEvents: 100, StreamAuditMaxEventBytes: 4096, }, diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 7becaa1b..0ff72fdd 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -87,6 +87,9 @@ func TestLoadMissingConfigUsesDefaultsAndEnv(t *testing.T) { 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) } diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 5aee47df..9c734a18 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -254,13 +254,11 @@ type Client struct { sessions map[string]string sessionProtocols map[string]string - // standaloneStreams holds, per upstream name, the shared "standalone" - // SSE GET connection used to receive server-initiated messages that - // aren't tied to any specific request — notably progress notifications - // from servers (e.g. the reference Python SDK) that don't attribute - // them to the request that triggered them. See standaloneStream. + // 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[string]*standaloneStream + standaloneStreams map[standaloneStreamKey]*standaloneStream } type InvokeResult struct { @@ -333,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), standaloneStreams: make(map[string]*standaloneStream)} + 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) { @@ -1667,8 +1665,12 @@ func writeRPC(w interface{ Write([]byte) (int, error) }, id int64, method string // 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 } @@ -1687,6 +1689,26 @@ func readRPC(reader *bufio.Reader, expectedID json.RawMessage) (rpcResponse, err } } +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 { diff --git a/internal/mcp/http_stream.go b/internal/mcp/http_stream.go index b264054c..1c451bf4 100644 --- a/internal/mcp/http_stream.go +++ b/internal/mcp/http_stream.go @@ -5,9 +5,11 @@ import ( "encoding/json" "fmt" "io" + "math/rand/v2" "net/http" "strings" "sync" + "sync/atomic" "time" ) @@ -62,7 +64,7 @@ func (c *Client) doHTTPToolCallStream(ctx context.Context, upstream Upstream, bo if !strings.Contains(strings.ToLower(h.contentType), "text/event-stream") { defer resp.Body.Close() - bodyBytes, err := io.ReadAll(resp.Body) + bodyBytes, err := readAllLimited(resp.Body, opts.maxMessageBytes()) if err != nil { if reason := guard.reason(); reason != "" { return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) @@ -79,7 +81,26 @@ func (c *Client) doHTTPToolCallStream(ctx context.Context, upstream Upstream, bo // 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) + 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 @@ -108,9 +129,9 @@ type postStreamPump struct { done chan struct{} } -func newPostStreamPump(c *Client, guard *callTimeoutGuard, upstream Upstream, resp *http.Response) *postStreamPump { +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) + go p.run(c, guard, upstream, maxMessageBytes) return p } @@ -158,11 +179,42 @@ func (p *postStreamPump) send(msg postStreamMsg) { } } -func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstream) { +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) - reader := newSSEEventReader(p.current.Body) + reader := newSSEEventReaderWithLimit(p.current.Body, maxMessageBytes) 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 @@ -191,7 +243,9 @@ func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstre p.send(postStreamMsg{err: fmt.Errorf("upstream %q closed the stream without a JSON-RPC response or resumable event id", upstream.Name)}) return } - if err := waitForSSEReconnect(guard.ctx, retryDelay); err != nil { + delay := sseReconnectDelay(retryDelay, reconnectAttempt) + reconnectAttempt++ + if err := waitForSSEReconnect(guard.ctx, delay); err != nil { if reason := guard.reason(); reason != "" { p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while waiting to resume: %w", upstream.Name, reason, ErrStreamTimeout)}) return @@ -212,11 +266,12 @@ func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstre _ = resumed.Body.Close() return } - reader = newSSEEventReader(resumed.Body) + reader = newSSEEventReaderWithLimit(resumed.Body, maxMessageBytes) resumedFrom = lastEventID continue } guard.resetIdle() + reconnectAttempt = 0 if evt.HasRetry { retryDelay = evt.Retry } @@ -241,7 +296,7 @@ func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstre // 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) (streamCallOutcome, error) { +func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progressCh <-chan StreamEvent, guard *callTimeoutGuard, upstream Upstream, sessionID string, maxMessageBytes int) (streamCallOutcome, error) { expectedID := json.RawMessage([]byte("1")) statusCode := resp.StatusCode relayed := 0 @@ -259,7 +314,7 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress return sink.Event(evt) } - pump := newPostStreamPump(c, guard, upstream, resp) + pump := newPostStreamPump(c, guard, upstream, resp, maxMessageBytes) defer pump.stop() for { @@ -424,20 +479,23 @@ func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool s // from either the POST response or the standalone stream. effectiveSink := sink var progressCh chan StreamEvent - if rewritten, wireToken, original, ok := c.rewriteProgressToken(merged); ok { + 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.acquireStandaloneStream(upstream) - standalone.registerWaiter(wireToken, progressWaiter{events: progressCh}) + standalone = c.acquireStandaloneStreamWithLimit(upstream, opts.maxMessageBytes()) + standalone.registerWaiter(wireToken, progressWaiter{events: progressCh, dropped: &standaloneDropped}) defer func() { - c.releaseStandaloneStream(upstream, standalone) - // The POST and standalone connections can finish out of order. - // Keep the unique-token waiter briefly so an in-flight progress - // message is not dropped after the terminal response arrives. - time.AfterFunc(standaloneWaiterGracePeriod, func() { - standalone.unregisterWaiter(wireToken) - }) + current := standalone + current.unregisterWaiter(wireToken) + c.releaseStandaloneStream(current) + if statsSink, ok := effectiveSink.(StreamStatsSink); ok { + statsSink.StreamStats(StreamStats{StandaloneEventsDropped: standaloneDropped.Load()}) + } }() } @@ -460,6 +518,12 @@ func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool s }); 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}) + } outcome, err = c.doHTTPToolCallStream(ctx, upstream, body, effectiveSink, progressCh, opts) if err != nil { return InvokeResult{}, err diff --git a/internal/mcp/http_stream_test.go b/internal/mcp/http_stream_test.go index 12a5c1fd..eae30a11 100644 --- a/internal/mcp/http_stream_test.go +++ b/internal/mcp/http_stream_test.go @@ -145,6 +145,52 @@ func TestInvokeStreamResumesAfterUpstreamClosesSSEBeforeTerminalResponse(t *test } } +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 diff --git a/internal/mcp/sse_reader.go b/internal/mcp/sse_reader.go index b4b3292f..b6ed8854 100644 --- a/internal/mcp/sse_reader.go +++ b/internal/mcp/sse_reader.go @@ -3,6 +3,7 @@ package mcp import ( "bufio" "encoding/json" + "errors" "fmt" "io" "strconv" @@ -17,6 +18,8 @@ import ( // 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 @@ -35,9 +38,17 @@ type sseWireEvent struct { } 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) - scanner.Buffer(make([]byte, 1024*1024), 4*1024*1024) - return &sseEventReader{scanner: scanner} + 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 @@ -47,10 +58,15 @@ func (r *sseEventReader) NextEvent() (sseWireEvent, error) { 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 } @@ -85,6 +101,9 @@ func (r *sseEventReader) NextEvent() (sseWireEvent, error) { } } 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 { @@ -108,6 +127,7 @@ func (r *sseEventReader) takeEvent() sseWireEvent { r.hasData = false r.hasID = false r.hasRetry = false + r.eventSize = 0 return evt } diff --git a/internal/mcp/sse_reader_test.go b/internal/mcp/sse_reader_test.go index ee67ef42..953fea22 100644 --- a/internal/mcp/sse_reader_test.go +++ b/internal/mcp/sse_reader_test.go @@ -3,6 +3,7 @@ package mcp import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "strings" @@ -57,3 +58,39 @@ func TestSSEEventReaderSaturatesOversizedRetryWithoutOverflow(t *testing.T) { 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) + } +} diff --git a/internal/mcp/standalone_stream.go b/internal/mcp/standalone_stream.go index e04123af..20a31646 100644 --- a/internal/mcp/standalone_stream.go +++ b/internal/mcp/standalone_stream.go @@ -3,11 +3,13 @@ package mcp import ( "context" "encoding/json" + "errors" "fmt" "io" "net/http" "strings" "sync" + "sync/atomic" "time" ) @@ -15,7 +17,8 @@ import ( // owns the sink. This prevents shared-reader writes from racing the call's // terminal response. type progressWaiter struct { - events chan StreamEvent + events chan StreamEvent + dropped *atomic.Int64 } // A full waiter buffer drops that call's event rather than blocking the one @@ -25,19 +28,13 @@ 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. -// One ref-counted stream serves all calls sharing an upstream session. It is -// not resumable; a later acquire opens a fresh connection. -// -// standaloneWaiterGracePeriod lets an in-flight progress message arrive after -// the POST terminal response without keeping its unique-token waiter forever. -const standaloneWaiterGracePeriod = 2 * time.Second - // terminalSettleWindow briefly drains progress that races the terminal across // the independent POST and standalone connections. Each arrival resets it so // a trailing burst is drained completely. const terminalSettleWindow = 25 * time.Millisecond type standaloneStream struct { + key standaloneStreamKey mu sync.Mutex refCount int cancel context.CancelFunc @@ -49,39 +46,72 @@ type standaloneStream struct { // 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 +} + +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[upstream.Name] + s := c.standaloneStreams[key] if s == nil { - s = &standaloneStream{waiters: make(map[string]progressWaiter)} - c.standaloneStreams[upstream.Name] = s + s = &standaloneStream{ + key: key, + waiters: make(map[string]progressWaiter), + maxBytes: maxMessageBytes, + } + c.standaloneStreams[key] = s } - c.standaloneMu.Unlock() - 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 = make(chan struct{}) - go c.runStandaloneStream(streamCtx, upstream, s) + 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. Once the last reference is gone, it cancels the -// reader goroutine, waits for it to fully exit, and evicts the entry so a -// future acquire opens a fresh connection (picking up, e.g., a session that -// was reinitialized in the meantime). -func (c *Client) releaseStandaloneStream(upstream Upstream, s *standaloneStream) { +// 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 @@ -92,19 +122,16 @@ func (c *Client) releaseStandaloneStream(upstream Upstream, s *standaloneStream) 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 } - if last { - c.standaloneMu.Lock() - if c.standaloneStreams[upstream.Name] == s { - delete(c.standaloneStreams, upstream.Name) - } - c.standaloneMu.Unlock() - } } func (s *standaloneStream) registerWaiter(token string, w progressWaiter) { @@ -122,17 +149,17 @@ func (s *standaloneStream) unregisterWaiter(token string) { // 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) (*http.Response, error) { +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 := c.getSessionProtocol(upstream.Name); protocol != "" { + if protocol != "" { req.Header.Set("MCP-Protocol-Version", protocol) } - if sessionID := c.getSession(upstream.Name); sessionID != "" { + if sessionID != "" { req.Header.Set("Mcp-Session-Id", sessionID) } applyAuthHeaders(req, upstream) @@ -143,36 +170,64 @@ func (c *Client) openStandaloneGET(ctx context.Context, upstream Upstream) (*htt if resp.StatusCode >= http.StatusBadRequest { defer resp.Body.Close() body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024)) - return nil, fmt.Errorf("upstream %q standalone stream returned HTTP %d: %s", upstream.Name, resp.StatusCode, extractErrorDetail(body)) + 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, fmt.Errorf("upstream %q standalone stream returned content type %q, want text/event-stream", upstream.Name, resp.Header.Get("Content-Type")) + 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) { - defer close(s.done) - resp, err := c.openStandaloneGET(ctx, upstream) - if err != nil { - c.debugf("standalone stream unavailable server=%s err=%v", upstream.Name, err) - s.mu.Lock() - s.unsupported = true - s.mu.Unlock() - return - } - defer resp.Body.Close() - reader := newSSEEventReader(resp.Body) +func (c *Client) runStandaloneStream(ctx context.Context, upstream Upstream, s *standaloneStream, done chan<- struct{}) { + defer close(done) + attempt := 0 for { - evt, err := reader.NextEvent() + resp, err := c.openStandaloneGET(ctx, upstream, s.key.sessionID, s.key.protocol) if err != nil { - return - } - if !evt.HasData { + 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 } - c.routeStandaloneEvent(s, evt.Data) + + 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++ } } @@ -221,6 +276,9 @@ func (c *Client) routeStandaloneEvent(s *standaloneStream, payload []byte) { // 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) + } } } @@ -312,3 +370,9 @@ func (s *callSink) Event(evt StreamEvent) error { } 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 index 7fd42ec6..f007e4b5 100644 --- a/internal/mcp/standalone_stream_test.go +++ b/internal/mcp/standalone_stream_test.go @@ -143,6 +143,25 @@ func TestRouteStandaloneEventAttributesTokenlessNotificationOnlyWhenUnambiguous( } } +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 @@ -298,6 +317,184 @@ func TestInvokeStreamStandaloneProgressResetsIdleTimeout(t *testing.T) { } } +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 @@ -480,11 +677,11 @@ func TestStandaloneStreamRefcountsSharedConnection(t *testing.T) { t.Fatalf("expected exactly 1 standalone connection while both waiters are active, got %d", got) } - client.releaseStandaloneStream(upstream, s1) + 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(upstream, s2) + client.releaseStandaloneStream(s2) s3 := client.acquireStandaloneStream(upstream) if s3 == s1 { @@ -497,7 +694,40 @@ func TestStandaloneStreamRefcountsSharedConnection(t *testing.T) { if got := atomic.LoadInt32(&connections); got != 2 { t.Fatalf("expected a new connection after full release + reacquire, got %d", got) } - client.releaseStandaloneStream(upstream, s3) + client.releaseStandaloneStream(s3) +} + +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 diff --git a/internal/mcp/stdio_stream.go b/internal/mcp/stdio_stream.go index f6a25ca5..da1cfef1 100644 --- a/internal/mcp/stdio_stream.go +++ b/internal/mcp/stdio_stream.go @@ -74,7 +74,7 @@ func (c *Client) invokeStdioStream(ctx context.Context, upstream Upstream, tool }); err != nil { return InvokeResult{}, err } - if _, err := readRPC(reader, rpcIDMessage(initID)); err != nil { + if _, err := readRPCWithLimit(reader, rpcIDMessage(initID), opts.maxMessageBytes()); err != nil { if reason := guard.reason(); reason != "" { return InvokeResult{}, fmt.Errorf("upstream %q: %s during stdio initialize: %w", upstream.Name, reason, ErrStreamTimeout) } @@ -95,13 +95,13 @@ func (c *Client) invokeStdioStream(ctx context.Context, upstream Upstream, tool } guard.armBodyTimeouts(opts.IdleTimeout, opts.MaxDuration) - return c.relayStdioToolCall(reader, sink, guard, upstream, callID, stderr) + 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) (InvokeResult, error) { +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() { @@ -111,7 +111,7 @@ func (c *Client) relayStdioToolCall(reader *bufio.Reader, sink StreamSink, guard } } for { - line, err := reader.ReadBytes('\n') + line, err := readLineLimited(reader, maxMessageBytes) if err != nil { if reason := guard.reason(); reason != "" { return InvokeResult{}, fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout) diff --git a/internal/mcp/stdio_stream_test.go b/internal/mcp/stdio_stream_test.go index d3968e14..69a03021 100644 --- a/internal/mcp/stdio_stream_test.go +++ b/internal/mcp/stdio_stream_test.go @@ -93,6 +93,28 @@ func TestInvokeStreamStdioTerminalOnlyResponseNeverTouchesSink(t *testing.T) { } } +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"+ diff --git a/internal/mcp/stream.go b/internal/mcp/stream.go index 2c680a5b..5ea2ef35 100644 --- a/internal/mcp/stream.go +++ b/internal/mcp/stream.go @@ -3,10 +3,17 @@ 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 @@ -23,6 +30,19 @@ type StreamEvent struct { 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 @@ -54,6 +74,16 @@ type StreamOptions struct { // 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 diff --git a/pkg/atryum/atryum.go b/pkg/atryum/atryum.go index 4265c275..7c24a357 100644 --- a/pkg/atryum/atryum.go +++ b/pkg/atryum/atryum.go @@ -268,9 +268,10 @@ func runServer(args []string, o options) error { } service.SetStreamOptions( mcp.StreamOptions{ - HeaderTimeout: time.Duration(streamHeaderTimeoutSeconds) * time.Second, - IdleTimeout: time.Duration(cfg.Defaults.StreamIdleTimeoutSeconds) * time.Second, - MaxDuration: time.Duration(cfg.Defaults.StreamMaxDurationSeconds) * time.Second, + HeaderTimeout: time.Duration(streamHeaderTimeoutSeconds) * 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, From 89bfdf38b1149cbf3820fa1aa93846606039dafb Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 14:25:58 -0400 Subject: [PATCH 07/18] fix: make stream auditing bounded and truthful Replace per-invocation audit goroutines with an ordered shared dispatcher, classify quiet downstream cancellation, record terminal delivery separately, and avoid claiming success before durable persistence. --- CHANGELOG.md | 12 +- docs/architecture.md | 43 +++-- internal/api/handlers.go | 25 ++- internal/api/handlers_test.go | 11 ++ internal/api/sse_relay_test.go | 33 ++++ internal/invocation/service.go | 25 ++- internal/invocation/stream_execution.go | 34 ++-- internal/invocation/stream_execution_test.go | 110 +++++++++++ internal/invocation/stream_sink.go | 174 +++++++++++------ .../invocation/stream_sink_internal_test.go | 180 ++++++++++++++++++ 10 files changed, 553 insertions(+), 94 deletions(-) create mode 100644 internal/invocation/stream_sink_internal_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 384cb1a2..2efed8d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,11 +19,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 are audited into `invocation_events` - (`invocation.stream_event` / `invocation.stream_completed`). New + 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_audit_max_events`, + `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 @@ -34,6 +36,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `_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. ## [0.2.0] - 2026-07-14 diff --git a/docs/architecture.md b/docs/architecture.md index 1db2a7cf..0fa2e052 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -246,6 +246,7 @@ Streaming is opt-in and keeps the old behavior as its fallback: | **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. | @@ -320,6 +321,7 @@ sequenceDiagram 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 ``` @@ -354,22 +356,24 @@ cross-delivery. |---|---| | `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 background queue. | +| `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`. -#### Time limits +#### 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 three upstream limits: +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 response headers or stdio session initialization. | +| 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 between upstream messages. It resets after every message. | | 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 @@ -381,7 +385,8 @@ other transport failures. #### 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 actually streams. + 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 @@ -393,25 +398,37 @@ other transport failures. - **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. -- **The standalone stream does not resume.** If that shared GET disconnects, calls can - still receive their terminal responses on their POST connections, but may miss - standalone progress until a later group of calls opens a new GET. + 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 handled by the relay - is offered to a bounded background audit queue. Slow or failed writes are counted in - the completion audit row; they do not delay the live event. +- **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. + 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 diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 686b93ae..36f3aef5 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -43,6 +43,7 @@ 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) @@ -1371,7 +1372,8 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server h.debugf("mcp tools.call error after stream started server=%s tool=%s err=%v", server, params.Name, err) errBody, _ := json.Marshal(map[string]any{"code": -32000, "message": err.Error()}) terminal, _ := json.Marshal(jsonRPCResponse{JSONRPC: "2.0", ID: req.ID, Error: errBody}) - _ = sink.finishStream(terminal) + deliveryErr := sink.finishStream(terminal) + h.recordStreamDelivery(r.Context(), resp.InvocationID, deliveryErr) return } h.writeRPCError(w, req.ID, -32000, err.Error()) @@ -1394,7 +1396,8 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server // goroutine — mandatory before returning from the handler. body, _ := json.Marshal(result) terminal, _ := json.Marshal(jsonRPCResponse{JSONRPC: "2.0", ID: req.ID, Result: body}) - _ = sink.finishStream(terminal) + deliveryErr := sink.finishStream(terminal) + h.recordStreamDelivery(r.Context(), resp.InvocationID, deliveryErr) return } h.writeRPCResult(w, req.ID, result) @@ -4244,3 +4247,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 1ef49793..9e308011 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -47,6 +47,10 @@ type stubService struct { recordReq *invocation.ExternalExecutionUpdate recordCtx context.Context + streamDeliveryInvocationID string + streamDeliveryStatus string + streamDeliveryMessage string + createSessionReq *invocation.CreateSessionRequest createSessionAgentID string @@ -75,6 +79,13 @@ func (s *stubService) InvokeStreaming(ctx context.Context, req invocation.Create } 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 } diff --git a/internal/api/sse_relay_test.go b/internal/api/sse_relay_test.go index 7f54b9b6..c5406eff 100644 --- a/internal/api/sse_relay_test.go +++ b/internal/api/sse_relay_test.go @@ -199,6 +199,39 @@ func TestSSERelaySinkHeartbeatFailureSurfacesOnNextEvent(t *testing.T) { } } +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) { diff --git a/internal/invocation/service.go b/internal/invocation/service.go index 8794ed88..1b1b50ab 100644 --- a/internal/invocation/service.go +++ b/internal/invocation/service.go @@ -180,9 +180,9 @@ type upstreamClient interface { // 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 audit sink's -// bounded queue remains a final backpressure guard even with MaxEvents zero, -// and reports any queue drops in invocation.stream_completed. +// 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 @@ -269,6 +269,25 @@ func (s *Service) SetStreamOptions(opts mcp.StreamOptions, auditLimits StreamAud s.streamAuditLimits = auditLimits } +// 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. diff --git a/internal/invocation/stream_execution.go b/internal/invocation/stream_execution.go index 3e79babc..a3f423ad 100644 --- a/internal/invocation/stream_execution.go +++ b/internal/invocation/stream_execution.go @@ -10,11 +10,10 @@ import ( "github.com/validmind/atryum/internal/mcp" ) -// finishExecutionStreaming is finishExecution's live-relay path: it wraps -// the caller's sink in an auditing decorator (so every relayed event and -// the call's outcome are recorded as invocation_events rows regardless of -// whether the downstream write later fails) and calls InvokeStream with -// s.streamOptions instead of the fixed s.defaultTimeout. +// 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.streamAuditLimits) result, err := s.client.InvokeStream(ctx, upstream, req.Tool, req.Input, req.RequestID, req.Meta, audited, s.streamOptions) @@ -23,14 +22,15 @@ func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, if err != nil { inv.Status = StatusFailed - reason, message := classifyStreamError(audited, err) + reason, message := classifyStreamError(ctx, audited, err) inv.Error = mustJSON(map[string]any{"message": message}) - audited.finish(completed, "failed") persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout) defer cancelPersist() if updateErr := s.invocations.UpdateResult(persistCtx, inv); updateErr != nil { - return InvocationResponse{}, fmt.Errorf("persist streaming invocation failure: %w", updateErr) + 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 @@ -47,7 +47,6 @@ func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, if result.Failed { inv.Status = StatusFailed inv.Error = result.Body - audited.finish(completed, "failed") 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)}), @@ -56,7 +55,6 @@ func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, } else { inv.Status = StatusSucceeded inv.Response = result.Body - audited.finish(completed, "succeeded") 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)}), @@ -66,7 +64,13 @@ func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, persistCtx, cancelPersist := context.WithTimeout(context.WithoutCancel(ctx), terminalPersistenceTimeout) defer cancelPersist() if err := s.invocations.UpdateResult(persistCtx, inv); err != nil { - return InvocationResponse{}, err + audited.finish(completed, "persistence_failed") + return s.toResponse(inv), err + } + if result.Failed { + audited.finish(completed, "failed") + } else { + audited.finish(completed, "succeeded") } _ = s.events.Create(persistCtx, terminalEvent) return s.toResponse(inv), nil @@ -79,15 +83,21 @@ func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, // 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(audited *auditingSink, err error) (reason string, message string) { +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) { + return "stream_aborted_downstream", 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 index 6641dc32..c9c56452 100644 --- a/internal/invocation/stream_execution_test.go +++ b/internal/invocation/stream_execution_test.go @@ -36,6 +36,17 @@ type blockingStreamEventRepo 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) }) @@ -391,6 +402,105 @@ func TestInvokeStreamingSinkAbortPersistsFailureAfterRequestContextCancellation( t.Fatal("expected persisted invocation.failed event with reason stream_aborted_downstream") } +func TestInvokeStreamingQuietRequestCancellationIsDownstreamAbort(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 { + if evt.Type == "invocation.failed" && jsonContains(evt.Data, "stream_aborted_downstream") { + return + } + } + t.Fatal("quiet request cancellation was not audited as stream_aborted_downstream") +} + +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) { diff --git a/internal/invocation/stream_sink.go b/internal/invocation/stream_sink.go index 0fc0a76b..4736d64d 100644 --- a/internal/invocation/stream_sink.go +++ b/internal/invocation/stream_sink.go @@ -10,11 +10,57 @@ import ( ) const ( - streamAuditQueueCapacity = 128 + 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)} + perWorkerCapacity := 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 { + return int(d.nextShard.Add(1)-1) % 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 @@ -23,26 +69,26 @@ const ( // 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 a bounded background queue rather than on the -// relay's hot path. A stalled audit repository must neither delay an event -// reaching the agent nor defeat the stream's idle/max-duration bounds. Each -// write and the final queue drain are time-bounded; failures and queue drops -// are summarized in invocation.stream_completed. +// 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 - auditCtx context.Context - cancelAudit context.CancelFunc - auditQueue chan Event - auditDone chan struct{} - closeOnce sync.Once - persisted atomic.Int64 - failed atomic.Int64 - dropped atomic.Int64 + 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 @@ -59,12 +105,8 @@ func newAuditingSink(inner mcp.StreamSink, events eventRepo, invocationID string requestID: requestID, upstreamName: upstreamName, limits: limits, - } - if events != nil { - a.auditCtx, a.cancelAudit = context.WithCancel(context.Background()) - a.auditQueue = make(chan Event, streamAuditQueueCapacity) - a.auditDone = make(chan struct{}) - go a.runAuditWriter() + auditShard: sharedStreamAuditDispatcher.assignShard(), + drained: make(chan struct{}), } return a } @@ -88,8 +130,12 @@ func (a *auditingSink) Event(evt mcp.StreamEvent) error { return nil } +func (a *auditingSink) StreamStats(stats mcp.StreamStats) { + a.standaloneDropped.Add(stats.StandaloneEventsDropped) +} + func (a *auditingSink) recordEvent(evt mcp.StreamEvent) { - if a.auditQueue == nil { + if a.events == nil { return } if a.limits.MaxEvents > 0 && a.seq > a.limits.MaxEvents { @@ -120,62 +166,68 @@ func (a *auditingSink) recordEvent(evt mcp.StreamEvent) { Payload: mustJSON(payload), CreatedAt: time.Now().UTC(), } - select { - case a.auditQueue <- record: - default: + 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) runAuditWriter() { - defer close(a.auditDone) - for evt := range a.auditQueue { - writeCtx, cancel := context.WithTimeout(a.auditCtx, streamAuditWriteTimeout) - err := a.events.Create(writeCtx, evt) - cancel() - if err != nil { - a.failed.Add(1) - continue - } +func (a *auditingSink) completeAuditWrite(err error) { + if err != nil { + a.failed.Add(1) + } else { a.persisted.Add(1) } + a.completePending() } -func (a *auditingSink) stopAuditWriter() { - if a.auditQueue == nil { - return +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) }) } - a.closeOnce.Do(func() { close(a.auditQueue) }) +} + +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.auditDone: - a.cancelAudit() + case <-a.drained: + return true case <-timer.C: - // Cancel the in-flight write and make every queued write fail fast. - // Do not wait indefinitely if an eventRepo violates context - // cancellation; terminal invocation persistence must remain bounded. - a.cancelAudit() - select { - case <-a.auditDone: - case <-time.After(streamAuditWriteTimeout): - } + return false } } // finish records the invocation.stream_completed totals row. terminal is -// "succeeded", "failed", or "aborted". +// "succeeded", "failed", or "persistence_failed". func (a *auditingSink) finish(completed time.Time, terminal string) { if a.events == nil { return } - a.stopAuditWriter() + 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(), - "terminal": terminal, + "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() 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") + } + }) +} From 4e11271949d495e55a80269cd87efb30ac4da78b Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 16:22:36 -0400 Subject: [PATCH 08/18] fix: address stream relay review findings Correctness: - guard the stdio stderr boundedBuffer with a mutex: os/exec's copy goroutine can still be writing when error paths read Len/String before cmd.Wait runs - flag server-to-client requests routed via the standalone stream as ServerRequest, matching the POST-stream classification, so audit doesn't mislabel sampling/elicitation requests as notifications - normalize CR/CRLF in downstream SSE framing: a raw CR inside upstream JSON would split the data frame for a compliant parser Safety and audit truthfulness: - substitute derived stream bounds when SetStreamOptions was never called, so an embedder wiring NewService directly never runs an unbounded relay - reset the idle timeout on SSE keepalive/comment lines so a busy-but- quiet upstream heartbeating through a long tool run isn't cut off - send a static message in the terminal SSE error frame; internal error text (SQL/driver detail) stays in the server log - audit a bare request-context cancel as stream_canceled rather than stream_aborted_downstream: shutdown and a quiet disconnect are indistinguishable without a failed downstream write - rename the failure audit's events_relayed to events_total to match stream_completed; cap the session-expired body drain at 64KB Readability: - extract callTimeoutGuard.timeoutErr, drainTrailingProgress, and a shared toolCallEnvelopeID; move the header-timeout fallback into config.EffectiveStreamHeaderTimeoutSeconds; document the stream_completed durability-before-delivery latency bound Tests: - fix unsynchronized test-server counters (latent -race flakes) - pin drainTrailingProgress settle-window semantics deterministically - add end-to-end MaxDuration, resume-failure, StreamStats delivery, and keepalive-liveness coverage; poll for pending_approval instead of sleeping in the approval-gate test --- atryum.example.toml | 5 +- docs/architecture.md | 2 +- internal/api/handlers.go | 7 +- internal/api/sse_relay.go | 14 +- internal/api/sse_relay_test.go | 29 +- internal/config/config.go | 12 + internal/config/config_test.go | 15 + internal/invocation/service.go | 38 ++- internal/invocation/stream_execution.go | 11 +- internal/invocation/stream_execution_test.go | 53 ++- internal/invocation/stream_sink.go | 11 + internal/mcp/client.go | 37 ++- internal/mcp/http_stream.go | 136 +++++--- internal/mcp/http_stream_test.go | 326 ++++++++++++++++++- internal/mcp/sse_reader.go | 9 + internal/mcp/sse_reader_test.go | 22 ++ internal/mcp/standalone_stream.go | 7 +- internal/mcp/standalone_stream_test.go | 74 +++++ internal/mcp/stdio_stream.go | 8 +- internal/mcp/stdio_stream_test.go | 2 + internal/mcp/stream.go | 9 +- internal/mcp/stream_timeout.go | 18 + pkg/atryum/atryum.go | 6 +- 23 files changed, 737 insertions(+), 114 deletions(-) diff --git a/atryum.example.toml b/atryum.example.toml index 4ffa8fd4..7b708095 100644 --- a/atryum.example.toml +++ b/atryum.example.toml @@ -57,8 +57,9 @@ 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 between upstream messages while reading the response. -# Resets on every message. Does not bound the call's total duration. +# 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 diff --git a/docs/architecture.md b/docs/architecture.md index 0fa2e052..5058d928 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -371,7 +371,7 @@ 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 between upstream messages. It resets after every message. | +| 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. | diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 36f3aef5..ec63c43b 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -1369,8 +1369,13 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server // 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; 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. h.debugf("mcp tools.call error after stream started server=%s tool=%s err=%v", server, params.Name, err) - errBody, _ := json.Marshal(map[string]any{"code": -32000, "message": err.Error()}) + 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) diff --git a/internal/api/sse_relay.go b/internal/api/sse_relay.go index 983ce668..3db43873 100644 --- a/internal/api/sse_relay.go +++ b/internal/api/sse_relay.go @@ -11,6 +11,18 @@ import ( "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. @@ -20,7 +32,7 @@ func writeSSEEvent(w io.Writer, flusher http.Flusher, event string, data []byte) return err } } - for _, line := range bytes.Split(data, []byte("\n")) { + for _, line := range splitSSELines(data) { if _, err := fmt.Fprintf(w, "data: %s\n", line); err != nil { return err } diff --git a/internal/api/sse_relay_test.go b/internal/api/sse_relay_test.go index c5406eff..92760bc9 100644 --- a/internal/api/sse_relay_test.go +++ b/internal/api/sse_relay_test.go @@ -255,8 +255,13 @@ func TestMCPToolsCallErrorAfterStreamStartedWritesTerminalErrorFrame(t *testing. 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, "persisting result failed") { - t.Fatalf("expected a terminal JSON-RPC error frame, 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) } } @@ -574,3 +579,23 @@ func TestMCPToolsCallEndToEndRelaysLiveBeforeTerminalResponseExists(t *testing.T 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/config/config.go b/internal/config/config.go index e77ec3a3..68f87a84 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -149,6 +149,18 @@ type DefaultsConfig struct { 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 { Name string `toml:"name"` Mode string `toml:"mode"` diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 0ff72fdd..0e92a243 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -169,3 +169,18 @@ recent_chat_messages_limit = 42 t.Fatalf("RecentChatMessagesLimit = %d", cfg.ManagedAgents[0].RecentChatMessagesLimit) } } + +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) + } +} diff --git a/internal/invocation/service.go b/internal/invocation/service.go index 1b1b50ab..502bfe61 100644 --- a/internal/invocation/service.go +++ b/internal/invocation/service.go @@ -218,9 +218,12 @@ type Service struct { pendingApprovals map[string]chan approvalDecision // streamOptions and streamAuditLimits govern InvokeStreaming's execution - // once a sink is present (see finishExecutionStreaming). The zero value - // of each disables its bounds; SetStreamOptions installs real values. + // 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 @@ -266,9 +269,40 @@ func (s *Service) SetInvocationSummarizer(client SummaryClient) { // 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 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 +) + +// 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, + } +} + // 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. diff --git a/internal/invocation/stream_execution.go b/internal/invocation/stream_execution.go index a3f423ad..e467e3d1 100644 --- a/internal/invocation/stream_execution.go +++ b/internal/invocation/stream_execution.go @@ -16,7 +16,7 @@ import ( // 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.streamAuditLimits) - result, err := s.client.InvokeStream(ctx, upstream, req.Tool, req.Input, req.RequestID, req.Meta, audited, s.streamOptions) + 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 @@ -38,7 +38,7 @@ func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, // 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_relayed": audited.seq}), + Payload: mustJSON(map[string]any{"reason": reason, "message": message, "events_total": audited.seq}), CreatedAt: completed, }) return s.toResponse(inv), nil @@ -88,7 +88,12 @@ func classifyStreamError(ctx context.Context, audited *auditingSink, err error) return "stream_aborted_downstream", audited.downstreamErr.Error() } if errors.Is(err, context.Canceled) || errors.Is(ctx.Err(), context.Canceled) { - return "stream_aborted_downstream", err.Error() + // 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() diff --git a/internal/invocation/stream_execution_test.go b/internal/invocation/stream_execution_test.go index c9c56452..668a7a99 100644 --- a/internal/invocation/stream_execution_test.go +++ b/internal/invocation/stream_execution_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "sync" + "sync/atomic" "testing" "time" @@ -402,7 +403,7 @@ func TestInvokeStreamingSinkAbortPersistsFailureAfterRequestContextCancellation( t.Fatal("expected persisted invocation.failed event with reason stream_aborted_downstream") } -func TestInvokeStreamingQuietRequestCancellationIsDownstreamAbort(t *testing.T) { +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") @@ -440,11 +441,15 @@ func TestInvokeStreamingQuietRequestCancellationIsDownstreamAbort(t *testing.T) t.Fatal(err) } for _, evt := range events.Items { - if evt.Type == "invocation.failed" && jsonContains(evt.Data, "stream_aborted_downstream") { + // 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_aborted_downstream") + t.Fatal("quiet request cancellation was not audited as stream_canceled") } func TestInvokeStreamingDoesNotAuditSuccessBeforeTerminalPersistence(t *testing.T) { @@ -552,9 +557,10 @@ func TestInvokeStreamingIdleTimeoutMarksFailedAsStreamTimeout(t *testing.T) { } func TestInvokeStreamingMidStreamSessionRetryRefusalMarksFailedWithDistinctReason(t *testing.T) { - var toolsCallCount int + // 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++ + 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}}`) @@ -575,8 +581,8 @@ func TestInvokeStreamingMidStreamSessionRetryRefusalMarksFailedWithDistinctReaso if resp.Status != invocation.StatusFailed { t.Fatalf("expected failed status, got %s", resp.Status) } - if toolsCallCount != 1 { - t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed mid-stream)", toolsCallCount) + 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{}) @@ -617,16 +623,35 @@ func TestInvokeStreamingApprovalGateDoesNotTouchSinkBeforeApproval(t *testing.T) sink := &recordingSink{} go func() { - time.Sleep(50 * time.Millisecond) - if sink.touched() { - t.Errorf("sink touched before approval — approval gating must precede any relay") + // 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 + deadline := time.Now().Add(10 * time.Second) + for pendingID == "" && time.Now().Before(deadline) { + list, err := service.List(context.Background(), invocation.InvocationListFilter{Limit: 10}) + if err == nil { + for _, item := range list.Items { + if item.Status == invocation.StatusPendingApproval { + pendingID = item.InvocationID + break + } + } + } + if pendingID == "" { + time.Sleep(5 * time.Millisecond) + } } - list, err := service.List(context.Background(), invocation.InvocationListFilter{Limit: 10}) - if err != nil || len(list.Items) == 0 { - t.Errorf("expected a pending invocation to approve") + if pendingID == "" { + t.Error("timed out waiting for a pending-approval invocation") return } - if err := service.Approve(context.Background(), list.Items[0].InvocationID, ""); err != nil { + 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) } }() diff --git a/internal/invocation/stream_sink.go b/internal/invocation/stream_sink.go index 4736d64d..3f83c6aa 100644 --- a/internal/invocation/stream_sink.go +++ b/internal/invocation/stream_sink.go @@ -215,6 +215,17 @@ func (a *auditingSink) waitForAuditWrites() bool { // 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 diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 9c734a18..938af063 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -633,10 +633,16 @@ func (c *Client) TestConnection(ctx context.Context, upstream Upstream) Connecti return result } +// 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 id "1" — Atryum's own request id, -// if any, travels only in the caller-facing InvocationResponse, not on the -// wire to the upstream. +// 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)) } @@ -650,7 +656,7 @@ func marshalToolCallEnvelopeWithMeta(tool string, input map[string]any, meta map if meta != nil { params["_meta"] = meta } - return json.Marshal(Envelope{JSONRPC: "2.0", ID: json.RawMessage([]byte("1")), Method: "tools/call", Params: mustRawJSON(params)}) + return json.Marshal(Envelope{JSONRPC: "2.0", ID: toolCallEnvelopeID, Method: "tools/call", Params: mustRawJSON(params)}) } // toolCallResultFromRPCResponse maps an already-decoded tools/call JSON-RPC @@ -673,7 +679,7 @@ func toolCallResultFromRPCResponse(rpcResp rpcResponse, statusCode int) (InvokeR // 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, json.RawMessage([]byte("1"))) + rpcResp, err := decodeRPCResponse(result, toolCallEnvelopeID) if err != nil { return InvokeResult{}, false, err } @@ -1096,7 +1102,13 @@ const stdioStderrCap = 64 * 1024 // 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 } @@ -1106,6 +1118,8 @@ func newBoundedBuffer(limit int) *boundedBuffer { } 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 @@ -1120,8 +1134,17 @@ func (b *boundedBuffer) Write(p []byte) (int, error) { return len(p), nil } -func (b *boundedBuffer) Len() int { return b.buf.Len() } -func (b *boundedBuffer) String() string { return b.buf.String() } +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 == "" { diff --git a/internal/mcp/http_stream.go b/internal/mcp/http_stream.go index 1c451bf4..fbc13b25 100644 --- a/internal/mcp/http_stream.go +++ b/internal/mcp/http_stream.go @@ -13,6 +13,32 @@ import ( "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 @@ -39,10 +65,7 @@ func (c *Client) doHTTPToolCallStream(ctx context.Context, upstream Upstream, bo h, err := c.doHTTPEnvelopeHeaders(guard.ctx, upstream, body, DefaultMCPProtocolVersion, true) guard.disarmSetupTimeout() if err != nil { - if reason := guard.reason(); reason != "" { - return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) - } - return streamCallOutcome{}, err + return streamCallOutcome{}, guard.timeoutErr(upstream.Name, "", err) } resp := h.resp @@ -58,7 +81,11 @@ func (c *Client) doHTTPToolCallStream(ctx context.Context, upstream Upstream, bo if h.sessionExpired { defer resp.Body.Close() - _, _ = io.Copy(io.Discard, resp.Body) + // 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 } @@ -66,10 +93,7 @@ func (c *Client) doHTTPToolCallStream(ctx context.Context, upstream Upstream, bo defer resp.Body.Close() bodyBytes, err := readAllLimited(resp.Body, opts.maxMessageBytes()) if err != nil { - if reason := guard.reason(); reason != "" { - return streamCallOutcome{}, fmt.Errorf("upstream %q: %s: %w", upstream.Name, reason, ErrStreamTimeout) - } - return streamCallOutcome{}, err + 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) @@ -211,7 +235,11 @@ func sseReconnectDelay(serverDelay time.Duration, attempt int) time.Duration { 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 @@ -231,8 +259,8 @@ func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstre if stopped { return } - if reason := guard.reason(); reason != "" { - p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout)}) + if timeoutErr := guard.timeoutErr(upstream.Name, "", nil); timeoutErr != nil { + p.send(postStreamMsg{err: timeoutErr}) return } if err != io.EOF { @@ -246,20 +274,12 @@ func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstre delay := sseReconnectDelay(retryDelay, reconnectAttempt) reconnectAttempt++ if err := waitForSSEReconnect(guard.ctx, delay); err != nil { - if reason := guard.reason(); reason != "" { - p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while waiting to resume: %w", upstream.Name, reason, ErrStreamTimeout)}) - return - } - p.send(postStreamMsg{err: err}) + 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 { - if reason := guard.reason(); reason != "" { - p.send(postStreamMsg{err: fmt.Errorf("upstream %q %s while resuming: %w", upstream.Name, reason, ErrStreamTimeout)}) - return - } - p.send(postStreamMsg{err: err}) + p.send(postStreamMsg{err: guard.timeoutErr(upstream.Name, "while resuming", err)}) return } if !p.setCurrent(resumed) { @@ -267,10 +287,10 @@ func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstre return } reader = newSSEEventReaderWithLimit(resumed.Body, maxMessageBytes) + reader.onActivity = guard.resetIdle resumedFrom = lastEventID continue } - guard.resetIdle() reconnectAttempt = 0 if evt.HasRetry { retryDelay = evt.Retry @@ -297,7 +317,7 @@ func (p *postStreamPump) run(c *Client, guard *callTimeoutGuard, upstream Upstre // 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 := json.RawMessage([]byte("1")) + expectedID := toolCallEnvelopeID statusCode := resp.StatusCode relayed := 0 started := false @@ -313,6 +333,12 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress 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() @@ -328,46 +354,26 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress continue } if err := deliver(evt); err != nil { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + return fail(err) } case msg, ok := <-pump.msgs: if !ok { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, fmt.Errorf("upstream %q: stream ended unexpectedly", upstream.Name) + return fail(fmt.Errorf("upstream %q: stream ended unexpectedly", upstream.Name)) } if msg.err != nil { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, msg.err + 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 streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + return fail(err) } invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, statusCode) if progressCh != nil { - // See terminalSettleWindow: give a notification already in - // flight on the standalone stream a brief, bounded chance - // to arrive before finalizing. - settle := time.NewTimer(terminalSettleWindow) - settleLoop: - for { - select { - case evt, ok := <-progressCh: - if !ok { - break settleLoop - } - if err := deliver(evt); err != nil { - settle.Stop() - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err - } - if !settle.Stop() { - <-settle.C - } - settle.Reset(terminalSettleWindow) - case <-settle.C: - break settleLoop - } + if err := drainTrailingProgress(progressCh, deliver, terminalSettleWindow); err != nil { + return fail(err) } } if !(missingSession && relayed == 0) { @@ -376,11 +382,11 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress return streamCallOutcome{invoke: invoke, missingSession: missingSession, eventsRelayed: relayed, sessionID: sessionID}, nil case rpcMessageServerRequest: if err := deliver(StreamEvent{Data: payload, ServerRequest: true}); err != nil { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + return fail(err) } case rpcMessageNotification: if err := deliver(StreamEvent{Data: payload}); err != nil { - return streamCallOutcome{eventsRelayed: relayed, sessionID: sessionID}, err + return fail(err) } default: // Unrecognized payload shape (e.g. a response to some other id). @@ -390,6 +396,34 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress } } +// drainTrailingProgress gives a standalone-stream notification already in +// flight a brief, bounded chance to arrive before the terminal response is +// finalized (the caller passes terminalSettleWindow as window). Each arrival +// resets the window so a trailing burst is drained completely. It returns +// once the window elapses with no new arrival, progressCh closes, or deliver +// fails. +func drainTrailingProgress(progressCh <-chan StreamEvent, deliver func(StreamEvent) error, window time.Duration) error { + settle := time.NewTimer(window) + defer settle.Stop() + for { + select { + case evt, ok := <-progressCh: + if !ok { + return nil + } + if err := deliver(evt); err != nil { + return err + } + if !settle.Stop() { + <-settle.C + } + settle.Reset(window) + case <-settle.C: + return nil + } + } +} + func waitForSSEReconnect(ctx context.Context, delay time.Duration) error { if delay <= 0 { return nil diff --git a/internal/mcp/http_stream_test.go b/internal/mcp/http_stream_test.go index eae30a11..d8f16404 100644 --- a/internal/mcp/http_stream_test.go +++ b/internal/mcp/http_stream_test.go @@ -8,6 +8,8 @@ import ( "net/http" "net/http/httptest" "strings" + "sync" + "sync/atomic" "testing" "time" ) @@ -82,10 +84,12 @@ func TestInvokeStreamRelaysEventsBeforeTerminalResponseExists(t *testing.T) { } func TestInvokeStreamResumesAfterUpstreamClosesSSEBeforeTerminalResponse(t *testing.T) { - var resumeRequests int + // 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++ + 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) } @@ -134,8 +138,8 @@ func TestInvokeStreamResumesAfterUpstreamClosesSSEBeforeTerminalResponse(t *test if err != nil { t.Fatalf("InvokeStream returned error: %v", err) } - if resumeRequests != 1 { - t.Fatalf("resume request count = %d, want 1", resumeRequests) + 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) @@ -466,6 +470,9 @@ func TestInvokeStreamMapsTerminalRPCErrorAfterRelayedEvents(t *testing.T) { } 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) { @@ -475,20 +482,25 @@ func TestInvokeStreamRetriesOnceWhenMissingSessionBeforeAnyEventRelayed(t *testi } 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 toolsCallCount == 1 { + if call == 1 { writeTestSSEEventFlush(w, flusher, `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"No session ID provided for non-initialization request"}}`) return } @@ -510,11 +522,14 @@ func TestInvokeStreamRetriesOnceWhenMissingSessionBeforeAnyEventRelayed(t *testi if err != nil { t.Fatalf("InvokeStream returned error: %v", err) } - if toolsCallCount != 2 { - t.Fatalf("tools/call count = %d, want 2", toolsCallCount) + stateMu.Lock() + gotCalls, gotSessions := toolsCallCount, len(sessions) + stateMu.Unlock() + if gotCalls != 2 { + t.Fatalf("tools/call count = %d, want 2", gotCalls) } - if len(sessions) != 2 { - t.Fatalf("initialize sessions = %#v, want two sessions", sessions) + 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) @@ -525,8 +540,9 @@ func TestInvokeStreamRetriesOnceWhenMissingSessionBeforeAnyEventRelayed(t *testi } func TestInvokeStreamRefusesRetryAfterEventsAlreadyRelayed(t *testing.T) { - var initializeCount int - var toolsCallCount int + // 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 { @@ -534,13 +550,13 @@ func TestInvokeStreamRefusesRetryAfterEventsAlreadyRelayed(t *testing.T) { } switch req.Method { case "initialize": - initializeCount++ + 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++ + 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}}`) @@ -565,11 +581,11 @@ func TestInvokeStreamRefusesRetryAfterEventsAlreadyRelayed(t *testing.T) { if !errors.Is(err, ErrStreamSessionRetryRefused) { t.Fatalf("expected errors.Is(err, ErrStreamSessionRetryRefused) to hold, got %v", err) } - if toolsCallCount != 1 { - t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed)", toolsCallCount) + if got := toolsCallCount.Load(); got != 1 { + t.Fatalf("tools/call count = %d, want 1 (no retry once events were relayed)", got) } - if initializeCount != 1 { - t.Fatalf("initialize count = %d, want 1 (no reinitialize attempt)", initializeCount) + 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)) @@ -616,3 +632,279 @@ func TestInvokeStreamIdleTimeoutAbortsRead(t *testing.T) { 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(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(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(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) + } +} + +func TestDrainTrailingProgressResetsWindowPerArrival(t *testing.T) { + // Three events spaced 300ms apart against a 500ms window. Each gap is + // under the window (200ms margin), but the cumulative spacing is not: + // without the per-arrival reset the single 500ms timer fires between the + // second and third event and the drain returns having delivered only 2. + const window = 500 * time.Millisecond + const gap = 300 * time.Millisecond + + progressCh := make(chan StreamEvent, 1) + go func() { + for range 3 { + progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)} + time.Sleep(gap) + } + }() + + delivered := 0 + deliver := func(StreamEvent) error { + delivered++ + return nil + } + if err := drainTrailingProgress(progressCh, deliver, window); err != nil { + t.Fatalf("drainTrailingProgress returned error: %v", err) + } + if delivered != 3 { + t.Fatalf("delivered %d events, want all 3 (window must reset on each arrival)", delivered) + } +} + +// 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. +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 index b6ed8854..9329b3e7 100644 --- a/internal/mcp/sse_reader.go +++ b/internal/mcp/sse_reader.go @@ -26,6 +26,12 @@ type sseEventReader struct { 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 { @@ -55,6 +61,9 @@ func newSSEEventReaderWithLimit(r io.Reader, maxBytes int) *sseEventReader { // 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 { diff --git a/internal/mcp/sse_reader_test.go b/internal/mcp/sse_reader_test.go index 953fea22..fa20b50f 100644 --- a/internal/mcp/sse_reader_test.go +++ b/internal/mcp/sse_reader_test.go @@ -94,3 +94,25 @@ func TestSSEEventReaderDoesNotAccumulateCommentBytesAcrossEvents(t *testing.T) { 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 index 20a31646..ec3c19e7 100644 --- a/internal/mcp/standalone_stream.go +++ b/internal/mcp/standalone_stream.go @@ -248,6 +248,11 @@ func (c *Client) routeStandaloneEvent(s *standaloneStream, payload []byte) { 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() @@ -271,7 +276,7 @@ func (c *Client) routeStandaloneEvent(s *standaloneStream, payload []byte) { // itself (matching on its own wireToken), so the raw payload is sent // through unmodified here. select { - case waiter.events <- StreamEvent{Data: payload}: + 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 diff --git a/internal/mcp/standalone_stream_test.go b/internal/mcp/standalone_stream_test.go index f007e4b5..ed0a5989 100644 --- a/internal/mcp/standalone_stream_test.go +++ b/internal/mcp/standalone_stream_test.go @@ -824,3 +824,77 @@ func TestInvokeStreamStandaloneStreamWrongContentTypeDoesNotFailCall(t *testing. 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_stream.go b/internal/mcp/stdio_stream.go index da1cfef1..0b0ab444 100644 --- a/internal/mcp/stdio_stream.go +++ b/internal/mcp/stdio_stream.go @@ -75,8 +75,8 @@ func (c *Client) invokeStdioStream(ctx context.Context, upstream Upstream, tool return InvokeResult{}, err } if _, err := readRPCWithLimit(reader, rpcIDMessage(initID), opts.maxMessageBytes()); err != nil { - if reason := guard.reason(); reason != "" { - return InvokeResult{}, fmt.Errorf("upstream %q: %s during stdio initialize: %w", upstream.Name, reason, ErrStreamTimeout) + 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())) @@ -113,8 +113,8 @@ func (c *Client) relayStdioToolCall(reader *bufio.Reader, sink StreamSink, guard for { line, err := readLineLimited(reader, maxMessageBytes) if err != nil { - if reason := guard.reason(); reason != "" { - return InvokeResult{}, fmt.Errorf("upstream %q %s: %w", upstream.Name, reason, ErrStreamTimeout) + 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())) diff --git a/internal/mcp/stdio_stream_test.go b/internal/mcp/stdio_stream_test.go index 69a03021..0a7e0999 100644 --- a/internal/mcp/stdio_stream_test.go +++ b/internal/mcp/stdio_stream_test.go @@ -1,3 +1,5 @@ +//go:build !windows + package mcp import ( diff --git a/internal/mcp/stream.go b/internal/mcp/stream.go index 5ea2ef35..7fd145e1 100644 --- a/internal/mcp/stream.go +++ b/internal/mcp/stream.go @@ -67,9 +67,12 @@ type StreamOptions struct { // 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, including events routed over - // the shared standalone HTTP stream. For a plain HTTP JSON response it - // bounds the complete body read. Zero disables the check. + // 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. diff --git a/internal/mcp/stream_timeout.go b/internal/mcp/stream_timeout.go index 98a67216..3eef1fec 100644 --- a/internal/mcp/stream_timeout.go +++ b/internal/mcp/stream_timeout.go @@ -2,6 +2,7 @@ package mcp import ( "context" + "fmt" "sync" "sync/atomic" "time" @@ -147,3 +148,20 @@ func (g *callTimeoutGuard) reason() string { 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/pkg/atryum/atryum.go b/pkg/atryum/atryum.go index 7c24a357..006cb255 100644 --- a/pkg/atryum/atryum.go +++ b/pkg/atryum/atryum.go @@ -262,13 +262,9 @@ func runServer(args []string, o options) error { service.SetInvocationSummarizer(&summaryAdapter{client: backendClient}) } service.SetSessionStore(store.NewExternalSessionRepoWithDialect(db, dialect)) - streamHeaderTimeoutSeconds := cfg.Defaults.StreamHeaderTimeoutSeconds - if streamHeaderTimeoutSeconds <= 0 { - streamHeaderTimeoutSeconds = cfg.Defaults.RequestTimeoutSeconds - } service.SetStreamOptions( mcp.StreamOptions{ - HeaderTimeout: time.Duration(streamHeaderTimeoutSeconds) * time.Second, + 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, From eb8f7380a1f480b8a4ce4f93db4199e664a04dec Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Thu, 23 Jul 2026 16:54:39 -0400 Subject: [PATCH 09/18] finalize buffered calls after disconnect; harden relay lifecycle - run buffered finalization writes on a bounded detached context so a downstream disconnect can't leave the row stuck executing (matches the streaming path; regression test added) - emit the mcp.tools.call trace event on the error-after-stream-started path (finalize_failed: true) - defer sseRelaySink.stopHeartbeat in the handler so the heartbeat goroutine can't outlive a panicking handler; count delivered (not attempted) frames in eventCount - shard index modulo in uint64; queue capacity floored at 1; wrap both persistence errors consistently - shared test upstream helpers fail cleanly (t.Errorf + HTTP error) instead of hanging on t.Fatalf from a handler goroutine - document relay behavior guarantees and the buffered fix in the changelog; note SetStreamRelayEnabled is startup-wiring only --- CHANGELOG.md | 14 ++++ docs/architecture.md | 7 +- internal/api/handlers.go | 20 +++++- internal/api/sse_relay.go | 21 +++++- internal/invocation/service.go | 41 ++++++++--- internal/invocation/stream_execution.go | 4 +- internal/invocation/stream_execution_test.go | 73 +++++++++++++++++++- internal/invocation/stream_sink.go | 8 ++- internal/mcp/http_stream_test.go | 22 ++++-- internal/mcp/standalone_stream_test.go | 12 ++-- internal/mcp/stdio_stream.go | 4 +- 11 files changed, 194 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2efed8d9..9aada60e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 diff --git a/docs/architecture.md b/docs/architecture.md index 5058d928..d9a07b2f 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -379,8 +379,11 @@ The downstream connection has a separate per-write deadline. If the downstream c 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, a downstream disconnect, and -other transport failures. +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 diff --git a/internal/api/handlers.go b/internal/api/handlers.go index ec63c43b..d837991b 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -737,7 +737,10 @@ func NewHandler(svc service, serverSvc serverService, policyRegistry *policy.Reg // 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. +// 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 } @@ -1350,6 +1353,10 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server 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 @@ -1369,12 +1376,19 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server // 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; the wire gets a + // 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. - h.debugf("mcp tools.call error after stream started server=%s tool=%s err=%v", server, params.Name, err) + 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) diff --git a/internal/api/sse_relay.go b/internal/api/sse_relay.go index 3db43873..5ff29e07 100644 --- a/internal/api/sse_relay.go +++ b/internal/api/sse_relay.go @@ -161,22 +161,37 @@ func (s *sseRelaySink) Event(evt mcp.StreamEvent) error { // now rather than waiting for this write to discover it again. return s.writeErr } - s.eventCount++ 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.stopOnce.Do(func() { close(s.heartbeatStop) }) - <-s.heartbeatDone + s.stopHeartbeat() s.mu.Lock() defer s.mu.Unlock() if s.writeErr != nil { diff --git a/internal/invocation/service.go b/internal/invocation/service.go index 502bfe61..a1e140cf 100644 --- a/internal/invocation/service.go +++ b/internal/invocation/service.go @@ -275,12 +275,15 @@ func (s *Service) SetStreamOptions(opts mcp.StreamOptions, auditLimits StreamAud // 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 in +// (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 + fallbackStreamIdleTimeout = 60 * time.Second + fallbackStreamMaxDuration = 10 * time.Minute + fallbackStreamAuditMaxEvents = 100 + fallbackStreamAuditMaxEvtBytes = 4096 ) // effectiveStreamOptions returns the configured stream bounds, or — when @@ -303,6 +306,20 @@ func (s *Service) effectiveStreamOptions() mcp.StreamOptions { } } +// 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. @@ -1220,17 +1237,23 @@ func (s *Service) finishExecutionBuffered(ctx context.Context, inv Invocation, u 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, @@ -1238,14 +1261,14 @@ func (s *Service) finishExecutionBuffered(ctx context.Context, inv Invocation, u } 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) } return s.toResponse(inv), nil } diff --git a/internal/invocation/stream_execution.go b/internal/invocation/stream_execution.go index e467e3d1..8384b61c 100644 --- a/internal/invocation/stream_execution.go +++ b/internal/invocation/stream_execution.go @@ -15,7 +15,7 @@ import ( // 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.streamAuditLimits) + 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 @@ -65,7 +65,7 @@ func (s *Service) finishExecutionStreaming(ctx context.Context, inv Invocation, defer cancelPersist() if err := s.invocations.UpdateResult(persistCtx, inv); err != nil { audited.finish(completed, "persistence_failed") - return s.toResponse(inv), err + return s.toResponse(inv), fmt.Errorf("persist streaming invocation result: %w", err) } if result.Failed { audited.finish(completed, "failed") diff --git a/internal/invocation/stream_execution_test.go b/internal/invocation/stream_execution_test.go index 668a7a99..0933e4b8 100644 --- a/internal/invocation/stream_execution_test.go +++ b/internal/invocation/stream_execution_test.go @@ -97,7 +97,13 @@ func sseToolCallUpstream(t *testing.T, callHandler func(w http.ResponseWriter, r 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.Fatalf("decode request: %v", err) + // 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": @@ -110,7 +116,8 @@ func sseToolCallUpstream(t *testing.T, callHandler func(w http.ResponseWriter, r case "tools/call": callHandler(w, r, body) default: - t.Fatalf("unexpected method %q", body["method"]) + t.Errorf("unexpected method %q", body["method"]) + http.Error(w, "unexpected method", http.StatusInternalServerError) } })) } @@ -705,3 +712,65 @@ func TestInvokeStreamingNilSinkMatchesInvoke(t *testing.T) { 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 index 3f83c6aa..ae227bb2 100644 --- a/internal/invocation/stream_sink.go +++ b/internal/invocation/stream_sink.go @@ -29,7 +29,9 @@ type streamAuditDispatcher struct { func newStreamAuditDispatcher() *streamAuditDispatcher { d := &streamAuditDispatcher{queues: make([]chan streamAuditWrite, streamAuditWorkerCount)} - perWorkerCapacity := streamAuditQueueCapacity / 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]) @@ -38,7 +40,9 @@ func newStreamAuditDispatcher() *streamAuditDispatcher { } func (d *streamAuditDispatcher) assignShard() int { - return int(d.nextShard.Add(1)-1) % len(d.queues) + // 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 { diff --git a/internal/mcp/http_stream_test.go b/internal/mcp/http_stream_test.go index d8f16404..a4cbdfe4 100644 --- a/internal/mcp/http_stream_test.go +++ b/internal/mcp/http_stream_test.go @@ -16,7 +16,10 @@ import ( // invokeStreamTestServer builds the initialize/notifications.initialized // scaffolding shared by the InvokeStream tests below, dispatching tools/call -// to callHandler. +// 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) { @@ -31,7 +34,9 @@ func invokeStreamTestServer(t *testing.T, sessionID string, callHandler func(w h } var req Envelope if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - t.Fatalf("decode request: %v", err) + t.Errorf("decode request: %v", err) + http.Error(w, "bad request body", http.StatusBadRequest) + return } switch req.Method { case "initialize": @@ -42,7 +47,8 @@ func invokeStreamTestServer(t *testing.T, sessionID string, callHandler func(w h case "tools/call": callHandler(w, r, req) default: - t.Fatalf("unexpected method %q", req.Method) + t.Errorf("unexpected method %q", req.Method) + http.Error(w, "unexpected method", http.StatusInternalServerError) } })) } @@ -704,9 +710,17 @@ func TestDrainTrailingProgressResetsWindowPerArrival(t *testing.T) { const gap = 300 * time.Millisecond progressCh := make(chan StreamEvent, 1) + senderStop := make(chan struct{}) + t.Cleanup(func() { close(senderStop) }) go func() { for range 3 { - progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)} + select { + case progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)}: + case <-senderStop: + // The drain exited early (a failure mode under test); don't + // leave this sender blocked on a channel nobody reads. + return + } time.Sleep(gap) } }() diff --git a/internal/mcp/standalone_stream_test.go b/internal/mcp/standalone_stream_test.go index ed0a5989..2e12f4eb 100644 --- a/internal/mcp/standalone_stream_test.go +++ b/internal/mcp/standalone_stream_test.go @@ -13,10 +13,14 @@ import ( "time" ) -// syncFakeStreamSink is fakeStreamSink's mutex-protected counterpart. It's -// needed wherever a test can have both relaySSEToolCall's own read loop and -// routeStandaloneEvent deliver to the same sink concurrently — the plain -// fakeStreamSink above assumes single-goroutine delivery and would race. +// 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 diff --git a/internal/mcp/stdio_stream.go b/internal/mcp/stdio_stream.go index 0b0ab444..82421df1 100644 --- a/internal/mcp/stdio_stream.go +++ b/internal/mcp/stdio_stream.go @@ -121,11 +121,13 @@ func (c *Client) relayStdioToolCall(reader *bufio.Reader, sink StreamSink, guard } 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 } - guard.resetIdle() switch classifyRPCMessage(line, expectedID) { case rpcMessageTerminalResponse: From f62ecc0eb602f1fdabf08bf1c45c4001249b2a7e Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Fri, 24 Jul 2026 12:05:32 -0400 Subject: [PATCH 10/18] fix: close standalone-stream connect race before sending tools/call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit acquireStandaloneStreamWithLimit only spawns the GET-connect goroutine and returns immediately; the caller then sent tools/call right away. If the upstream's first progress notification fired before that GET reached the server, it had nowhere to land and was silently dropped — reproduced as 2/3 progress notifications instead of 3 in TestMCPToolsCallAgainstRealStandaloneStreamServer. Add a ready signal closed after the stream's first connect attempt (success or failure), and wait on it before sending the dependent tools/call, bounded by HeaderTimeout so a hung upstream can't stall the call indefinitely. --- internal/mcp/http_stream.go | 2 ++ internal/mcp/standalone_stream.go | 35 +++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/internal/mcp/http_stream.go b/internal/mcp/http_stream.go index fbc13b25..4a330bc7 100644 --- a/internal/mcp/http_stream.go +++ b/internal/mcp/http_stream.go @@ -523,6 +523,7 @@ func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool s 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) @@ -557,6 +558,7 @@ func (c *Client) invokeHTTPStream(ctx context.Context, upstream Upstream, tool s 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 { diff --git a/internal/mcp/standalone_stream.go b/internal/mcp/standalone_stream.go index ec3c19e7..b985ea06 100644 --- a/internal/mcp/standalone_stream.go +++ b/internal/mcp/standalone_stream.go @@ -47,6 +47,12 @@ type standaloneStream struct { // 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 { @@ -87,6 +93,7 @@ func (c *Client) acquireStandaloneStreamWithLimit(upstream Upstream, maxMessageB key: key, waiters: make(map[string]progressWaiter), maxBytes: maxMessageBytes, + ready: make(chan struct{}), } c.standaloneStreams[key] = s } @@ -134,6 +141,29 @@ func (c *Client) releaseStandaloneStream(s *standaloneStream) { } } +// 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 @@ -186,8 +216,13 @@ func (c *Client) openStandaloneGET(ctx context.Context, upstream Upstream, sessi 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 From 38ce5182b5dd280ffbfb36e57d9ec601822a4471 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Fri, 24 Jul 2026 14:06:38 -0400 Subject: [PATCH 11/18] docs: document SSE relay control flow in pseudocode --- docs/architecture.md | 303 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 303 insertions(+) diff --git a/docs/architecture.md b/docs/architecture.md index d9a07b2f..8c75cdf6 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -362,6 +362,309 @@ cross-delivery. 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 +``` + +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 +``` + +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 From 4864436906a8ba5402d225f07f8fd2261f9b13c0 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Fri, 24 Jul 2026 14:07:00 -0400 Subject: [PATCH 12/18] test: add streaming coverage to the harness integration suite The integration suite had zero coverage of the SSE relay: no streaming MCP target, no timing-aware verification, and fake_agent.py never requested SSE. Adds an everything-streaming target (the official reference "everything" server over stdio) and a fake_agent.py --stream mode that reads the tools/call response incrementally, timestamps each frame, and fails unless progress measurably arrives before the terminal result -- proving live delivery over a real process/network boundary, distinct from the existing Go E2E tests (mcp-everything-test, mcp-standalone-stream-test), which prove the same for the upstream-facing HTTP relay internals against real SDKs. Sharing the server-everything dependency with mcp-everything-test is an accepted, minor redundancy in exchange for not maintaining a second hand-rolled MCP fixture. Also fixes two pre-existing bugs this surfaced, both blocking every case on stock macOS (bash 3.2 + BSD awk), not just the new target: render_atryum_config piping a multi-line TOML block through `awk -v` (swapped for a small Python substitution), and five call sites expanding empty bash arrays under `set -u` (fixed with the portable "${arr[@]+"${arr[@]}"}" idiom). --- integrations/config/mcp-targets.yaml | 36 ++++- integrations/lib/atryum.sh | 29 +++- integrations/lib/harness.sh | 6 +- scripts/fake_agent.py | 190 +++++++++++++++++++++++++-- 4 files changed, 245 insertions(+), 16 deletions(-) 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/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}") From f8d3cc602ece05afc2f8334fdfe0b0d7d7b15740 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Fri, 24 Jul 2026 14:19:29 -0400 Subject: [PATCH 13/18] docs: clarify upstream/downstream terminology in system context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit "Upstream"/"downstream" were used throughout the doc (Runtime ingress, Component boundaries, the SSE relay section) but only ever defined ~200 lines in, under a feature-specific subsection. The System context diagram also had two different nodes both labeled "MCP client" — the downstream agent's connection and Atryum's own upstream-facing client — one of them literally named "Upstream", colliding with what "upstream" means everywhere else in the doc (the external tool server, the diagram's separate "Tools" node). --- docs/architecture.md | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 8c75cdf6..09f9ce23 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -11,12 +11,18 @@ harnesses, an MCP-compatible proxy, and an admin API/UI. All ingress paths share 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] @@ -27,7 +33,7 @@ flowchart LR Service[Invocation service] Rules[Rule evaluation] Store[(SQLite or PostgreSQL)] - Upstream[MCP client] + MCPClient[Upstream MCP client] end Stock --> Bootstrap @@ -42,8 +48,8 @@ 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 `cmd/atryum`, `pkg/atryum`, `pkg/migrations`, From af2efe82d11b3cc3f8cfaef640ccee9145d86598 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Mon, 27 Jul 2026 09:41:44 -0400 Subject: [PATCH 14/18] improved docs with more scenarios --- docs/architecture.md | 163 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 162 insertions(+), 1 deletion(-) diff --git a/docs/architecture.md b/docs/architecture.md index 09f9ce23..ddc086ac 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -145,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 @@ -170,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. @@ -336,6 +387,55 @@ 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 @@ -458,6 +558,38 @@ finishExecutionStreaming(invocation, sink): 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. @@ -551,6 +683,35 @@ Path A reader (its own goroutine, one per call): 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: ``` From 8f22e355d8ce036c51e65dd64b3bb0891ba0f3a4 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Mon, 27 Jul 2026 09:47:02 -0400 Subject: [PATCH 15/18] test: pin the multi-replica approval race and diagram it inline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds TestMultiReplicaApprovalRaceOverwritesApprovedRowAsFailed: two Service instances sharing one DB show that an approval landing on a different process's pendingApprovals map never wakes the original waiter, and that the caller's later context cancellation overwrites the row to failed even though it briefly read approved. This behavior was previously undocumented in code and untested — only docs/architecture.md described it. Also adds ASCII sequence diagrams in the comments above waitForHumanApproval and the standalone-stream token-collision test, so both scenarios are readable from the code itself instead of only from the architecture doc or inferred from synchronization-primitive names. --- internal/invocation/service.go | 29 ++++++ internal/invocation/service_test.go | 120 +++++++++++++++++++++++++ internal/mcp/standalone_stream_test.go | 25 ++++++ 3 files changed, 174 insertions(+) diff --git a/internal/invocation/service.go b/internal/invocation/service.go index ba09e72a..597a0a58 100644 --- a/internal/invocation/service.go +++ b/internal/invocation/service.go @@ -1208,6 +1208,35 @@ func (s *Service) executeNow(ctx context.Context, inv Invocation, upstream mcp.U } // waitForHumanApproval blocks until an operator approves or denies, or the context is cancelled. +// +// 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() diff --git a/internal/invocation/service_test.go b/internal/invocation/service_test.go index 629f25be..fc234500 100644 --- a/internal/invocation/service_test.go +++ b/internal/invocation/service_test.go @@ -197,6 +197,126 @@ 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. +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 + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + 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 + break + } + time.Sleep(10 * time.Millisecond) + } + 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 diff --git a/internal/mcp/standalone_stream_test.go b/internal/mcp/standalone_stream_test.go index 2e12f4eb..0cf0c0de 100644 --- a/internal/mcp/standalone_stream_test.go +++ b/internal/mcp/standalone_stream_test.go @@ -546,6 +546,31 @@ func TestInvokeStreamRestoresCallerProgressTokenOnPOSTResponseStreamToo(t *testi // 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{} From c6abf0c387367700b54f5338ea2b921ac947dab5 Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Mon, 27 Jul 2026 09:59:13 -0400 Subject: [PATCH 16/18] fix: bound terminal stream cleanup Stop continuous trailing progress from extending completed calls, honor cancellation during terminal draining, and close superseded SSE response bodies after reconnection. --- internal/mcp/http_stream.go | 47 ++++++++----- internal/mcp/http_stream_test.go | 113 +++++++++++++++++++++++------- internal/mcp/standalone_stream.go | 4 +- 3 files changed, 120 insertions(+), 44 deletions(-) diff --git a/internal/mcp/http_stream.go b/internal/mcp/http_stream.go index 4a330bc7..cdb41c63 100644 --- a/internal/mcp/http_stream.go +++ b/internal/mcp/http_stream.go @@ -177,17 +177,25 @@ func (p *postStreamPump) stop() { }) } -// setCurrent installs resp as the response the pump is currently reading -// from (after a resume). 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. +// 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() - defer p.mu.Unlock() 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 } @@ -372,8 +380,8 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress } invoke, missingSession := toolCallResultFromRPCResponse(rpcResp, statusCode) if progressCh != nil { - if err := drainTrailingProgress(progressCh, deliver, terminalSettleWindow); err != nil { - return fail(err) + 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) { @@ -396,28 +404,31 @@ func (c *Client) relaySSEToolCall(resp *http.Response, sink StreamSink, progress } } -// drainTrailingProgress gives a standalone-stream notification already in -// flight a brief, bounded chance to arrive before the terminal response is -// finalized (the caller passes terminalSettleWindow as window). Each arrival -// resets the window so a trailing burst is drained completely. It returns -// once the window elapses with no new arrival, progressCh closes, or deliver -// fails. -func drainTrailingProgress(progressCh <-chan StreamEvent, deliver func(StreamEvent) error, window time.Duration) error { +// 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 := deliver(evt); err != nil { + if err := ctx.Err(); err != nil { return err } - if !settle.Stop() { - <-settle.C + if err := deliver(evt); err != nil { + return err } - settle.Reset(window) case <-settle.C: return nil } diff --git a/internal/mcp/http_stream_test.go b/internal/mcp/http_stream_test.go index a4cbdfe4..5d227051 100644 --- a/internal/mcp/http_stream_test.go +++ b/internal/mcp/http_stream_test.go @@ -155,6 +155,41 @@ func TestInvokeStreamResumesAfterUpstreamClosesSSEBeforeTerminalResponse(t *test } } +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) @@ -654,7 +689,7 @@ func TestDrainTrailingProgressDrainsBufferedBurstThenStopsAtClose(t *testing.T) // 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(progressCh, deliver, time.Hour); err != nil { + if err := drainTrailingProgress(context.Background(), progressCh, deliver, time.Hour); err != nil { t.Fatalf("drainTrailingProgress returned error: %v", err) } if len(delivered) != 3 { @@ -676,7 +711,7 @@ func TestDrainTrailingProgressStopsAtDeliverError(t *testing.T) { deliverCalls++ return sinkErr } - if err := drainTrailingProgress(progressCh, deliver, time.Hour); !errors.Is(err, 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 { @@ -691,7 +726,7 @@ func TestDrainTrailingProgressReturnsOnceWindowElapsesWithNoArrival(t *testing.T return nil } start := time.Now() - if err := drainTrailingProgress(progressCh, deliver, 20*time.Millisecond); err != nil { + 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 @@ -701,40 +736,70 @@ func TestDrainTrailingProgressReturnsOnceWindowElapsesWithNoArrival(t *testing.T } } -func TestDrainTrailingProgressResetsWindowPerArrival(t *testing.T) { - // Three events spaced 300ms apart against a 500ms window. Each gap is - // under the window (200ms margin), but the cumulative spacing is not: - // without the per-arrival reset the single 500ms timer fires between the - // second and third event and the drain returns having delivered only 2. - const window = 500 * time.Millisecond - const gap = 300 * time.Millisecond +func TestDrainTrailingProgressUsesAbsoluteWindowDuringContinuousArrivals(t *testing.T) { + const ( + window = 30 * time.Millisecond + gap = 2 * time.Millisecond + ) progressCh := make(chan StreamEvent, 1) senderStop := make(chan struct{}) - t.Cleanup(func() { close(senderStop) }) + senderDone := make(chan struct{}) go func() { - for range 3 { + defer close(senderDone) + ticker := time.NewTicker(gap) + defer ticker.Stop() + for { select { - case progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)}: + case <-ticker.C: + select { + case progressCh <- StreamEvent{Data: []byte(`{"progress":1}`)}: + case <-senderStop: + return + } case <-senderStop: - // The drain exited early (a failure mode under test); don't - // leave this sender blocked on a channel nobody reads. return } - time.Sleep(gap) } }() - delivered := 0 - deliver := func(StreamEvent) error { - delivered++ - return nil - } - if err := drainTrailingProgress(progressCh, deliver, window); err != nil { + 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) } - if delivered != 3 { - t.Fatalf("delivered %d events, want all 3 (window must reset on each arrival)", delivered) +} + +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) } } diff --git a/internal/mcp/standalone_stream.go b/internal/mcp/standalone_stream.go index b985ea06..6e904ed1 100644 --- a/internal/mcp/standalone_stream.go +++ b/internal/mcp/standalone_stream.go @@ -29,8 +29,8 @@ const standaloneWaiterEventBuffer = 32 // 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. Each arrival resets it so -// a trailing burst is drained completely. +// 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 { From 28d30ae51526c1dd01931f9563149c17831bdaea Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Mon, 27 Jul 2026 10:07:10 -0400 Subject: [PATCH 17/18] test: extract pollUntil and newTestAgentServer to cut duplication Six near-identical "poll until condition" loops across internal/invocation and internal/api collapse into one pollUntil(t, timeout, interval, cond) helper per package. Three identical db->resolver->service->handler-> httptest.Server construction blocks in internal/api collapse into newTestAgentServer. Left two structurally different lookalike loops (stream_timeout_test.go's continuous-pressure test, sse_relay_test.go's heartbeat-corruption stress test) untouched since they hammer for a fixed duration rather than poll for a condition. --- internal/api/agent_server_helpers_test.go | 74 ++++++++++++++++++++ internal/api/mcp_everything_test.go | 53 +++----------- internal/api/mcp_standalone_stream_test.go | 53 +++----------- internal/api/sse_relay_test.go | 27 +------ internal/invocation/service_test.go | 35 +++++---- internal/invocation/stream_execution_test.go | 22 +++--- 6 files changed, 126 insertions(+), 138 deletions(-) create mode 100644 internal/api/agent_server_helpers_test.go 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/mcp_everything_test.go b/internal/api/mcp_everything_test.go index 14c8f77f..bc805470 100644 --- a/internal/api/mcp_everything_test.go +++ b/internal/api/mcp_everything_test.go @@ -23,22 +23,13 @@ package api import ( "bufio" - "context" - "database/sql" "fmt" "net" "net/http" - "net/http/httptest" "os/exec" "strings" "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" ) // startEverythingServer launches the real @modelcontextprotocol/server-everything @@ -67,49 +58,23 @@ func startEverythingServer(t *testing.T) (baseURL string) { }) baseURL = fmt.Sprintf("http://127.0.0.1:%d/mcp", port) - deadline := time.Now().Add(30 * time.Second) - for time.Now().Before(deadline) { + 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 { - _ = conn.Close() - return baseURL + if err != nil { + return false } - time.Sleep(200 * time.Millisecond) + _ = 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)") } - 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 "" + return baseURL } func TestMCPToolsCallAgainstRealEverythingServer(t *testing.T) { baseURL := startEverythingServer(t) - db, err := sql.Open("sqlite", ":memory:") - if err != nil { - t.Fatalf("open db: %v", err) - } - defer 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: "everything", Mode: "http", BaseURL: baseURL, Enabled: true, TimeoutSeconds: 30}}, - }) - 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{}, 30*time.Second, nil, nil, nil, nil, - ) - 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()) - defer agentServer.Close() + 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)) diff --git a/internal/api/mcp_standalone_stream_test.go b/internal/api/mcp_standalone_stream_test.go index 93c5f0fb..277990a6 100644 --- a/internal/api/mcp_standalone_stream_test.go +++ b/internal/api/mcp_standalone_stream_test.go @@ -26,22 +26,13 @@ package api import ( "bufio" - "context" - "database/sql" "fmt" "net" "net/http" - "net/http/httptest" "os/exec" "strings" "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" ) // startStandaloneFixtureServer launches the real FastMCP-based fixture @@ -71,49 +62,23 @@ func startStandaloneFixtureServer(t *testing.T) (baseURL string) { }) baseURL = fmt.Sprintf("http://127.0.0.1:%d/mcp", port) - deadline := time.Now().Add(30 * time.Second) - for time.Now().Before(deadline) { + 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 { - _ = conn.Close() - return baseURL + if err != nil { + return false } - time.Sleep(200 * time.Millisecond) + _ = 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)") } - 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 "" + return baseURL } func TestMCPToolsCallAgainstRealStandaloneStreamServer(t *testing.T) { baseURL := startStandaloneFixtureServer(t) - db, err := sql.Open("sqlite", ":memory:") - if err != nil { - t.Fatalf("open db: %v", err) - } - defer 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: "standalone-fixture", Mode: "http", BaseURL: baseURL, Enabled: true, TimeoutSeconds: 30}}, - }) - 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{}, 30*time.Second, nil, nil, nil, nil, - ) - 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()) - defer agentServer.Close() + 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)) diff --git a/internal/api/sse_relay_test.go b/internal/api/sse_relay_test.go index 92760bc9..0632b23d 100644 --- a/internal/api/sse_relay_test.go +++ b/internal/api/sse_relay_test.go @@ -3,7 +3,6 @@ package api import ( "bufio" "context" - "database/sql" "encoding/json" "fmt" "net/http" @@ -13,9 +12,7 @@ import ( "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" ) @@ -518,29 +515,7 @@ func TestMCPToolsCallEndToEndRelaysLiveBeforeTerminalResponseExists(t *testing.T })) defer upstream.Close() - db, err := sql.Open("sqlite", ":memory:") - if err != nil { - t.Fatalf("open db: %v", err) - } - defer 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: "demo", Mode: "http", BaseURL: upstream.URL, Enabled: true, TimeoutSeconds: 5}}, - }) - 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{}, 5*time.Second, nil, nil, nil, nil, - ) - - h := NewHandler(svc, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil) - agentServer := httptest.NewServer(h.Routes()) - defer agentServer.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)) diff --git a/internal/invocation/service_test.go b/internal/invocation/service_test.go index fc234500..d0782bda 100644 --- a/internal/invocation/service_test.go +++ b/internal/invocation/service_test.go @@ -257,15 +257,14 @@ func TestMultiReplicaApprovalRaceOverwritesApprovedRowAsFailed(t *testing.T) { }() var invocationID string - deadline := time.Now().Add(2 * time.Second) - for time.Now().Before(deadline) { + 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 - break + return true } - time.Sleep(10 * time.Millisecond) - } + return false + }) if invocationID == "" { t.Fatal("invocation never reached pending_approval") } @@ -496,17 +495,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) } @@ -1521,6 +1516,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_test.go b/internal/invocation/stream_execution_test.go index 0933e4b8..95490ba0 100644 --- a/internal/invocation/stream_execution_test.go +++ b/internal/invocation/stream_execution_test.go @@ -636,21 +636,19 @@ func TestInvokeStreamingApprovalGateDoesNotTouchSinkBeforeApproval(t *testing.T) // InvokeStreaming block on the approval channel until the suite // timeout. var pendingID string - deadline := time.Now().Add(10 * time.Second) - for pendingID == "" && time.Now().Before(deadline) { + pollUntil(t, 10*time.Second, 5*time.Millisecond, func() bool { list, err := service.List(context.Background(), invocation.InvocationListFilter{Limit: 10}) - if err == nil { - for _, item := range list.Items { - if item.Status == invocation.StatusPendingApproval { - pendingID = item.InvocationID - break - } - } + if err != nil { + return false } - if pendingID == "" { - time.Sleep(5 * time.Millisecond) + 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 From 6909bd92172bc5f8f2da2d31ddb7e6087dc3bd2a Mon Sep 17 00:00:00 2001 From: Kam Lall Date: Mon, 27 Jul 2026 10:35:04 -0400 Subject: [PATCH 18/18] add more comments on the sequences of things --- internal/api/mcp_everything_test.go | 17 +++++ internal/api/mcp_standalone_stream_test.go | 18 +++++ internal/api/sse_relay_test.go | 50 ++++++++++++++ internal/invocation/service_test.go | 28 ++++++++ internal/invocation/stream_execution_test.go | 49 ++++++++++++++ internal/mcp/http_stream_test.go | 42 ++++++++++++ internal/mcp/standalone_stream_test.go | 70 ++++++++++++++++++++ 7 files changed, 274 insertions(+) diff --git a/internal/api/mcp_everything_test.go b/internal/api/mcp_everything_test.go index bc805470..09a1b005 100644 --- a/internal/api/mcp_everything_test.go +++ b/internal/api/mcp_everything_test.go @@ -71,6 +71,23 @@ func startEverythingServer(t *testing.T) (baseURL string) { 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) diff --git a/internal/api/mcp_standalone_stream_test.go b/internal/api/mcp_standalone_stream_test.go index 277990a6..b6e8ed31 100644 --- a/internal/api/mcp_standalone_stream_test.go +++ b/internal/api/mcp_standalone_stream_test.go @@ -75,6 +75,24 @@ func startStandaloneFixtureServer(t *testing.T) (baseURL string) { 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) diff --git a/internal/api/sse_relay_test.go b/internal/api/sse_relay_test.go index 0632b23d..5aafe660 100644 --- a/internal/api/sse_relay_test.go +++ b/internal/api/sse_relay_test.go @@ -162,6 +162,32 @@ func (f *switchableFailingWriter) Write(p []byte) (int, error) { // 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) @@ -459,6 +485,30 @@ func readNextSSEFrame(t *testing.T, reader *bufio.Reader) sseEventFrame { // 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) { diff --git a/internal/invocation/service_test.go b/internal/invocation/service_test.go index d0782bda..1602fbac 100644 --- a/internal/invocation/service_test.go +++ b/internal/invocation/service_test.go @@ -207,6 +207,34 @@ func TestInvokeFailsClosedWhenRuleLoadFails(t *testing.T) { // 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 diff --git a/internal/invocation/stream_execution_test.go b/internal/invocation/stream_execution_test.go index 95490ba0..a7b6bd31 100644 --- a/internal/invocation/stream_execution_test.go +++ b/internal/invocation/stream_execution_test.go @@ -271,6 +271,33 @@ func TestInvokeStreamingAuditCapsEnforcedWithoutSuppressingRelay(t *testing.T) { } } +// 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") @@ -607,6 +634,28 @@ func TestInvokeStreamingMidStreamSessionRetryRefusalMarksFailedWithDistinctReaso } } +// 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") diff --git a/internal/mcp/http_stream_test.go b/internal/mcp/http_stream_test.go index 5d227051..b8b43bdf 100644 --- a/internal/mcp/http_stream_test.go +++ b/internal/mcp/http_stream_test.go @@ -736,6 +736,30 @@ func TestDrainTrailingProgressReturnsOnceWindowElapsesWithNoArrival(t *testing.T } } +// 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 @@ -808,6 +832,24 @@ func TestDrainTrailingProgressStopsWhenCallIsCanceled(t *testing.T) { // 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) { diff --git a/internal/mcp/standalone_stream_test.go b/internal/mcp/standalone_stream_test.go index 0cf0c0de..8ea26e2e 100644 --- a/internal/mcp/standalone_stream_test.go +++ b/internal/mcp/standalone_stream_test.go @@ -243,6 +243,26 @@ func TestInvokeStreamStandaloneStreamRelaysProgressNotification(t *testing.T) { } } +// 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{}) @@ -321,6 +341,36 @@ func TestInvokeStreamStandaloneProgressResetsIdleTimeout(t *testing.T) { } } +// 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 @@ -726,6 +776,26 @@ func TestStandaloneStreamRefcountsSharedConnection(t *testing.T) { 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}