Skip to main content

max / makenotwork

kberg: add MCP resources, subscriptions, server push, and cancellation Adds the resource + push half of MCP so the tools-only server can expose live objects an agent subscribes to (deox exposes editor buffers this way). - resource module: dynamic ResourceRegistry (register/remove at runtime, async read handler per resource), ResourceDescriptor, ResourceContents. Serves resources/list + resources/read. - Subscriptions: resources/subscribe + unsubscribe (session-scoped); ResourceRegistry::notify_updated pushes notifications/resources/updated over SSE, fanned out to sessions that filter by their subscription set. - Cancellation: client-sent notifications/cancelled aborts the matching in-flight tools/call; BoundServer::cancel_all / cancel_session abort in-flight calls and push notifications/cancelled to the client. - Transport: /mcp now serves POST + a GET SSE stream + DELETE, with Mcp-Session-Id session management (new server/session.rs). Subscriptions and cancellation are HTTP-only; stdio keeps tools + resource read. 25 unit + 6 over-the-wire integration tests (incl. e2e SSE update push and app-cancel-over-SSE); clippy clean; no-default-features builds clean. New dep: tokio-stream (server feature). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Author: Max Johnson <me@maxj.phd> · 2026-07-14 21:00 UTC
Commit: 2e8fb44b5532fefa9cdf93c84285072e1a5fbf62
Parent: b3d1651
12 files changed, +1170 insertions, -33 deletions
M shared/kberg/Cargo.lock +13
@@ -580,6 +580,7 @@ dependencies = [
580 580 "serde_json",
581 581 "thiserror",
582 582 "tokio",
583 + "tokio-stream",
583 584 "tower",
584 585 "tower-http",
585 586 "tracing",
@@ -1208,6 +1209,18 @@ dependencies = [
1208 1209 ]
1209 1210
1210 1211 [[package]]
1212 + name = "tokio-stream"
1213 + version = "0.1.18"
1214 + source = "registry+https://github.com/rust-lang/crates.io-index"
1215 + checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70"
1216 + dependencies = [
1217 + "futures-core",
1218 + "pin-project-lite",
1219 + "tokio",
1220 + "tokio-util",
1221 + ]
1222 +
1223 + [[package]]
1211 1224 name = "tokio-util"
1212 1225 version = "0.7.18"
1213 1226 source = "registry+https://github.com/rust-lang/crates.io-index"
M shared/kberg/Cargo.toml +2 -1
@@ -10,7 +10,7 @@ categories = ["api-bindings", "development-tools"]
10 10
11 11 [features]
12 12 default = ["server", "ollama"]
13 - server = ["dep:axum", "dep:tokio", "dep:tower", "dep:tower-http"]
13 + server = ["dep:axum", "dep:tokio", "dep:tokio-stream", "dep:tower", "dep:tower-http"]
14 14 ollama = ["dep:reqwest", "dep:tokio"]
15 15
16 16 [dependencies]
@@ -28,6 +28,7 @@ tower-http = { version = "0.6", optional = true, features = ["trace"] }
28 28 reqwest = { version = "0.12", optional = true, features = ["json"] }
29 29
30 30 tokio = { version = "1", optional = true, features = ["rt-multi-thread", "macros", "net", "sync", "time", "io-std", "io-util"] }
31 + tokio-stream = { version = "0.1", optional = true, features = ["sync"] }
31 32
32 33 [dev-dependencies]
33 34 tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "sync", "time"] }
M shared/kberg/README.md +11 -4
@@ -6,8 +6,9 @@ Named after Königsberg — the city whose seven bridges seeded graph theory (Eu
6 6
7 7 ## Ships
8 8
9 - - `Tool` trait + `ToolRegistry` — the app-side surface.
10 - - Streamable HTTP MCP server (JSON-RPC 2.0) — how any MCP-speaking client reaches your app.
9 + - `Tool` trait + `ToolRegistry` — the app-side action surface.
10 + - `ResourceRegistry` — the app-side object surface: live things a client lists, reads, and subscribes to. Dynamic (register/remove at runtime); `notify_updated` pushes a change to subscribed clients.
11 + - Streamable HTTP MCP server (JSON-RPC 2.0) — how any MCP-speaking client reaches your app. POST for requests, a GET SSE stream for server-initiated notifications, DELETE to end a session.
11 12 - Provider-agnostic `Agent` + `InferenceProvider` trait — the tool-use loop is decoupled from the LLM backend.
12 13 - `provider::ollama::OllamaProvider` — reference implementation for a local Ollama server.
13 14
@@ -34,7 +35,10 @@ Capability IDs are a public contract the app owns. Renaming means a new capabili
34 35 |---|---|---|
35 36 | `Tool` + `ToolRegistry` | shipped | register app functions as MCP tools |
36 37 | `ToolKind::Read / Write(capability)` | shipped | webhook-style write permissions |
37 - | Streamable HTTP MCP server | shipped | speak MCP over local HTTP |
38 + | Streamable HTTP MCP server | shipped | speak MCP over local HTTP (POST + GET SSE + DELETE) |
39 + | `ResourceRegistry` (list / read) | shipped | expose live objects a client lists and reads |
40 + | Resource subscriptions + server push | shipped | `resources/subscribe` + `notifications/resources/updated` over SSE |
41 + | Request cancellation | shipped | `notifications/cancelled` both ways; app aborts an in-flight call and tells the client |
38 42 | `InferenceProvider` + `Agent` | shipped | provider-agnostic tool-use loop |
39 43 | Ollama provider | shipped | drive tools from a local Ollama model |
40 44 | Hard refusals (`Refusal`) | shipped | registered "not offered" tools with stable reason |
@@ -54,10 +58,13 @@ src/
54 58 lib.rs
55 59 error.rs
56 60 tool.rs # Tool, ToolKind, WriteCapability, ToolRegistry, Refusal, SurfaceProjection
61 + resource.rs # ResourceRegistry, ResourceDescriptor, ResourceContents (+ notify_updated)
57 62 agent.rs # Agent, InferenceProvider, Message, RunOutcome
58 63 server/ # feature = "server"
59 - mod.rs
64 + mod.rs # POST/GET-SSE/DELETE routes, JSON-RPC dispatch
60 65 protocol.rs
66 + session.rs # per-session subscriptions, SSE out channel, cancellation
67 + stdio.rs
61 68 provider/ # provider adapters (feature-gated)
62 69 mod.rs
63 70 ollama.rs # feature = "ollama"
M shared/kberg/examples/toy.rs +1
@@ -113,6 +113,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
113 113 server::ServeConfig {
114 114 projection: kberg::SurfaceProjection::Full,
115 115 grants: Some(HashSet::from(["remember_project".to_string()])),
116 + resources: None,
116 117 },
117 118 )
118 119 .await?;
M shared/kberg/src/error.rs +9
@@ -7,6 +7,15 @@ pub enum Error {
7 7 #[error("tool not found: {0}")]
8 8 ToolNotFound(String),
9 9
10 + #[error("resource not found: {0}")]
11 + ResourceNotFound(String),
12 +
13 + /// An in-flight request was cancelled — by the client (`notifications/
14 + /// cancelled`) or by the app (a human took the buffer back). The `String`
15 + /// is the cancelled request id.
16 + #[error("request cancelled: {0}")]
17 + Cancelled(String),
18 +
10 19 #[error("invalid arguments for tool `{tool}`: {message}")]
11 20 InvalidArgs { tool: String, message: String },
12 21
M shared/kberg/src/lib.rs +5 -1
@@ -6,7 +6,9 @@
6 6 //!
7 7 //! # Shape
8 8 //!
9 - //! - [`Tool`] + [`ToolRegistry`] — app-side surface.
9 + //! - [`Tool`] + [`ToolRegistry`] — app-side action surface.
10 + //! - [`resource::ResourceRegistry`] — app-side object surface: live things a
11 + //! client lists and reads (deox exposes editor buffers this way).
10 12 //! - [`server`] — Streamable HTTP MCP server (JSON-RPC 2.0). Any MCP-speaking
11 13 //! client can reach a Kberg-served app.
12 14 //! - [`agent::Agent`] + [`agent::InferenceProvider`] — provider-agnostic
@@ -34,10 +36,12 @@
34 36
35 37 pub mod agent;
36 38 mod error;
39 + pub mod resource;
37 40 pub mod tool;
38 41
39 42 pub use agent::{Agent, AgentConfig, InferenceProvider, Message, Role, RunOutcome, ToolCall};
40 43 pub use error::{Error, Result};
44 + pub use resource::{ResourceContents, ResourceDescriptor, ResourceRegistry};
41 45 pub use tool::{
42 46 ContentPart, Refusal, SurfaceProjection, Tool, ToolCallResult, ToolKind, ToolRegistry,
43 47 ToolSpec, WriteCapability,
A shared/kberg/src/resource.rs +294
@@ -0,0 +1,294 @@
1 + //! Resources — the read side of the MCP object surface.
2 + //!
3 + //! Where [`crate::tool`] exposes an app's *actions*, resources expose its live
4 + //! *objects*: things a client reads and (in the push layer) subscribes to. deox
5 + //! exposes editor buffers this way; a content-addressed store could expose its
6 + //! samples.
7 + //!
8 + //! Unlike tools, resources are **dynamic** — a buffer opens and closes at
9 + //! runtime — so [`ResourceRegistry`] is interior-mutable and cloneable:
10 + //! [`register`](ResourceRegistry::register) and [`remove`](ResourceRegistry::remove)
11 + //! take `&self` and any clone observes the change. Each resource carries an
12 + //! async read handler the registry invokes on `resources/read`; the app owns the
13 + //! underlying state (for deox, behind the main-loop channel).
14 + //!
15 + //! The subscription + `notifications/resources/updated` push layer builds on
16 + //! this registry; see the server module.
17 +
18 + use std::collections::HashMap;
19 + use std::future::Future;
20 + use std::pin::Pin;
21 + use std::sync::{Arc, RwLock};
22 +
23 + use serde::Serialize;
24 +
25 + use crate::error::{Error, Result};
26 +
27 + type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
28 + type ReadHandler = Arc<dyn Fn(String) -> BoxFuture<Result<ResourceContents>> + Send + Sync>;
29 +
30 + /// Metadata for a single resource, as it appears in `resources/list`.
31 + ///
32 + /// `uri` is the stable identifier the client passes back to `resources/read`
33 + /// and `resources/subscribe`. Choose a scheme the app owns (deox uses
34 + /// `buffer://<path>`); it is an addressing contract, not a filesystem path.
35 + #[derive(Debug, Clone, Serialize)]
36 + pub struct ResourceDescriptor {
37 + pub uri: String,
38 + pub name: String,
39 + #[serde(skip_serializing_if = "Option::is_none")]
40 + pub description: Option<String>,
41 + #[serde(rename = "mimeType", skip_serializing_if = "Option::is_none")]
42 + pub mime_type: Option<String>,
43 + }
44 +
45 + impl ResourceDescriptor {
46 + pub fn new(uri: impl Into<String>, name: impl Into<String>) -> Self {
47 + Self {
48 + uri: uri.into(),
49 + name: name.into(),
50 + description: None,
51 + mime_type: None,
52 + }
53 + }
54 +
55 + pub fn with_description(mut self, description: impl Into<String>) -> Self {
56 + self.description = Some(description.into());
57 + self
58 + }
59 +
60 + pub fn with_mime_type(mut self, mime_type: impl Into<String>) -> Self {
61 + self.mime_type = Some(mime_type.into());
62 + self
63 + }
64 + }
65 +
66 + /// The contents of one resource, as returned inside a `resources/read` result.
67 + ///
68 + /// Mirrors MCP's resource-contents shape: a `uri`, an optional `mimeType`, and
69 + /// exactly one of `text` (UTF-8) or `blob` (base64). Use [`text`](Self::text)
70 + /// for buffers and other textual resources; [`blob`](Self::blob) for binary.
71 + #[derive(Debug, Clone, Serialize)]
72 + pub struct ResourceContents {
73 + pub uri: String,
74 + #[serde(rename = "mimeType", skip_serializing_if = "Option::is_none")]
75 + pub mime_type: Option<String>,
76 + #[serde(skip_serializing_if = "Option::is_none")]
77 + pub text: Option<String>,
78 + #[serde(skip_serializing_if = "Option::is_none")]
79 + pub blob: Option<String>,
80 + }
81 +
82 + impl ResourceContents {
83 + /// Textual contents (the common case: a buffer, a note, a config file).
84 + pub fn text(uri: impl Into<String>, text: impl Into<String>) -> Self {
85 + Self {
86 + uri: uri.into(),
87 + mime_type: None,
88 + text: Some(text.into()),
89 + blob: None,
90 + }
91 + }
92 +
93 + /// Binary contents, already base64-encoded.
94 + pub fn blob(uri: impl Into<String>, blob_base64: impl Into<String>) -> Self {
95 + Self {
96 + uri: uri.into(),
97 + mime_type: None,
98 + text: None,
99 + blob: Some(blob_base64.into()),
100 + }
101 + }
102 +
103 + pub fn with_mime_type(mut self, mime_type: impl Into<String>) -> Self {
104 + self.mime_type = Some(mime_type.into());
105 + self
106 + }
107 + }
108 +
109 + struct Entry {
110 + descriptor: ResourceDescriptor,
111 + handler: ReadHandler,
112 + }
113 +
114 + /// Registry of resources a client can list and read.
115 + ///
116 + /// Cheaply cloneable and interior-mutable: the map lives behind `Arc<RwLock>`,
117 + /// so the app can [`register`](Self::register) and [`remove`](Self::remove)
118 + /// resources at runtime while the server holds its own clone. This is the
119 + /// deliberate contrast with [`crate::tool::ToolRegistry`], whose surface is
120 + /// fixed before serving.
121 + ///
122 + /// With the `server` feature, [`notify_updated`](Self::notify_updated) pushes a
123 + /// change to subscribed clients over the SSE stream.
124 + #[derive(Clone)]
125 + pub struct ResourceRegistry {
126 + entries: Arc<RwLock<HashMap<String, Entry>>>,
127 + /// Fan-out of updated URIs to every connected SSE session, which filters by
128 + /// its own subscription set. Capacity bounds how far a slow session may lag
129 + /// before it drops updates (it then re-reads current state on next use).
130 + #[cfg(feature = "server")]
131 + updates: tokio::sync::broadcast::Sender<String>,
132 + }
133 +
134 + impl Default for ResourceRegistry {
135 + fn default() -> Self {
136 + Self::new()
137 + }
138 + }
139 +
140 + impl ResourceRegistry {
141 + pub fn new() -> Self {
142 + Self {
143 + entries: Arc::new(RwLock::new(HashMap::new())),
144 + #[cfg(feature = "server")]
145 + updates: tokio::sync::broadcast::channel(256).0,
146 + }
147 + }
148 +
149 + /// Register (or replace) a resource. `handler` is invoked on each
150 + /// `resources/read` for this uri and returns the resource's current
151 + /// contents — it should fetch live state, not a snapshot taken at
152 + /// registration. Registering an existing uri replaces its descriptor and
153 + /// handler.
154 + pub fn register<F, Fut>(&self, descriptor: ResourceDescriptor, handler: F)
155 + where
156 + F: Fn(String) -> Fut + Send + Sync + 'static,
157 + Fut: Future<Output = Result<ResourceContents>> + Send + 'static,
158 + {
159 + let handler: ReadHandler = Arc::new(move |uri| Box::pin(handler(uri)));
160 + let uri = descriptor.uri.clone();
161 + self.entries
162 + .write()
163 + .unwrap()
164 + .insert(uri, Entry { descriptor, handler });
165 + }
166 +
167 + /// Remove a resource. Returns whether it was present.
168 + pub fn remove(&self, uri: &str) -> bool {
169 + self.entries.write().unwrap().remove(uri).is_some()
170 + }
171 +
172 + pub fn contains(&self, uri: &str) -> bool {
173 + self.entries.read().unwrap().contains_key(uri)
174 + }
175 +
176 + /// Descriptors for every registered resource, for `resources/list`.
177 + pub fn list(&self) -> Vec<ResourceDescriptor> {
178 + self.entries
179 + .read()
180 + .unwrap()
181 + .values()
182 + .map(|e| e.descriptor.clone())
183 + .collect()
184 + }
185 +
186 + /// Read one resource's current contents by uri. The read handler runs
187 + /// without the registry lock held, so a handler may itself touch the
188 + /// registry.
189 + pub async fn read(&self, uri: &str) -> Result<ResourceContents> {
190 + let handler = {
191 + let map = self.entries.read().unwrap();
192 + map.get(uri).map(|e| e.handler.clone())
193 + };
194 + match handler {
195 + Some(handler) => handler(uri.to_string()).await,
196 + None => Err(Error::ResourceNotFound(uri.to_string())),
197 + }
198 + }
199 + }
200 +
201 + #[cfg(feature = "server")]
202 + impl ResourceRegistry {
203 + /// Announce that a resource changed. Subscribed clients receive a
204 + /// `notifications/resources/updated` for `uri` over their SSE stream; they
205 + /// then re-`read` it to get the new contents (the notification carries the
206 + /// uri, not the body). The app calls this after mutating the underlying
207 + /// state — for deox, on a buffer or selection change. A no-op when no
208 + /// session is subscribed.
209 + pub fn notify_updated(&self, uri: impl Into<String>) {
210 + // `send` errors only when there are zero receivers; that is the normal
211 + // "nobody is streaming" case, so the error is deliberately ignored.
212 + let _ = self.updates.send(uri.into());
213 + }
214 +
215 + /// A receiver for the update fan-out, used by each SSE session.
216 + pub(crate) fn subscribe_updates(&self) -> tokio::sync::broadcast::Receiver<String> {
217 + self.updates.subscribe()
218 + }
219 + }
220 +
221 + /// `resources/read` param object: `{ "uri": "..." }`.
222 + #[cfg(feature = "server")]
223 + #[derive(Debug, serde::Deserialize)]
224 + pub(crate) struct ReadResourceParams {
225 + pub uri: String,
226 + }
227 +
228 + /// Wrap a single [`ResourceContents`] into the `resources/read` result shape
229 + /// (`{ "contents": [ ... ] }`).
230 + #[cfg(feature = "server")]
231 + pub(crate) fn read_result(contents: ResourceContents) -> serde_json::Value {
232 + serde_json::json!({ "contents": [contents] })
233 + }
234 +
235 + #[cfg(test)]
236 + mod tests {
237 + use super::*;
238 + use serde_json::json;
239 +
240 + #[tokio::test]
241 + async fn register_list_read_roundtrip() {
242 + let reg = ResourceRegistry::new();
243 + reg.register(
244 + ResourceDescriptor::new("buffer://main.rs", "main.rs").with_mime_type("text/x-rust"),
245 + |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
246 + );
247 +
248 + let list = reg.list();
249 + assert_eq!(list.len(), 1);
250 + assert_eq!(list[0].uri, "buffer://main.rs");
251 +
252 + let contents = reg.read("buffer://main.rs").await.unwrap();
253 + assert_eq!(contents.text.as_deref(), Some("fn main() {}"));
254 +
255 + // Serialized read result matches the MCP `contents` array shape.
256 + let v = read_result(contents);
257 + assert_eq!(v["contents"][0]["uri"], "buffer://main.rs");
258 + assert_eq!(v["contents"][0]["mimeType"], json!(null)); // not set on this content
259 + assert_eq!(v["contents"][0]["text"], "fn main() {}");
260 + }
261 +
262 + #[tokio::test]
263 + async fn read_unknown_uri_is_not_found() {
264 + let reg = ResourceRegistry::new();
265 + let err = reg.read("buffer://gone.rs").await.unwrap_err();
266 + assert!(matches!(err, Error::ResourceNotFound(_)));
267 + }
268 +
269 + #[tokio::test]
270 + async fn remove_is_observed_through_a_clone() {
271 + let reg = ResourceRegistry::new();
272 + let clone = reg.clone();
273 + reg.register(ResourceDescriptor::new("buffer://a", "a"), |uri| async move {
274 + Ok(ResourceContents::text(uri, "x"))
275 + });
276 + assert!(clone.contains("buffer://a"));
277 + assert!(reg.remove("buffer://a"));
278 + assert!(!clone.contains("buffer://a"));
279 + assert!(!reg.remove("buffer://a")); // second remove reports absent
280 + }
281 +
282 + #[tokio::test]
283 + async fn register_replaces_existing_uri() {
284 + let reg = ResourceRegistry::new();
285 + reg.register(ResourceDescriptor::new("buffer://a", "a"), |uri| async move {
286 + Ok(ResourceContents::text(uri, "old"))
287 + });
288 + reg.register(ResourceDescriptor::new("buffer://a", "a"), |uri| async move {
289 + Ok(ResourceContents::text(uri, "new"))
290 + });
291 + assert_eq!(reg.list().len(), 1);
292 + assert_eq!(reg.read("buffer://a").await.unwrap().text.as_deref(), Some("new"));
293 + }
294 + }
M shared/kberg/src/server/mod.rs +298 -24
@@ -1,42 +1,67 @@
1 1 //! Streamable HTTP MCP server.
2 2 //!
3 - //! Exposes a single POST `/mcp` endpoint speaking JSON-RPC 2.0. Handles
4 - //! `initialize`, `tools/list`, and `tools/call`. No SSE for MVP.
3 + //! One `/mcp` route, three methods:
4 + //! - **POST** — JSON-RPC requests: `initialize` (opens a session, returns its
5 + //! id in `Mcp-Session-Id`), `tools/list`, `tools/call`, `resources/list`,
6 + //! `resources/read`, `resources/subscribe`/`unsubscribe`, and the
7 + //! `notifications/cancelled` notification.
8 + //! - **GET** — opens the server→client SSE stream for a session, down which
9 + //! `notifications/resources/updated` and `notifications/cancelled` are
10 + //! pushed.
11 + //! - **DELETE** — ends a session.
5 12 //!
6 13 //! Configuration is passed at startup via [`ServeConfig`]:
7 14 //! - `projection` — Full or Compact surface exposed to clients on this port.
8 15 //! - `grants` — which write capabilities are honored; unmatched writes return
9 16 //! `CapabilityDenied`. `None` bypasses the check (fully-trusted context).
17 + //! - `resources` — an optional [`ResourceRegistry`]; when set, the resource
18 + //! methods above are served and the `resources` capability is advertised.
10 19
11 20 use std::collections::HashSet;
21 + use std::convert::Infallible;
12 22 use std::net::SocketAddr;
13 23 use std::sync::Arc;
14 24
15 25 use axum::{
16 26 Json, Router,
17 27 extract::State,
18 - http::StatusCode,
19 - response::IntoResponse,
28 + http::{HeaderMap, StatusCode},
29 + response::{
30 + IntoResponse,
31 + sse::{Event, KeepAlive, Sse},
32 + },
20 33 routing::post,
21 34 };
22 35 use serde_json::{Value, json};
23 36 use tokio::net::TcpListener;
24 37 use tokio::task::JoinHandle;
38 + use tokio_stream::wrappers::UnboundedReceiverStream;
39 + use tokio_stream::{Stream, StreamExt};
25 40
26 41 use crate::error::{Error, Result};
42 + use crate::resource::{ReadResourceParams, ResourceRegistry, read_result};
27 43 use crate::tool::{SurfaceProjection, ToolRegistry};
28 44
29 45 pub mod protocol;
46 + mod session;
30 47 pub mod stdio;
31 48
32 49 use protocol::{
33 - CallToolParams, InitializeResult, JsonRpcRequest, JsonRpcResponse, MCP_PROTOCOL_VERSION,
34 - ServerCapabilities, ServerInfo, ToolsCapability, codes,
50 + CallToolParams, CancelledParams, InitializeResult, JsonRpcRequest, JsonRpcResponse,
51 + MCP_PROTOCOL_VERSION, ResourcesCapability, ServerCapabilities, ServerInfo, SubscribeParams,
52 + ToolsCapability, codes,
35 53 };
54 + use session::{Session, Sessions, resource_updated_notification};
55 +
56 + /// Header carrying the MCP session id, per the Streamable HTTP transport.
57 + const SESSION_HEADER: &str = "mcp-session-id";
36 58
37 59 pub struct ServeConfig {
38 60 pub projection: SurfaceProjection,
39 61 pub grants: Option<HashSet<String>>,
62 + /// Resource surface exposed alongside the tools. `None` means the server
63 + /// advertises no `resources` capability and rejects `resources/*` methods.
64 + pub resources: Option<ResourceRegistry>,
40 65 }
41 66
42 67 impl Default for ServeConfig {
@@ -44,6 +69,7 @@ impl Default for ServeConfig {
44 69 Self {
45 70 projection: SurfaceProjection::Full,
46 71 grants: Some(HashSet::new()),
72 + resources: None,
47 73 }
48 74 }
49 75 }
@@ -56,6 +82,7 @@ impl ServeConfig {
56 82 Self {
57 83 projection: SurfaceProjection::Full,
58 84 grants: Some(ids.into_iter().map(Into::into).collect()),
85 + resources: None,
59 86 }
60 87 }
61 88
@@ -63,6 +90,13 @@ impl ServeConfig {
63 90 self.projection = projection;
64 91 self
65 92 }
93 +
94 + /// Attach a resource registry, enabling `resources/list` and
95 + /// `resources/read` and advertising the `resources` capability.
96 + pub fn with_resources(mut self, resources: ResourceRegistry) -> Self {
97 + self.resources = Some(resources);
98 + self
99 + }
66 100 }
67 101
68 102 /// Handle to a running MCP server.
@@ -72,6 +106,7 @@ impl ServeConfig {
72 106 pub struct BoundServer {
73 107 pub addr: SocketAddr,
74 108 handle: JoinHandle<std::io::Result<()>>,
109 + sessions: Sessions,
75 110 }
76 111
77 112 impl BoundServer {
@@ -80,6 +115,30 @@ impl BoundServer {
80 115 format!("http://{}/mcp", self.addr)
81 116 }
82 117
118 + /// Cancel every in-flight `tools/call` across all sessions, pushing a
119 + /// `notifications/cancelled` to each affected client. The app-initiated
120 + /// path — a human taking a buffer back in deox. Returns the count
121 + /// cancelled. For v0 (one agent = one session) this is the whole cancel
122 + /// surface; [`cancel_session`](Self::cancel_session) targets one client
123 + /// once an app tracks session ids.
124 + pub fn cancel_all(&self, reason: &str) -> usize {
125 + self.sessions.cancel_all(reason)
126 + }
127 +
128 + /// Cancel in-flight `tools/call`s for one session by its `Mcp-Session-Id`.
129 + /// Returns the count cancelled (0 if the session is unknown or idle).
130 + pub fn cancel_session(&self, session_id: &str, reason: &str) -> usize {
131 + self.sessions
132 + .get(session_id)
133 + .map(|s| s.cancel_all(reason))
134 + .unwrap_or(0)
135 + }
136 +
137 + /// Ids of the currently-connected sessions.
138 + pub fn sessions(&self) -> Vec<String> {
139 + self.sessions.ids()
140 + }
141 +
83 142 pub fn shutdown(self) {
84 143 self.handle.abort();
85 144 }
@@ -103,6 +162,7 @@ impl BoundServer {
103 162 struct AppState {
104 163 registry: ToolRegistry,
105 164 config: Arc<ServeConfig>,
165 + sessions: Sessions,
106 166 }
107 167
108 168 /// Start an MCP server. Returns once the listener is bound; the server runs
@@ -115,22 +175,36 @@ pub async fn serve(
115 175 let listener = TcpListener::bind(bind).await?;
116 176 let addr = listener.local_addr()?;
117 177
178 + let sessions = Sessions::default();
118 179 let state = AppState {
119 180 registry,
120 181 config: Arc::new(config),
182 + sessions: sessions.clone(),
121 183 };
122 184
123 185 let app = Router::new()
124 - .route("/mcp", post(handle_rpc))
186 + .route("/mcp", post(handle_rpc).get(handle_sse).delete(handle_delete))
125 187 .with_state(state);
126 188
127 189 let handle = tokio::spawn(async move { axum::serve(listener, app).await });
128 190
129 191 tracing::info!(%addr, "kberg server listening");
130 - Ok(BoundServer { addr, handle })
192 + Ok(BoundServer { addr, handle, sessions })
193 + }
194 +
195 + /// Look up the session named by the `Mcp-Session-Id` header, if any.
196 + fn session_from_headers(state: &AppState, headers: &HeaderMap) -> Option<Arc<Session>> {
197 + headers
198 + .get(SESSION_HEADER)
199 + .and_then(|h| h.to_str().ok())
200 + .and_then(|id| state.sessions.get(id))
131 201 }
132 202
133 - async fn handle_rpc(State(state): State<AppState>, Json(req): Json<Value>) -> impl IntoResponse {
203 + async fn handle_rpc(
204 + State(state): State<AppState>,
205 + headers: HeaderMap,
206 + Json(req): Json<Value>,
207 + ) -> impl IntoResponse {
134 208 let req: JsonRpcRequest = match serde_json::from_value(req) {
135 209 Ok(r) => r,
136 210 Err(e) => {
@@ -146,44 +220,158 @@ async fn handle_rpc(State(state): State<AppState>, Json(req): Json<Value>) -> im
146 220 }
147 221 };
148 222
149 - match dispatch(
223 + // `initialize` opens a session; every other method rides an existing one
224 + // named by the header (absent for a bare client, e.g. over stdio-style use).
225 + let is_initialize = req.method == "initialize";
226 + let session = session_from_headers(&state, &headers);
227 +
228 + let dispatched = dispatch(
150 229 &state.registry,
230 + state.config.resources.as_ref(),
151 231 state.config.projection,
152 232 state.config.grants.as_ref(),
233 + session.as_deref(),
153 234 req,
154 235 )
155 - .await
156 - {
236 + .await;
237 +
238 + match dispatched {
157 239 // Notifications (no `id`) get a 202 with no body.
158 240 None => StatusCode::ACCEPTED.into_response(),
159 - Some(response) => (StatusCode::OK, Json(response)).into_response(),
241 + Some(response) => {
242 + let mut http = (StatusCode::OK, Json(response)).into_response();
243 + if is_initialize {
244 + let new_session = state.sessions.create();
245 + if let Ok(value) = new_session.id.parse() {
246 + http.headers_mut().insert(SESSION_HEADER, value);
247 + }
248 + }
249 + http
250 + }
160 251 }
161 252 }
162 253
254 + /// GET `/mcp` — open the server→client SSE stream for a session. This is the
255 + /// channel `notifications/resources/updated` and `notifications/cancelled`
256 + /// travel down. Requires a valid `Mcp-Session-Id`.
257 + async fn handle_sse(
258 + State(state): State<AppState>,
259 + headers: HeaderMap,
260 + ) -> axum::response::Response {
261 + let Some(session) = session_from_headers(&state, &headers) else {
262 + return (StatusCode::BAD_REQUEST, "missing or unknown Mcp-Session-Id").into_response();
263 + };
264 +
265 + let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
266 + session.attach_out(tx.clone());
267 +
268 + // Forward resource updates this session is subscribed to. The task lives as
269 + // long as the SSE stream: once the client disconnects, `rx` drops and the
270 + // `tx.send` below fails, ending the loop.
271 + if let Some(resources) = &state.config.resources {
272 + let mut updates = resources.subscribe_updates();
273 + let session = session.clone();
274 + tokio::spawn(async move {
275 + loop {
276 + match updates.recv().await {
277 + Ok(uri) => {
278 + if session.is_subscribed(&uri)
279 + && tx.send(resource_updated_notification(&uri)).is_err()
280 + {
281 + break; // client gone
282 + }
283 + }
284 + // A slow session that lagged past the buffer drops those
285 + // updates and keeps streaming; it re-reads on next use.
286 + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
287 + Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
288 + }
289 + }
290 + });
291 + }
292 +
293 + let stream = UnboundedReceiverStream::new(rx).map(|note| {
294 + Ok::<Event, Infallible>(Event::default().data(note.to_string()))
295 + });
296 + sse_response(stream)
297 + }
298 +
299 + fn sse_response<S>(stream: S) -> axum::response::Response
300 + where
301 + S: Stream<Item = std::result::Result<Event, Infallible>> + Send + 'static,
302 + {
303 + Sse::new(stream).keep_alive(KeepAlive::default()).into_response()
304 + }
305 +
306 + /// DELETE `/mcp` — end a session (the client is done). Idempotent.
307 + async fn handle_delete(State(state): State<AppState>, headers: HeaderMap) -> impl IntoResponse {
308 + if let Some(id) = headers.get(SESSION_HEADER).and_then(|h| h.to_str().ok()) {
309 + state.sessions.remove(id);
310 + }
311 + StatusCode::NO_CONTENT
312 + }
313 +
163 314 /// Transport-agnostic JSON-RPC dispatch. Handles `initialize`, `tools/list`,
164 315 /// and `tools/call` against a registry. Returns `None` for a notification
165 316 /// (a request with no `id`), which carries no response. Shared by the HTTP
166 317 /// handler and the [`stdio`] transport.
167 318 pub(crate) async fn dispatch(
168 319 registry: &ToolRegistry,
320 + resources: Option<&ResourceRegistry>,
169 321 projection: SurfaceProjection,
170 322 grants: Option<&HashSet<String>>,
323 + session: Option<&Session>,
171 324 req: JsonRpcRequest,
172 325 ) -> Option<JsonRpcResponse> {
173 - let id = req.id.clone()?;
326 + // Notifications (no `id`) carry no response, but some carry a side effect.
327 + let Some(id) = req.id.clone() else {
328 + if req.method == "notifications/cancelled"
329 + && let Some(session) = session
330 + && let Ok(params) = serde_json::from_value::<CancelledParams>(req.params.unwrap_or(Value::Null))
331 + {
332 + session.cancel_request(&params.request_id.to_string());
333 + }
334 + return None;
335 + };
336 +
174 337 let response = match req.method.as_str() {
175 - "initialize" => JsonRpcResponse::ok(id, initialize_result()),
338 + "initialize" => JsonRpcResponse::ok(id, initialize_result(resources.is_some())),
176 339 "tools/list" => {
177 340 let specs = registry.specs_projected(projection);
178 341 JsonRpcResponse::ok(id, json!({ "tools": specs }))
179 342 }
180 343 "tools/call" => match serde_json::from_value::<CallToolParams>(req.params.unwrap_or(Value::Null)) {
181 - Ok(params) => match registry.call(&params.name, params.arguments, grants).await {
182 - Ok(r) => JsonRpcResponse::ok(id, serde_json::to_value(r).unwrap()),
183 - Err(e) => rpc_err_from(id, e),
184 - },
344 + Ok(params) => {
345 + let call = registry.call(&params.name, params.arguments, grants);
346 + // With a session, run so a human/client can cancel mid-flight.
347 + let result = match session {
348 + Some(session) => session.run_cancellable(&id, call).await,
349 + None => call.await,
350 + };
351 + match result {
352 + Ok(r) => JsonRpcResponse::ok(id, serde_json::to_value(r).unwrap()),
353 + Err(e) => rpc_err_from(id, e),
354 + }
355 + }
185 356 Err(e) => JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
186 357 },
358 + "resources/list" => match resources {
359 + Some(res) => JsonRpcResponse::ok(id, json!({ "resources": res.list() })),
360 + None => resources_unsupported(id),
361 + },
362 + "resources/read" => match resources {
363 + Some(res) => match serde_json::from_value::<ReadResourceParams>(req.params.unwrap_or(Value::Null)) {
364 + Ok(params) => match res.read(&params.uri).await {
365 + Ok(contents) => JsonRpcResponse::ok(id, read_result(contents)),
366 + Err(e) => rpc_err_from(id, e),
367 + },
368 + Err(e) => JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
369 + },
370 + None => resources_unsupported(id),
371 + },
372 + method @ ("resources/subscribe" | "resources/unsubscribe") => {
373 + dispatch_subscription(method, resources, session, id, req.params)
374 + }
187 375 other => JsonRpcResponse::err(
188 376 id,
189 377 codes::METHOD_NOT_FOUND,
@@ -193,9 +381,57 @@ pub(crate) async fn dispatch(
193 381 Some(response)
194 382 }
195 383
384 + /// `resources/subscribe` / `resources/unsubscribe`. Both need a resource
385 + /// registry (to validate the uri exists) and a session (to hold the
386 + /// subscription) — the latter means these are HTTP-only.
387 + fn dispatch_subscription(
388 + method: &str,
389 + resources: Option<&ResourceRegistry>,
390 + session: Option<&Session>,
391 + id: Value,
392 + params: Option<Value>,
393 + ) -> JsonRpcResponse {
394 + let Some(resources) = resources else {
395 + return resources_unsupported(id);
396 + };
397 + let Some(session) = session else {
398 + return JsonRpcResponse::err(
399 + id,
400 + codes::INVALID_PARAMS,
401 + "resource subscriptions require a session (HTTP transport)",
402 + );
403 + };
404 + let params: SubscribeParams = match serde_json::from_value(params.unwrap_or(Value::Null)) {
405 + Ok(p) => p,
406 + Err(e) => return JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
407 + };
408 + if method == "resources/subscribe" {
409 + if !resources.contains(&params.uri) {
410 + return rpc_err_from(id, Error::ResourceNotFound(params.uri));
411 + }
412 + session.subscribe(params.uri);
413 + } else {
414 + session.unsubscribe(&params.uri);
415 + }
416 + // MCP subscribe/unsubscribe return an empty result on success.
417 + JsonRpcResponse::ok(id, json!({}))
418 + }
419 +
420 + /// A `resources/*` method arrived at a server with no resource registry
421 + /// configured. It is a method the server does not implement, not a bad request.
422 + fn resources_unsupported(id: Value) -> JsonRpcResponse {
423 + JsonRpcResponse::err(
424 + id,
425 + codes::METHOD_NOT_FOUND,
426 + "server exposes no resources",
427 + )
428 + }
429 +
196 430 fn rpc_err_from(id: Value, e: Error) -> JsonRpcResponse {
197 431 let (code, message) = match &e {
198 432 Error::ToolNotFound(_) => (codes::METHOD_NOT_FOUND, e.to_string()),
433 + Error::ResourceNotFound(_) => (codes::RESOURCE_NOT_FOUND, e.to_string()),
434 + Error::Cancelled(_) => (codes::REQUEST_CANCELLED, e.to_string()),
199 435 Error::InvalidArgs { .. } => (codes::INVALID_PARAMS, e.to_string()),
200 436 Error::CapabilityDenied { .. } | Error::Refused { .. } => (codes::INVALID_PARAMS, e.to_string()),
201 437 _ => (codes::INTERNAL_ERROR, e.to_string()),
@@ -203,11 +439,14 @@ fn rpc_err_from(id: Value, e: Error) -> JsonRpcResponse {
203 439 JsonRpcResponse::err(id, code, message)
204 440 }
205 441
206 - fn initialize_result() -> Value {
442 + fn initialize_result(resources_enabled: bool) -> Value {
207 443 serde_json::to_value(InitializeResult {
208 444 protocol_version: MCP_PROTOCOL_VERSION,
209 445 capabilities: ServerCapabilities {
210 446 tools: ToolsCapability { list_changed: false },
447 + // Over HTTP, resources support subscribe + server-push updates.
448 + resources: resources_enabled
449 + .then_some(ResourcesCapability { subscribe: true, list_changed: false }),
211 450 },
212 451 server_info: ServerInfo {
213 452 name: "kberg".to_string(),
@@ -220,6 +459,7 @@ fn initialize_result() -> Value {
220 459 #[cfg(test)]
221 460 mod tests {
222 461 use super::*;
462 + use crate::resource::{ResourceContents, ResourceDescriptor};
223 463 use crate::tool::{Tool, ToolCallResult, ToolKind};
224 464 use async_trait::async_trait;
225 465
@@ -264,13 +504,15 @@ mod tests {
264 504 let r = registry();
265 505 let grants = HashSet::new();
266 506
267 - let init = dispatch(&r, SurfaceProjection::Full, Some(&grants), request(Some(json!(1)), "initialize", Value::Null))
507 + let init = dispatch(&r, None, SurfaceProjection::Full, Some(&grants), None, request(Some(json!(1)), "initialize", Value::Null))
268 508 .await
269 509 .unwrap();
270 510 let v = serde_json::to_value(init).unwrap();
271 511 assert_eq!(v["result"]["serverInfo"]["name"], "kberg");
512 + // No resource registry → no `resources` capability advertised.
513 + assert!(v["result"]["capabilities"]["resources"].is_null());
272 514
273 - let list = dispatch(&r, SurfaceProjection::Full, Some(&grants), request(Some(json!(2)), "tools/list", Value::Null))
515 + let list = dispatch(&r, None, SurfaceProjection::Full, Some(&grants), None, request(Some(json!(2)), "tools/list", Value::Null))
274 516 .await
275 517 .unwrap();
276 518 let v = serde_json::to_value(list).unwrap();
@@ -278,8 +520,10 @@ mod tests {
278 520
279 521 let call = dispatch(
280 522 &r,
523 + None,
281 524 SurfaceProjection::Full,
282 525 Some(&grants),
526 + None,
283 527 request(Some(json!(3)), "tools/call", json!({ "name": "echo", "arguments": {} })),
284 528 )
285 529 .await
@@ -291,14 +535,88 @@ mod tests {
291 535 #[tokio::test]
292 536 async fn dispatch_notification_has_no_response() {
293 537 let r = registry();
294 - let resp = dispatch(&r, SurfaceProjection::Full, None, request(None, "initialize", Value::Null)).await;
538 + let resp = dispatch(&r, None, SurfaceProjection::Full, None, None, request(None, "initialize", Value::Null)).await;
295 539 assert!(resp.is_none());
296 540 }
297 541
298 542 #[tokio::test]
299 543 async fn dispatch_unknown_method_errors() {
300 544 let r = registry();
301 - let resp = dispatch(&r, SurfaceProjection::Full, None, request(Some(json!(9)), "bogus", Value::Null))
545 + let resp = dispatch(&r, None, SurfaceProjection::Full, None, None, request(Some(json!(9)), "bogus", Value::Null))
546 + .await
547 + .unwrap();
548 + let v = serde_json::to_value(resp).unwrap();
549 + assert_eq!(v["error"]["code"], codes::METHOD_NOT_FOUND);
550 + }
551 +
552 + fn resources() -> ResourceRegistry {
553 + let res = ResourceRegistry::new();
554 + res.register(
555 + ResourceDescriptor::new("buffer://main.rs", "main.rs").with_mime_type("text/x-rust"),
556 + |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
557 + );
558 + res
559 + }
560 +
561 + #[tokio::test]
562 + async fn dispatch_initialize_advertises_resources_when_configured() {
563 + let r = registry();
564 + let res = resources();
565 + let init = dispatch(&r, Some(&res), SurfaceProjection::Full, None, None, request(Some(json!(1)), "initialize", Value::Null))
566 + .await
567 + .unwrap();
568 + let v = serde_json::to_value(init).unwrap();
569 + assert_eq!(v["result"]["capabilities"]["resources"]["subscribe"], true);
570 + }
571 +
572 + #[tokio::test]
573 + async fn dispatch_resources_list_and_read() {
574 + let r = registry();
575 + let res = resources();
Lines truncated
M shared/kberg/src/server/protocol.rs +33
@@ -65,6 +65,11 @@ pub mod codes {
65 65 pub const METHOD_NOT_FOUND: i32 = -32601;
66 66 pub const INVALID_PARAMS: i32 = -32602;
67 67 pub const INTERNAL_ERROR: i32 = -32603;
68 + /// MCP-defined: a `resources/read` (or subscribe) named an unknown uri.
69 + pub const RESOURCE_NOT_FOUND: i32 = -32002;
70 + /// An in-flight request was cancelled before it completed. Same code LSP
71 + /// uses for `ContentModified`/cancellation.
72 + pub const REQUEST_CANCELLED: i32 = -32800;
68 73 }
69 74
70 75 #[derive(Debug, Serialize)]
@@ -79,6 +84,10 @@ pub struct InitializeResult {
79 84 #[derive(Debug, Serialize)]
80 85 pub struct ServerCapabilities {
81 86 pub tools: ToolsCapability,
87 + /// Only advertised when the server was configured with a resource
88 + /// registry. `subscribe` stays `false` until the push layer lands.
89 + #[serde(skip_serializing_if = "Option::is_none")]
90 + pub resources: Option<ResourcesCapability>,
82 91 }
83 92
84 93 #[derive(Debug, Serialize)]
@@ -88,6 +97,13 @@ pub struct ToolsCapability {
88 97 }
89 98
90 99 #[derive(Debug, Serialize)]
100 + pub struct ResourcesCapability {
101 + pub subscribe: bool,
102 + #[serde(rename = "listChanged")]
103 + pub list_changed: bool,
104 + }
105 +
106 + #[derive(Debug, Serialize)]
91 107 pub struct ServerInfo {
92 108 pub name: String,
93 109 pub version: String,
@@ -99,3 +115,20 @@ pub struct CallToolParams {
99 115 #[serde(default)]
100 116 pub arguments: Value,
101 117 }
118 +
119 + /// `resources/subscribe` and `resources/unsubscribe` param object.
120 + #[derive(Debug, Deserialize)]
121 + pub struct SubscribeParams {
122 + pub uri: String,
123 + }
124 +
125 + /// `notifications/cancelled` param object (client → server). `requestId` is the
126 + /// id of the request being cancelled; `reason` is optional free text.
127 + #[derive(Debug, Deserialize)]
128 + pub struct CancelledParams {
129 + #[serde(rename = "requestId")]
130 + pub request_id: Value,
131 + #[serde(default)]
132 + #[allow(dead_code)]
133 + pub reason: Option<String>,
134 + }
A shared/kberg/src/server/session.rs +255
@@ -0,0 +1,255 @@
1 + //! Per-connection session state for the Streamable HTTP transport.
2 + //!
3 + //! A session is created on `initialize` (the server returns its id in the
4 + //! `Mcp-Session-Id` header) and referenced by that id on every later request.
5 + //! It holds the three things the push layer needs that a stateless request
6 + //! cannot:
7 + //!
8 + //! - the set of resource uris this client has **subscribed** to,
9 + //! - the **out** channel to this client's SSE stream (present while a GET
10 + //! stream is open), down which the server pushes notifications,
11 + //! - the **in-flight** `tools/call` requests, each with a trigger to cancel it.
12 + //!
13 + //! stdio has no sessions (one implicit peer, no `Mcp-Session-Id`), so its
14 + //! dispatch passes `None` and the session-scoped methods are unavailable there.
15 +
16 + use std::collections::{HashMap, HashSet};
17 + use std::future::Future;
18 + use std::sync::{Arc, Mutex};
19 +
20 + use serde_json::Value;
21 + use tokio::sync::{mpsc, oneshot};
22 +
23 + use crate::error::{Error, Result};
24 + use crate::tool::ToolCallResult;
25 +
26 + /// JSON-RPC notification value pushed to a session's SSE stream.
27 + pub(crate) type OutSender = mpsc::UnboundedSender<Value>;
28 +
29 + struct Inflight {
30 + /// Original JSON-RPC id, preserved with its type (number stays a number)
31 + /// so a `notifications/cancelled` echoes the exact id the client sent.
32 + id: Value,
33 + cancel: oneshot::Sender<()>,
34 + }
35 +
36 + /// One client connection's mutable state. Shared (`Arc`) between the POST
37 + /// handler that mutates it and the GET SSE handler that drains its out channel.
38 + pub(crate) struct Session {
39 + pub id: String,
40 + subscriptions: Mutex<HashSet<String>>,
41 + out: Mutex<Option<OutSender>>,
42 + inflight: Mutex<HashMap<String, Inflight>>,
43 + }
44 +
45 + impl Session {
46 + fn new(id: String) -> Self {
47 + Self {
48 + id,
49 + subscriptions: Mutex::new(HashSet::new()),
50 + out: Mutex::new(None),
51 + inflight: Mutex::new(HashMap::new()),
52 + }
53 + }
54 +
55 + pub fn subscribe(&self, uri: impl Into<String>) {
56 + self.subscriptions.lock().unwrap().insert(uri.into());
57 + }
58 +
59 + /// Remove a subscription. Returns whether it was present.
60 + pub fn unsubscribe(&self, uri: &str) -> bool {
61 + self.subscriptions.lock().unwrap().remove(uri)
62 + }
63 +
64 + pub fn is_subscribed(&self, uri: &str) -> bool {
65 + self.subscriptions.lock().unwrap().contains(uri)
66 + }
67 +
68 + /// Attach the SSE out channel (called when the client opens GET `/mcp`).
69 + /// Replaces any prior channel — a reconnect supersedes the stale stream.
70 + pub fn attach_out(&self, tx: OutSender) {
71 + *self.out.lock().unwrap() = Some(tx);
72 + }
73 +
74 + /// Push a notification to the SSE stream if one is connected.
75 + pub fn push(&self, notification: Value) {
76 + if let Some(tx) = self.out.lock().unwrap().as_ref() {
77 + let _ = tx.send(notification);
78 + }
79 + }
80 +
81 + /// Run a `tools/call` future so it can be cancelled while in flight. The
82 + /// call is registered under `id` until it completes; a concurrent
83 + /// [`cancel_request`](Self::cancel_request) or [`cancel_all`](Self::cancel_all)
84 + /// aborts it with [`Error::Cancelled`].
85 + pub async fn run_cancellable<F>(&self, id: &Value, fut: F) -> Result<ToolCallResult>
86 + where
87 + F: Future<Output = Result<ToolCallResult>>,
88 + {
89 + let key = id.to_string();
90 + let (tx, rx) = oneshot::channel();
91 + self.inflight
92 + .lock()
93 + .unwrap()
94 + .insert(key.clone(), Inflight { id: id.clone(), cancel: tx });
95 +
96 + let outcome = tokio::select! {
97 + biased;
98 + _ = rx => Err(Error::Cancelled(key.clone())),
99 + result = fut => result,
100 + };
101 +
102 + self.inflight.lock().unwrap().remove(&key);
103 + outcome
104 + }
105 +
106 + /// Cancel one in-flight request by its (stringified) id. Triggers the
107 + /// select in [`run_cancellable`](Self::run_cancellable). No-op if the id is
108 + /// not in flight (already finished or never existed).
109 + pub fn cancel_request(&self, id_key: &str) {
110 + if let Some(entry) = self.inflight.lock().unwrap().remove(id_key) {
111 + let _ = entry.cancel.send(());
112 + }
113 + }
114 +
115 + /// Cancel every in-flight request for this session and tell the client,
116 + /// pushing a `notifications/cancelled` for each. This is the app-initiated
117 + /// path — a human taking a buffer back. Returns the count cancelled.
118 + pub fn cancel_all(&self, reason: &str) -> usize {
119 + let drained: Vec<Inflight> = {
120 + let mut map = self.inflight.lock().unwrap();
121 + map.drain().map(|(_, v)| v).collect()
122 + };
123 + let n = drained.len();
124 + for entry in drained {
125 + let _ = entry.cancel.send(());
126 + self.push(cancelled_notification(&entry.id, reason));
127 + }
128 + n
129 + }
130 + }
131 +
132 + /// The set of live sessions, shared by the server's handlers and its
133 + /// [`BoundServer`](super::BoundServer) cancel handle.
134 + #[derive(Clone, Default)]
135 + pub(crate) struct Sessions {
136 + map: Arc<Mutex<HashMap<String, Arc<Session>>>>,
137 + }
138 +
139 + impl Sessions {
140 + /// Create a fresh session with a random id and return it.
141 + pub fn create(&self) -> Arc<Session> {
142 + let id = uuid::Uuid::new_v4().to_string();
143 + let session = Arc::new(Session::new(id.clone()));
144 + self.map.lock().unwrap().insert(id, session.clone());
145 + session
146 + }
147 +
148 + pub fn get(&self, id: &str) -> Option<Arc<Session>> {
149 + self.map.lock().unwrap().get(id).cloned()
150 + }
151 +
152 + /// Remove a session (client sent DELETE `/mcp`). Returns whether present.
153 + pub fn remove(&self, id: &str) -> bool {
154 + self.map.lock().unwrap().remove(id).is_some()
155 + }
156 +
157 + pub fn ids(&self) -> Vec<String> {
158 + self.map.lock().unwrap().keys().cloned().collect()
159 + }
160 +
161 + /// Cancel in-flight ops across every session. Returns the total cancelled.
162 + pub fn cancel_all(&self, reason: &str) -> usize {
163 + let sessions: Vec<Arc<Session>> = self.map.lock().unwrap().values().cloned().collect();
164 + sessions.iter().map(|s| s.cancel_all(reason)).sum()
165 + }
166 + }
167 +
168 + /// Build a server→client `notifications/cancelled` for `id`.
169 + pub(crate) fn cancelled_notification(id: &Value, reason: &str) -> Value {
170 + serde_json::json!({
171 + "jsonrpc": "2.0",
172 + "method": "notifications/cancelled",
173 + "params": { "requestId": id, "reason": reason },
174 + })
175 + }
176 +
177 + /// Build a server→client `notifications/resources/updated` for `uri`.
178 + pub(crate) fn resource_updated_notification(uri: &str) -> Value {
179 + serde_json::json!({
180 + "jsonrpc": "2.0",
181 + "method": "notifications/resources/updated",
182 + "params": { "uri": uri },
183 + })
184 + }
185 +
186 + #[cfg(test)]
187 + mod tests {
188 + use super::*;
189 +
190 + #[tokio::test]
191 + async fn run_cancellable_returns_result_when_not_cancelled() {
192 + let s = Session::new("s1".into());
193 + let out = s
194 + .run_cancellable(&serde_json::json!(1), async { Ok(ToolCallResult::text("done")) })
195 + .await
196 + .unwrap();
197 + assert_eq!(out.content.len(), 1);
198 + }
199 +
200 + #[tokio::test]
201 + async fn cancel_request_aborts_an_in_flight_call() {
202 + let s = Arc::new(Session::new("s1".into()));
203 + let id = serde_json::json!(7);
204 + let s2 = s.clone();
205 + // A call that would never finish on its own.
206 + let never = async { std::future::pending::<Result<ToolCallResult>>().await };
207 + let handle = tokio::spawn(async move { s2.run_cancellable(&id, never).await });
208 + // Let the call register itself, then cancel it.
209 + tokio::task::yield_now().await;
210 + s.cancel_request("7");
211 + let outcome = handle.await.unwrap();
212 + assert!(matches!(outcome, Err(Error::Cancelled(_))));
213 + }
214 +
215 + #[tokio::test]
216 + async fn cancel_all_notifies_the_client_with_original_id_type() {
217 + let s = Arc::new(Session::new("s1".into()));
218 + let (tx, mut rx) = mpsc::unbounded_channel();
219 + s.attach_out(tx);
220 +
221 + let id = serde_json::json!(42); // numeric id
222 + let s2 = s.clone();
223 + let never = async { std::future::pending::<Result<ToolCallResult>>().await };
224 + let handle = tokio::spawn(async move { s2.run_cancellable(&id, never).await });
225 + tokio::task::yield_now().await;
226 +
227 + assert_eq!(s.cancel_all("human took the buffer back"), 1);
228 + assert!(matches!(handle.await.unwrap(), Err(Error::Cancelled(_))));
229 +
230 + let note = rx.recv().await.unwrap();
231 + assert_eq!(note["method"], "notifications/cancelled");
232 + assert_eq!(note["params"]["requestId"], 42); // stayed a JSON number
233 + assert_eq!(note["params"]["reason"], "human took the buffer back");
234 + }
235 +
236 + #[test]
237 + fn subscribe_unsubscribe_tracks_membership() {
238 + let s = Session::new("s1".into());
239 + assert!(!s.is_subscribed("buffer://a"));
240 + s.subscribe("buffer://a");
241 + assert!(s.is_subscribed("buffer://a"));
242 + assert!(s.unsubscribe("buffer://a"));
243 + assert!(!s.unsubscribe("buffer://a"));
244 + }
245 +
246 + #[test]
247 + fn sessions_create_get_remove() {
248 + let sessions = Sessions::default();
249 + let s = sessions.create();
250 + assert!(sessions.get(&s.id).is_some());
251 + assert_eq!(sessions.ids().len(), 1);
252 + assert!(sessions.remove(&s.id));
253 + assert!(sessions.get(&s.id).is_none());
254 + }
255 + }
M shared/kberg/src/server/stdio.rs +13 -1
@@ -29,7 +29,19 @@ pub async fn serve(registry: ToolRegistry, config: ServeConfig) -> Result<()> {
29 29 continue;
30 30 }
31 31 let response = match serde_json::from_str(line) {
32 - Ok(req) => super::dispatch(&registry, config.projection, config.grants.as_ref(), req).await,
32 + Ok(req) => {
33 + // stdio has no session, so subscriptions and cancellation are
34 + // unavailable here; the resource read surface still works.
35 + super::dispatch(
36 + &registry,
37 + config.resources.as_ref(),
38 + config.projection,
39 + config.grants.as_ref(),
40 + None,
41 + req,
42 + )
43 + .await
44 + }
33 45 Err(e) => Some(JsonRpcResponse::err(
34 46 Value::Null,
35 47 codes::PARSE_ERROR,
M shared/kberg/tests/server.rs +236 -2
@@ -6,12 +6,13 @@
6 6 #![cfg(feature = "server")]
7 7
8 8 use std::collections::HashSet;
9 + use std::time::Duration;
9 10
10 11 use async_trait::async_trait;
11 12 use kberg::server::{self, ServeConfig};
12 13 use kberg::{
13 - Refusal, Result, SurfaceProjection, Tool, ToolCallResult, ToolKind, ToolRegistry,
14 - WriteCapability,
14 + Refusal, ResourceContents, ResourceDescriptor, ResourceRegistry, Result, SurfaceProjection,
15 + Tool, ToolCallResult, ToolKind, ToolRegistry, WriteCapability,
15 16 };
16 17 use serde_json::{Value, json};
17 18
@@ -67,6 +68,28 @@ fn registry() -> ToolRegistry {
67 68 r
68 69 }
69 70
71 + /// A read tool that never returns on its own — only cancellation ends it.
72 + struct Hang;
73 + #[async_trait]
74 + impl Tool for Hang {
75 + fn name(&self) -> &str {
76 + "hang"
77 + }
78 + fn description(&self) -> &str {
79 + "read-only; blocks until cancelled"
80 + }
81 + fn kind(&self) -> ToolKind {
82 + ToolKind::Read
83 + }
84 + fn input_schema(&self) -> Value {
85 + json!({ "type": "object", "properties": {} })
86 + }
87 + async fn call(&self, _args: Value) -> Result<ToolCallResult> {
88 + std::future::pending::<()>().await;
89 + unreachable!()
90 + }
91 + }
92 +
70 93 async fn rpc(client: &reqwest::Client, url: &str, method: &str, params: Value) -> Value {
71 94 client
72 95 .post(url)
@@ -79,6 +102,66 @@ async fn rpc(client: &reqwest::Client, url: &str, method: &str, params: Value) -
79 102 .unwrap()
80 103 }
81 104
105 + /// POST a request under a session id and return the parsed JSON-RPC response.
106 + async fn rpc_in_session(
107 + client: &reqwest::Client,
108 + url: &str,
109 + session: &str,
110 + id: i64,
111 + method: &str,
112 + params: Value,
113 + ) -> Value {
114 + client
115 + .post(url)
116 + .header("Mcp-Session-Id", session)
117 + .json(&json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }))
118 + .send()
119 + .await
120 + .unwrap()
121 + .json()
122 + .await
123 + .unwrap()
124 + }
125 +
126 + /// Initialize and return the server-assigned `Mcp-Session-Id`.
127 + async fn open_session(client: &reqwest::Client, url: &str) -> String {
128 + let resp = client
129 + .post(url)
130 + .json(&json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize" }))
131 + .send()
132 + .await
133 + .unwrap();
134 + resp.headers()
135 + .get("Mcp-Session-Id")
136 + .expect("initialize returns a session id")
137 + .to_str()
138 + .unwrap()
139 + .to_string()
140 + }
141 +
142 + /// Read the next `data:` JSON payload off an SSE response, buffering across
143 + /// chunks. Returns `None` at end of stream.
144 + async fn next_sse_json(resp: &mut reqwest::Response, buf: &mut String) -> Option<Value> {
145 + loop {
146 + if let Some(pos) = buf.find("data:")
147 + && let Some(nl) = buf[pos..].find('\n')
148 + {
149 + let line = buf[pos + 5..pos + nl].trim().to_string();
150 + *buf = buf[pos + nl + 1..].to_string();
151 + if !line.is_empty()
152 + && let Ok(v) = serde_json::from_str::<Value>(&line)
153 + {
154 + return Some(v);
155 + }
156 + continue;
157 + }
158 + match resp.chunk().await.ok()? {
159 + Some(bytes) => buf.push_str(&String::from_utf8_lossy(&bytes)),
160 + None => return None,
161 + }
162 + }
163 + }
164 +
82 165 #[tokio::test]
83 166 async fn full_surface_round_trip() {
84 167 let bound = server::serve(
@@ -159,6 +242,49 @@ async fn ungranted_write_is_denied_over_the_wire() {
159 242 }
160 243
161 244 #[tokio::test]
245 + async fn resources_listed_read_and_removed_over_the_wire() {
246 + // A resource registry the app mutates at runtime, mirroring deox opening
247 + // and closing buffers while the server stays up.
248 + let resources = ResourceRegistry::new();
249 + resources.register(
250 + ResourceDescriptor::new("buffer://main.rs", "main.rs").with_mime_type("text/x-rust"),
251 + |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
252 + );
253 +
254 + let bound = server::serve(
255 + registry(),
256 + "127.0.0.1:0",
257 + ServeConfig::granting(Vec::<String>::new()).with_resources(resources.clone()),
258 + )
259 + .await
260 + .unwrap();
261 + let url = bound.url();
262 + let client = reqwest::Client::new();
263 +
264 + // initialize advertises the resources capability with subscribe support.
265 + let init = rpc(&client, &url, "initialize", Value::Null).await;
266 + assert_eq!(init["result"]["capabilities"]["resources"]["subscribe"], true);
267 +
268 + // resources/list
269 + let list = rpc(&client, &url, "resources/list", Value::Null).await;
270 + let items = list["result"]["resources"].as_array().unwrap();
271 + assert_eq!(items.len(), 1);
272 + assert_eq!(items[0]["uri"], "buffer://main.rs");
273 + assert_eq!(items[0]["mimeType"], "text/x-rust");
274 +
275 + // resources/read
276 + let read = rpc(&client, &url, "resources/read", json!({ "uri": "buffer://main.rs" })).await;
277 + assert_eq!(read["result"]["contents"][0]["text"], "fn main() {}");
278 +
279 + // Close the buffer at runtime; the running server observes it (shared Arc).
280 + assert!(resources.remove("buffer://main.rs"));
281 + let gone = rpc(&client, &url, "resources/read", json!({ "uri": "buffer://main.rs" })).await;
282 + assert_eq!(gone["error"]["code"], -32002); // RESOURCE_NOT_FOUND
283 +
284 + bound.shutdown();
285 + }
286 +
287 + #[tokio::test]
162 288 async fn compact_projection_hides_write_tools() {
163 289 let bound = server::serve(
164 290 registry(),
@@ -183,3 +309,111 @@ async fn compact_projection_hides_write_tools() {
183 309
184 310 bound.shutdown();
185 311 }
312 +
313 + #[tokio::test]
314 + async fn resource_update_is_pushed_over_sse() {
315 + // The deox flow: agent subscribes to a buffer, the human edits it, the
316 + // agent's SSE stream receives a `resources/updated` push.
317 + let resources = ResourceRegistry::new();
318 + resources.register(
319 + ResourceDescriptor::new("buffer://main.rs", "main.rs"),
320 + |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
321 + );
322 + let bound = server::serve(
323 + registry(),
324 + "127.0.0.1:0",
325 + ServeConfig::granting(Vec::<String>::new()).with_resources(resources.clone()),
326 + )
327 + .await
328 + .unwrap();
329 + let url = bound.url();
330 + let client = reqwest::Client::new();
331 +
332 + let session = open_session(&client, &url).await;
333 +
334 + // Open the server->client SSE stream (registers the update forwarder).
335 + let mut sse = client
336 + .get(&url)
337 + .header("Mcp-Session-Id", &session)
338 + .send()
339 + .await
340 + .unwrap();
341 + let mut buf = String::new();
342 +
343 + // Subscribe to the buffer.
344 + let sub = rpc_in_session(
345 + &client,
346 + &url,
347 + &session,
348 + 2,
349 + "resources/subscribe",
350 + json!({ "uri": "buffer://main.rs" }),
351 + )
352 + .await;
353 + assert!(sub["result"].is_object(), "subscribe returns an empty result: {sub}");
354 +
355 + // The app mutates the buffer and announces it.
356 + resources.notify_updated("buffer://main.rs");
357 +
358 + let note = tokio::time::timeout(Duration::from_secs(2), next_sse_json(&mut sse, &mut buf))
359 + .await
360 + .expect("an update arrives before the timeout")
361 + .expect("the stream yields a notification");
362 + assert_eq!(note["method"], "notifications/resources/updated");
363 + assert_eq!(note["params"]["uri"], "buffer://main.rs");
364 +
365 + bound.shutdown();
366 + }
367 +
368 + #[tokio::test]
369 + async fn app_cancel_aborts_call_and_notifies_over_sse() {
370 + // The deox cancel: a call is in flight, the human cancels, the server
371 + // aborts the call and tells the agent over SSE.
372 + let mut reg = ToolRegistry::new();
373 + reg.register(Hang);
374 + let bound = server::serve(reg, "127.0.0.1:0", ServeConfig::granting(Vec::<String>::new()))
375 + .await
376 + .unwrap();
377 + let url = bound.url();
378 + let client = reqwest::Client::new();
379 +
380 + let session = open_session(&client, &url).await;
381 + let mut sse = client
382 + .get(&url)
383 + .header("Mcp-Session-Id", &session)
384 + .send()
385 + .await
386 + .unwrap();
387 + let mut buf = String::new();
388 +
389 + // Fire a call that hangs until cancelled.
390 + let call = {
391 + let client = client.clone();
392 + let url = url.clone();
393 + let session = session.clone();
394 + tokio::spawn(async move {
395 + rpc_in_session(&client, &url, &session, 7, "tools/call", json!({ "name": "hang", "arguments": {} })).await
396 + })
397 + };
398 +
399 + // Let it register as in-flight, then cancel it app-side.
400 + tokio::time::sleep(Duration::from_millis(200)).await;
401 + assert_eq!(bound.cancel_all("human took the buffer back"), 1);
402 +
403 + // The agent's SSE stream receives a cancelled notification echoing id 7.
404 + let note = tokio::time::timeout(Duration::from_secs(2), next_sse_json(&mut sse, &mut buf))
405 + .await
406 + .expect("a cancelled notification arrives")
407 + .expect("the stream yields a notification");
408 + assert_eq!(note["method"], "notifications/cancelled");
409 + assert_eq!(note["params"]["requestId"], 7);
410 +
411 + // The blocked POST completes with a request-cancelled JSON-RPC error.
412 + let resp = tokio::time::timeout(Duration::from_secs(2), call)
413 + .await
414 + .expect("the call returns after cancellation")
415 + .unwrap();
416 + assert_eq!(resp["error"]["code"], -32800);
417 +
418 + bound.shutdown();
419 + }