feat(904): MCP progress notifications + SSE for long-running tool calls
Follow-up to bug 903. The attach fix made run_tests retries safe, but
agents still observed the underlying MCP transport timeout as a
tool-call error and had to handle it via retry. Implement the proper
fix: MCP `notifications/progress` events keep the client's transport
timer alive so the call never errors from the agent's perspective.
What changed:
server/src/http/mcp/progress.rs (new)
- `ProgressEmitter` (progressToken + mpsc sender) installed in a
`tokio::task_local!` scope by the SSE response path.
- `emit_progress(progress, total, message)` builds a JSON-RPC
`notifications/progress` message and sends it via the channel.
No-op when no emitter is in scope (plain JSON path / tests / API
runtimes), so tool handlers can call it unconditionally.
server/src/http/mcp/mod.rs
- mcp_post_handler now detects `Accept: text/event-stream` AND a
`params._meta.progressToken` on tools/call. When both are present,
routes through `sse_tools_call` instead of the plain JSON path.
- sse_tools_call: spawns the dispatch task with the emitter installed,
builds an SSE stream that interleaves incoming progress events with
the final JSON-RPC response, with a 15s keep-alive interval as a
backstop for tools that don't emit their own progress.
- Plain JSON behaviour is unchanged for non-SSE clients and for
everything other than tools/call.
server/src/http/mcp/shell_tools/script.rs
- tool_run_tests poll loop emits `notifications/progress` every 25s
of elapsed time (well below the typical ~60s MCP transport
timeout). Attached callers (the bug 903 fix path) also emit so
their MCP socket stays alive while waiting for the in-flight job.
- Output filtering: on a passing run the response now returns a
one-line summary ("All N tests passed.") instead of the full
`cargo test` stdout, which was pure noise that burned agent
tokens. Failure output is unchanged (truncated tail with the
`failures:` section and final test_result line). CRDT entry
stores the same filtered value so attached callers see it too.
Tests (3 new):
- emit_progress_no_op_without_emitter — calling outside scope is safe
- emit_progress_sends_notification_when_emitter_installed — full path
- emit_progress_omits_optional_fields — total/message optional
Not changed: coder system_prompts still tell agents to retry on
transport-timeout errors. That advice is now belt-and-braces — if
claude-code's HTTP MCP client honours progress notifications, no agent
will ever observe the error; if not, retry is still safe post-903. We
can drop the retry advice once we've observed the SSE path working in
the field.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,123 @@
|
||||
//! Task-local plumbing for MCP `notifications/progress` events.
|
||||
//!
|
||||
//! Long-running tool handlers (notably `tool_run_tests`) emit progress events
|
||||
//! during execution so the MCP client's transport timer resets and the call
|
||||
//! never surfaces as a tool-call error to the agent. The HTTP MCP handler
|
||||
//! installs an emitter via `tokio::task_local!` before dispatching the tool
|
||||
//! and converts emitted events into SSE messages on the response stream.
|
||||
//!
|
||||
//! Tool handlers call [`emit_progress`] unconditionally — when no emitter is
|
||||
//! installed (plain JSON response path, or test code), the call is a no-op.
|
||||
//! That keeps handler code free of `if let Some(...) =` ceremony for the
|
||||
//! progress-aware branch.
|
||||
//!
|
||||
//! Wire format follows MCP 2025-03-26: the emitter wraps the data in a
|
||||
//! standard JSON-RPC notification with method `notifications/progress` and a
|
||||
//! `progressToken` echoed back to the client.
|
||||
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
|
||||
/// Per-request channel for emitting progress notifications back to an MCP
|
||||
/// client. Installed in a `tokio::task_local!` scope by the SSE response
|
||||
/// path and consumed by the SSE stream producer.
|
||||
#[derive(Clone)]
|
||||
pub struct ProgressEmitter {
|
||||
/// The client-supplied opaque token from `params._meta.progressToken`.
|
||||
/// Echoed back unchanged in every progress notification so the client
|
||||
/// can correlate progress with the originating request.
|
||||
pub token: Value,
|
||||
/// Channel into the SSE response stream. Each value sent is a
|
||||
/// fully-formed `notifications/progress` JSON-RPC message ready to be
|
||||
/// serialised as an SSE `data:` field.
|
||||
pub tx: UnboundedSender<Value>,
|
||||
}
|
||||
|
||||
tokio::task_local! {
|
||||
/// Set by the SSE response path before dispatching a tool call. Unset
|
||||
/// in the plain JSON path and in tests, where [`emit_progress`] no-ops.
|
||||
pub static EMITTER: ProgressEmitter;
|
||||
}
|
||||
|
||||
/// Emit a progress notification to the current request's SSE stream, if one
|
||||
/// is attached. No-op when no emitter is in scope (plain JSON path).
|
||||
///
|
||||
/// `progress` is a monotonically increasing value (typically seconds elapsed
|
||||
/// for long-running tools); `total` is optional and omitted when unknown;
|
||||
/// `message` is a short human-readable description of the current state.
|
||||
pub fn emit_progress(progress: f64, total: Option<f64>, message: Option<&str>) {
|
||||
let _ = EMITTER.try_with(|e| {
|
||||
let mut params = json!({
|
||||
"progressToken": e.token,
|
||||
"progress": progress,
|
||||
});
|
||||
if let Some(t) = total {
|
||||
params["total"] = json!(t);
|
||||
}
|
||||
if let Some(m) = message {
|
||||
params["message"] = json!(m);
|
||||
}
|
||||
let notification = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/progress",
|
||||
"params": params,
|
||||
});
|
||||
// Send is fire-and-forget. If the receiver dropped (client
|
||||
// disconnected mid-stream), we don't care — the tool dispatch
|
||||
// task keeps running and writes its final state to the CRDT.
|
||||
let _ = e.tx.send(notification);
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::sync::mpsc::unbounded_channel;
|
||||
|
||||
#[tokio::test]
|
||||
async fn emit_progress_no_op_without_emitter() {
|
||||
// Calling outside a task_local scope must not panic.
|
||||
emit_progress(1.0, None, Some("hello"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn emit_progress_sends_notification_when_emitter_installed() {
|
||||
let (tx, mut rx) = unbounded_channel();
|
||||
let emitter = ProgressEmitter {
|
||||
token: json!("test-token"),
|
||||
tx,
|
||||
};
|
||||
EMITTER
|
||||
.scope(emitter, async {
|
||||
emit_progress(5.0, Some(10.0), Some("halfway"));
|
||||
})
|
||||
.await;
|
||||
|
||||
let notif = rx.recv().await.expect("notification must be delivered");
|
||||
assert_eq!(notif["method"], "notifications/progress");
|
||||
assert_eq!(notif["params"]["progressToken"], "test-token");
|
||||
assert_eq!(notif["params"]["progress"], 5.0);
|
||||
assert_eq!(notif["params"]["total"], 10.0);
|
||||
assert_eq!(notif["params"]["message"], "halfway");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn emit_progress_omits_optional_fields() {
|
||||
let (tx, mut rx) = unbounded_channel();
|
||||
let emitter = ProgressEmitter {
|
||||
token: json!(42),
|
||||
tx,
|
||||
};
|
||||
EMITTER
|
||||
.scope(emitter, async {
|
||||
emit_progress(1.0, None, None);
|
||||
})
|
||||
.await;
|
||||
|
||||
let notif = rx.recv().await.unwrap();
|
||||
assert_eq!(notif["params"]["progressToken"], 42);
|
||||
assert_eq!(notif["params"]["progress"], 1.0);
|
||||
assert!(notif["params"].get("total").is_none());
|
||||
assert!(notif["params"].get("message").is_none());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user