fix: add SSE heartbeats and defer trace completion to prevent deadlock on MCP SSE streaming

[Pratham-Mishra04]

Summary

Fixes a deadlock and connection reliability issue in the MCP server's SSE handler. Without this fix, post-hook middleware (transport-plugin and tracing) would attempt to materialize the SSE response body during cleanup, blocking indefinitely waiting for an EOF that never arrives while the streaming goroutine is still running. Additionally, idle SSE connections had no heartbeat mechanism, making it impossible to detect client disconnects since fasthttp never cancels the bifrost context on its own.

Changes

• Sets BifrostContextKeyDeferTraceCompletion on the request context before streaming begins, signaling to transport-plugin and tracing middlewares to skip eager body materialization that would deadlock on the open SSE stream.
• Adds X-Accel-Buffering: no response header to prevent nginx and similar reverse proxies from buffering SSE events.
• Replaces the bare <-(*bifrostCtx).Done() wait with a 15-second ticker that sends SSE comment heartbeats (: ping\n\n). This keeps idle connections alive through proxies and detects client disconnects via reader.Send() returning false.
• Guards the initial SSE init message send with a reader.Send() return check so a disconnected client causes an early exit rather than continuing the goroutine.

Type of change

• Bug fix

Affected areas

• Transports (HTTP)

How to test

Connect an SSE client to the MCP server endpoint and verify:

1. Events are received without buffering through a reverse proxy.
2. Disconnecting the client causes the server-side goroutine to exit cleanly within ~15 seconds (next heartbeat tick).
3. Tracing and transport-plugin middleware complete without hanging after the SSE stream ends.

go test ./...

Breaking changes

• No

Related issues

Security considerations

None.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Pratham-Mishra04]

• refactor: move virtual key validation and context attachment before passthrough check #3213
• fix: add SSE heartbeats and defer trace completion to prevent deadlock on MCP SSE streaming #3212 👈 (View in Graphite)
• fix: handle per-user OAuth re-auth, refresh token expiry, and reconnect UX #3211
• feat: remove direct key bypass from HTTP gateway and Go SDK #3241
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 5/5

Safe to merge — the fix correctly mirrors the deferred-completion pattern already proven in inference.go, and the defer ordering (runCompleter → cancel → reader.Done → traceCompleter) ensures post-hooks and spans are always finalised cleanly.

All three concerns raised in the previous review round (leaked trace span, magic number heartbeat constant, cancel-before-completer ordering) are fully addressed. The goroutine's cleanup sequence now matches inference.go. The 100ms atomic poll for the completer slot is a bounded busy-wait; in the absence of a TransportInterceptorMiddleware the slot stays nil and the wait fires on every disconnect, but this is a cosmetic latency concern rather than a correctness problem. No race conditions were found: completerSlot escapes to the heap, traceCompleter is captured before goroutine launch (tracing middleware sets it before calling next(ctx)), and ctx is not accessed inside the goroutine.

No files require special attention.

Important Files Changed

Filename	Overview
transports/bifrost-http/handlers/mcpserver.go	Adds deferred trace completion, transport post-hook completer slot, SSE heartbeat loop, and X-Accel-Buffering header to fix deadlock and connection reliability issues; structure mirrors inference.go and the implementation looks sound with one minor liveness concern around the 100ms completer poll

_{Reviews (8): Last reviewed commit: "fix: mcp sse hang fixes" | Re-trigger Greptile}

[coderabbitai]

Warning

Rate limit exceeded

@Pratham-Mishra04 has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 9 minutes and 21 seconds before requesting another review.

To continue reviewing without waiting, purchase usage credits in the billing tab.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: ab0eaa7b-782e-44f7-ab6a-fbe51734dbdf

📥 Commits

Reviewing files that changed from the base of the PR and between e2bc0b5 and 83ff8a4.

📒 Files selected for processing (1)

• transports/bifrost-http/handlers/mcpserver.go

📝 Walkthrough

Walkthrough

Signals streaming context, exposes an atomic transport post-hook completer slot, defers trace completion, sets X-Accel-Buffering: no, verifies initial SSE frame delivery, emits 15s SSE comment heartbeats to detect disconnects, and on stop runs the transport completer to collect plugin logs.

Changes

SSE Streaming Robustness & Tracing Coordination

Layer / File(s)	Summary
Constants / Config transports/bifrost-http/handlers/mcpserver.go	Adds const sseHeartbeatInterval = 15 * time.Second.
Imports transports/bifrost-http/handlers/mcpserver.go	Adds sync/atomic for atomic.Value usage.
Context Decoration transports/bifrost-http/handlers/mcpserver.go	Sets schemas.BifrostContextKeyDeferTraceCompletion, stores an atomic.Value under schemas.BifrostContextKeyTransportPostHookCompleter, and reads schemas.BifrostContextKeyTraceCompleter from request context.
Response Headers transports/bifrost-http/handlers/mcpserver.go	Adds X-Accel-Buffering: no header in SSE responses.
Initial Frame Delivery transports/bifrost-http/handlers/mcpserver.go	Sends initial connection/opened SSE frame and aborts streaming if reader.Send(buf) fails.
Heartbeat & Streaming Loop transports/bifrost-http/handlers/mcpserver.go	Adds active loop with ticker sending SSE comment heartbeats (: ping\n\n) every 15s; exits when send fails or bifrost context is done.
Shutdown / Completer Invocation transports/bifrost-http/handlers/mcpserver.go	Adds runCompleter closure that waits briefly (~100ms) for a transport post-hook completer in the atomic slot, invokes it once to collect transport/plugin logs, defers cancel() and reader.Done(), and calls the trace completer with captured logs.
Manifests go.mod	Manifest touched; no tests changed in this diff.

Sequence Diagram(s)

sequenceDiagram
participant Client
participant SSEHandler as Bifrost HTTP handler
participant TransportPostHook as Transport Post-hook (atomic.Value)
participant TraceCompleter as Trace completer (ctx.UserValue)

Client->>SSEHandler: Open SSE connection
SSEHandler->>Client: Send "connection/opened" SSE frame
alt send fails
  SSEHandler-->>Client: Abort streaming / exit
else send succeeds
  loop every 15s
    SSEHandler->>Client: ": ping" SSE comment heartbeat
  end
  alt client disconnect or bifrostCtx.Done
    SSEHandler->>TransportPostHook: wait ~100ms for completer
    TransportPostHook-->>SSEHandler: provide completer/logs (if set)
    SSEHandler->>TraceCompleter: invoke with collected logs
    SSEHandler-->>Client: cleanup (cancel, reader.Done)
  end
end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

I’m a rabbit by the wire, I ping and softly hum,
I wait a beat for post-hooks, then gather what they’ve done.
I tuck the trace in ribbons, append the plugin logs,
Then bow and hop away as quiet goes the logs. 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main changes: adding SSE heartbeats and deferring trace completion to fix a deadlock in MCP SSE streaming.
Description check	✅ Passed	The description comprehensively covers the problem, changes, testing approach, and affected areas following the template structure with all critical sections completed.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 05-04-fix_mcp_sse_hang_fixes

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@transports/bifrost-http/handlers/mcpserver.go`:
- Around line 150-155: The SSE handler sets
BifrostContextKeyDeferTraceCompletion but does not seed the deferred post-hook
completer, so middlewares.go skips snapshot/completion and pooled objects leak;
before returning the stream in the handler in mcpserver.go, create and
SetUserValue for BifrostContextKeyTransportPostHookCompleter with the same
completer type used by other streaming handlers, and ensure the stream teardown
path drains/invokes that completer (mirror the logic used in the other streaming
handlers) so the deferred transport/tracing path runs and pooled objects are
released.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: eab98ee8-e58b-4652-981f-71a6cf38d843

📥 Commits

Reviewing files that changed from the base of the PR and between 09c4f02 and 08b9a17.

📒 Files selected for processing (1)

• transports/bifrost-http/handlers/mcpserver.go

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

transports/bifrost-http/handlers/mcpserver.go (1)

71-75: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Duplicate tool filter registration.

The same WithToolFilter call is registered twice on globalMCPServer with identical comments. This appears to be an accidental duplication that could cause the filter to run twice per request.

🔧 Proposed fix

 	// Register per-request tool filter so x-bf-mcp-include-clients and x-bf-mcp-include-tools are respected on tools/list
 	server.WithToolFilter(handler.makeIncludeClientsFilter())(handler.globalMCPServer)
-
-	// Register per-request tool filter so x-bf-mcp-include-clients and x-bf-mcp-include-tools are respected on tools/list
-	server.WithToolFilter(handler.makeIncludeClientsFilter())(handler.globalMCPServer)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@transports/bifrost-http/handlers/mcpserver.go` around lines 71 - 75, The
WithToolFilter registration for handler.makeIncludeClientsFilter() is duplicated
on handler.globalMCPServer; remove the redundant call so the filter is only
registered once (keep a single
server.WithToolFilter(handler.makeIncludeClientsFilter())(handler.globalMCPServer)
and delete the duplicate) to prevent the filter executing twice per request.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@transports/bifrost-http/handlers/mcpserver.go`:
- Around line 71-75: The WithToolFilter registration for
handler.makeIncludeClientsFilter() is duplicated on handler.globalMCPServer;
remove the redundant call so the filter is only registered once (keep a single
server.WithToolFilter(handler.makeIncludeClientsFilter())(handler.globalMCPServer)
and delete the duplicate) to prevent the filter executing twice per request.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 21380091-38fb-49a8-b202-2559e1566838

📥 Commits

Reviewing files that changed from the base of the PR and between 3dcc042 and e2bc0b5.

📒 Files selected for processing (1)

• transports/bifrost-http/handlers/mcpserver.go

[akshaydeo]

Merge activity

• May 6, 8:41 AM UTC: A user started a stack merge that includes this pull request via Graphite.
• May 6, 8:43 AM UTC: @akshaydeo merged this pull request with Graphite.

The base branch was changed.