From a445aad3c1918c80f8367fe4a4cbedf6d9d23e4b Mon Sep 17 00:00:00 2001 From: Mohamed Mansour Date: Thu, 17 Sep 2026 10:56:12 -0700 Subject: [PATCH 1/4] perf(rust): reduce notification payload cloning Consume owned router parameters and borrow events during session dispatch to avoid two redundant deep copies. Add an allocation benchmark and subscriber ownership regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/Cargo.lock | 7 + rust/Cargo.toml | 8 + rust/README.md | 22 +++ rust/benches/session_notifications.rs | 203 ++++++++++++++++++++++++++ rust/src/router.rs | 17 ++- rust/src/session.rs | 4 +- rust/tests/session_test.rs | 32 ++++ 7 files changed, 283 insertions(+), 10 deletions(-) create mode 100644 rust/benches/session_notifications.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index b91eebd06c..4c327dcf8b 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -445,6 +445,7 @@ dependencies = [ "serde_json", "serial_test", "sha2", + "stats_alloc", "tar", "tempfile", "tokio", @@ -1518,6 +1519,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "stats_alloc" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c0e04424e733e69714ca1bbb9204c1a57f09f5493439520f9f68c132ad25eec" + [[package]] name = "subtle" version = "2.6.1" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 4495d3928c..d7b280af36 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -14,6 +14,7 @@ license = "MIT" include = [ "src/**/*", "build/**/*", + "benches/**/*", "examples/**/*", "tests/**/*", "build.rs", @@ -81,6 +82,7 @@ windows-sys = { version = "0.61", default-features = false, features = [ rusqlite = { version = "0.35", features = ["bundled"] } schemars = "1" serial_test = "3" +stats_alloc = "0.1.10" tempfile = "3" tokio = { version = "1", features = ["rt-multi-thread"] } @@ -121,6 +123,12 @@ bench = false [[test]] name = "prepared_session_test" required-features = ["test-support"] + +[[bench]] +name = "session_notifications" +harness = false +required-features = ["test-support"] + [build-dependencies] dirs = "5" flate2 = "1" diff --git a/rust/README.md b/rust/README.md index 11d9637b22..1e35adf562 100644 --- a/rust/README.md +++ b/rust/README.md @@ -1223,3 +1223,25 @@ npm ci cd rust cargo test --features test-support ``` + +### Notification allocation benchmark + +```bash +cd rust +cargo bench --no-default-features --features test-support --bench session_notifications +``` + +This optimized-build benchmark drives Content-Length-framed notifications through +the SDK's real JSON-RPC reader, router, session loop, and subscriptions using +in-memory streams. It needs no running CLI or model requests. It covers 128-byte, +4-KiB, and 256-KiB payloads with zero, one, and two session observers, while also +exercising lifecycle routing. Each case uses four warmup batches and fifteen +measured batches of 64 events, with barriers ensuring both dispatchers finish. + +The CSV reports minimum, median, and maximum nanoseconds per event, allocation +counts, and allocated bytes per event. Frame construction and client startup are +outside the measured region; framing writes, dispatch, and payload verification +are included. The allocation counter is benchmark-only. Allocated bytes measure +heap allocation traffic, not peak heap usage or process RSS, and instrumented +timings are not production latency estimates. Compare repeated runs on the same +machine and toolchain. diff --git a/rust/benches/session_notifications.rs b/rust/benches/session_notifications.rs new file mode 100644 index 0000000000..9457e89c69 --- /dev/null +++ b/rust/benches/session_notifications.rs @@ -0,0 +1,203 @@ +#![allow(clippy::unwrap_used)] + +use std::alloc::System; +use std::hint::black_box; +use std::time::{Duration, Instant}; + +use github_copilot_sdk::{Client, SessionConfig}; +use serde_json::{Value, json}; +use stats_alloc::{INSTRUMENTED_SYSTEM, Region, StatsAlloc}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, DuplexStream, duplex}; + +#[global_allocator] +static ALLOCATOR: &StatsAlloc = &INSTRUMENTED_SYSTEM; + +const BATCH_SIZE: usize = 64; +const WARMUP_BATCHES: usize = 4; +const SAMPLES: usize = 15; + +async fn write_frame(writer: &mut DuplexStream, body: &[u8]) { + let header = format!("Content-Length: {}\r\n\r\n", body.len()); + writer.write_all(header.as_bytes()).await.unwrap(); + writer.write_all(body).await.unwrap(); +} + +async fn read_frame(reader: &mut BufReader) -> Value { + let mut header = String::new(); + reader.read_line(&mut header).await.unwrap(); + let length: usize = header + .trim() + .strip_prefix("Content-Length: ") + .unwrap() + .parse() + .unwrap(); + let mut separator = String::new(); + reader.read_line(&mut separator).await.unwrap(); + assert_eq!(separator, "\r\n"); + let mut body = vec![0; length]; + reader.read_exact(&mut body).await.unwrap(); + serde_json::from_slice(&body).unwrap() +} + +async fn measure(payload_size: usize, subscribers: usize) { + let workdir = tempfile::tempdir().unwrap(); + let (client_write, server_read) = duplex(8192); + let (mut server_write, client_read) = duplex(8192); + let client = + Client::from_streams(client_read, client_write, workdir.path().to_path_buf()).unwrap(); + let mut server_read = BufReader::new(server_read); + let create = client.create_session(SessionConfig::default()); + let respond = async { + let request = read_frame(&mut server_read).await; + assert_eq!(request["method"], "session.create"); + write_frame( + &mut server_write, + &serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": {"sessionId": request["params"]["sessionId"]} + })) + .unwrap(), + ) + .await; + }; + let (session, ()) = tokio::join!(create, respond); + let session = session.unwrap(); + let mut subscriptions: Vec<_> = (0..subscribers).map(|_| session.subscribe()).collect(); + let content = "x".repeat(payload_size); + let frames: Vec<_> = (0..BATCH_SIZE) + .map(|index| { + serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "method": "session.event", + "params": { + "sessionId": session.id(), + "event": { + "id": index.to_string(), + "timestamp": "2026-01-01T00:00:00Z", + "parentId": "previous-event", + "agentId": "test-agent", + "ephemeral": true, + "type": "assistant.message_delta", + "data": { + "messageId": "message", + "deltaContent": content, + "extra": {"nested": [1, true, null, {"key": "value"}]} + } + } + } + })) + .unwrap() + }) + .collect(); + // A lifecycle barrier exercises the second internal notification consumer + // and ensures its work is included even when there are no session observers. + let mut lifecycle = client.subscribe_lifecycle(); + let barrier = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "method": "session.lifecycle", + "params": {"type": "session.updated", "sessionId": session.id()} + })) + .unwrap(); + let session_barriers: Vec<_> = [false, true] + .map(|elicitation| { + serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "method": "session.event", + "params": { + "sessionId": session.id(), + "event": { + "id": "barrier", + "timestamp": "2026-01-01T00:00:00Z", + "type": "capabilities.changed", + "data": {"ui": {"elicitation": elicitation}} + } + } + })) + .unwrap() + }) + .into(); + let mut timings = Vec::with_capacity(SAMPLES); + let mut allocations = Vec::with_capacity(SAMPLES); + let mut allocated_bytes = Vec::with_capacity(SAMPLES); + + for sample in 0..WARMUP_BATCHES + SAMPLES { + let region = Region::new(ALLOCATOR); + let started = Instant::now(); + tokio::time::timeout(Duration::from_secs(30), async { + let send = async { + for frame in &frames { + write_frame(&mut server_write, frame).await; + } + write_frame(&mut server_write, &session_barriers[sample % 2]).await; + write_frame(&mut server_write, &barrier).await; + }; + let receive = async { + for index in 0..BATCH_SIZE { + for subscription in &mut subscriptions { + let event = subscription.recv().await.unwrap(); + assert_eq!(event.id.parse::().unwrap(), index); + assert_eq!(event.event_type, "assistant.message_delta"); + assert_eq!(event.data["deltaContent"].as_str().unwrap(), content); + assert_eq!(event.data["extra"]["nested"][3]["key"], "value"); + assert_eq!(event.agent_id.as_deref(), Some("test-agent")); + assert_eq!(event.parent_id.as_deref(), Some("previous-event")); + assert_eq!(event.ephemeral, Some(true)); + black_box(event); + } + } + for subscription in &mut subscriptions { + assert_eq!(subscription.recv().await.unwrap().id, "barrier"); + } + while session.capabilities().ui.and_then(|ui| ui.elicitation) + != Some(sample % 2 != 0) + { + tokio::task::yield_now().await; + } + lifecycle.recv().await.unwrap(); + }; + tokio::join!(send, receive); + }) + .await + .expect("notification pipeline stalled"); + let elapsed = started.elapsed().as_nanos() as f64 / BATCH_SIZE as f64; + let stats = region.change(); + if sample >= WARMUP_BATCHES { + timings.push(elapsed); + allocations.push(stats.allocations as f64 / BATCH_SIZE as f64); + allocated_bytes.push(stats.bytes_allocated as f64 / BATCH_SIZE as f64); + } + } + timings.sort_by(f64::total_cmp); + allocations.sort_by(f64::total_cmp); + allocated_bytes.sort_by(f64::total_cmp); + println!( + "{payload_size},{subscribers},{:.0},{:.0},{:.0},{:.1},{:.0}", + timings[0], + timings[SAMPLES / 2], + timings[SAMPLES - 1], + allocations[SAMPLES / 2], + allocated_bytes[SAMPLES / 2], + ); + session.stop_event_loop().await; + drop(session); + drop(client); + drop(server_write); + drop(server_read); + tokio::task::yield_now().await; +} + +fn main() { + println!( + "payload_bytes,subscribers,min_ns_per_event,median_ns_per_event,max_ns_per_event,allocations_per_event,allocated_bytes_per_event" + ); + for payload_size in [128, 4096, 262_144] { + for subscribers in [0, 1, 2] { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(measure(payload_size, subscribers)); + } + } +} diff --git a/rust/src/router.rs b/rust/src/router.rs index ce83d5f3ed..6adc2b7549 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -174,12 +174,12 @@ impl SessionRouter { // callback (if any) registered at client construction. if notification.method == "gitHubTelemetry.event" { if let Some(ref callback) = github_telemetry { - let Some(ref params) = notification.params else { + let Some(params) = notification.params else { continue; }; match serde_json::from_value::< crate::github_telemetry::GitHubTelemetryNotification, - >(params.clone()) + >(params) { Ok(telemetry) => { if std::panic::catch_unwind(std::panic::AssertUnwindSafe( @@ -206,7 +206,7 @@ impl SessionRouter { if notification.method != "session.event" { continue; } - let Some(ref params) = notification.params else { + let Some(params) = notification.params else { continue; }; let Some(session_id) = params.get("sessionId").and_then(|v| v.as_str()) @@ -216,18 +216,19 @@ impl SessionRouter { let sender = { let guard = sessions.lock(); - guard.get(session_id).map(|s| s.notifications.clone()) + guard + .get_key_value(session_id) + .map(|(id, s)| (id.clone(), s.notifications.clone())) }; - if let Some(sender) = sender { - match serde_json::from_value::(params.clone()) - { + if let Some((session_id, sender)) = sender { + match serde_json::from_value::(params) { Ok(event_notification) => { let _ = sender.send(event_notification); } Err(e) => { warn!( error = %e, - session_id = session_id, + session_id = %session_id, "failed to deserialize session event notification" ); } diff --git a/rust/src/session.rs b/rust/src/session.rs index 82e97b6606..68cc630462 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -2268,7 +2268,7 @@ async fn handle_notification( pending_external_tools: &PendingExternalTools, ) { let dispatch_start = Instant::now(); - let event = notification.event.clone(); + let event = ¬ification.event; let event_type = event.parsed_type(); if event_type == SessionEventType::PermissionRequested { tracing::debug!( @@ -2298,7 +2298,7 @@ async fn handle_notification( } waiter.last_assistant_message = Some(event.clone()); } - SessionEventType::SessionIdle if is_autopilot_continuation_idle(&event) => {} + SessionEventType::SessionIdle if is_autopilot_continuation_idle(event) => {} SessionEventType::SessionIdle | SessionEventType::SessionError => { if let Some(waiter) = guard.take() { if event_type == SessionEventType::SessionIdle { diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index 0a8bf52d5a..83af18379b 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -3857,6 +3857,38 @@ async fn session_event_notification_reaches_handler() { assert_eq!(event.event_type, "session.idle"); } +#[tokio::test] +async fn routed_event_preserves_owned_payload_for_each_subscriber() { + let (session, mut server) = create_session_pair().await; + let mut first = session.subscribe(); + let mut second = session.subscribe(); + let data = serde_json::json!({ + "deltaContent": "content".repeat(32 * 1024), + "extra": {"nested": [1, true, null, {"value": "preserved"}]}, + }); + + // A malformed notification must not prevent subsequent valid delivery. + server + .send_notification( + "session.event", + serde_json::json!({"sessionId": server.session_id, "event": {"data": data}}), + ) + .await; + server + .send_event("assistant.message_delta", data.clone()) + .await; + + let mut first_event = timeout(TIMEOUT, first.recv()).await.unwrap().unwrap(); + assert_eq!(first_event.data, data); + first_event.data["extra"]["nested"][3]["value"] = serde_json::json!("changed"); + first_event.data["deltaContent"] = serde_json::json!("replaced"); + + let second_event = timeout(TIMEOUT, second.recv()).await.unwrap().unwrap(); + assert_eq!(second_event.id, first_event.id); + assert_eq!(second_event.event_type, "assistant.message_delta"); + assert_eq!(second_event.data, data); +} + #[tokio::test] async fn router_routes_to_correct_session() { let (client, mut server_read, mut server_write) = make_client(); From 25e6fe302db1af1534b8fb384ea725c8a3bf1949 Mon Sep 17 00:00:00 2001 From: Mohamed Mansour Date: Thu, 17 Sep 2026 10:58:17 -0700 Subject: [PATCH 2/4] Remove benchmark dependency and simplify notification coverage Keep the ownership optimization focused: remove stats_alloc and benchmark documentation, and extend the existing subscriber test with a small payload instead of adding a separate large scenario. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/Cargo.lock | 7 - rust/Cargo.toml | 8 - rust/README.md | 22 --- rust/benches/session_notifications.rs | 203 -------------------------- rust/tests/session_test.rs | 29 +--- 5 files changed, 3 insertions(+), 266 deletions(-) delete mode 100644 rust/benches/session_notifications.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 4c327dcf8b..b91eebd06c 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -445,7 +445,6 @@ dependencies = [ "serde_json", "serial_test", "sha2", - "stats_alloc", "tar", "tempfile", "tokio", @@ -1519,12 +1518,6 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" -[[package]] -name = "stats_alloc" -version = "0.1.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c0e04424e733e69714ca1bbb9204c1a57f09f5493439520f9f68c132ad25eec" - [[package]] name = "subtle" version = "2.6.1" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index d7b280af36..4495d3928c 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -14,7 +14,6 @@ license = "MIT" include = [ "src/**/*", "build/**/*", - "benches/**/*", "examples/**/*", "tests/**/*", "build.rs", @@ -82,7 +81,6 @@ windows-sys = { version = "0.61", default-features = false, features = [ rusqlite = { version = "0.35", features = ["bundled"] } schemars = "1" serial_test = "3" -stats_alloc = "0.1.10" tempfile = "3" tokio = { version = "1", features = ["rt-multi-thread"] } @@ -123,12 +121,6 @@ bench = false [[test]] name = "prepared_session_test" required-features = ["test-support"] - -[[bench]] -name = "session_notifications" -harness = false -required-features = ["test-support"] - [build-dependencies] dirs = "5" flate2 = "1" diff --git a/rust/README.md b/rust/README.md index 1e35adf562..11d9637b22 100644 --- a/rust/README.md +++ b/rust/README.md @@ -1223,25 +1223,3 @@ npm ci cd rust cargo test --features test-support ``` - -### Notification allocation benchmark - -```bash -cd rust -cargo bench --no-default-features --features test-support --bench session_notifications -``` - -This optimized-build benchmark drives Content-Length-framed notifications through -the SDK's real JSON-RPC reader, router, session loop, and subscriptions using -in-memory streams. It needs no running CLI or model requests. It covers 128-byte, -4-KiB, and 256-KiB payloads with zero, one, and two session observers, while also -exercising lifecycle routing. Each case uses four warmup batches and fifteen -measured batches of 64 events, with barriers ensuring both dispatchers finish. - -The CSV reports minimum, median, and maximum nanoseconds per event, allocation -counts, and allocated bytes per event. Frame construction and client startup are -outside the measured region; framing writes, dispatch, and payload verification -are included. The allocation counter is benchmark-only. Allocated bytes measure -heap allocation traffic, not peak heap usage or process RSS, and instrumented -timings are not production latency estimates. Compare repeated runs on the same -machine and toolchain. diff --git a/rust/benches/session_notifications.rs b/rust/benches/session_notifications.rs deleted file mode 100644 index 9457e89c69..0000000000 --- a/rust/benches/session_notifications.rs +++ /dev/null @@ -1,203 +0,0 @@ -#![allow(clippy::unwrap_used)] - -use std::alloc::System; -use std::hint::black_box; -use std::time::{Duration, Instant}; - -use github_copilot_sdk::{Client, SessionConfig}; -use serde_json::{Value, json}; -use stats_alloc::{INSTRUMENTED_SYSTEM, Region, StatsAlloc}; -use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, DuplexStream, duplex}; - -#[global_allocator] -static ALLOCATOR: &StatsAlloc = &INSTRUMENTED_SYSTEM; - -const BATCH_SIZE: usize = 64; -const WARMUP_BATCHES: usize = 4; -const SAMPLES: usize = 15; - -async fn write_frame(writer: &mut DuplexStream, body: &[u8]) { - let header = format!("Content-Length: {}\r\n\r\n", body.len()); - writer.write_all(header.as_bytes()).await.unwrap(); - writer.write_all(body).await.unwrap(); -} - -async fn read_frame(reader: &mut BufReader) -> Value { - let mut header = String::new(); - reader.read_line(&mut header).await.unwrap(); - let length: usize = header - .trim() - .strip_prefix("Content-Length: ") - .unwrap() - .parse() - .unwrap(); - let mut separator = String::new(); - reader.read_line(&mut separator).await.unwrap(); - assert_eq!(separator, "\r\n"); - let mut body = vec![0; length]; - reader.read_exact(&mut body).await.unwrap(); - serde_json::from_slice(&body).unwrap() -} - -async fn measure(payload_size: usize, subscribers: usize) { - let workdir = tempfile::tempdir().unwrap(); - let (client_write, server_read) = duplex(8192); - let (mut server_write, client_read) = duplex(8192); - let client = - Client::from_streams(client_read, client_write, workdir.path().to_path_buf()).unwrap(); - let mut server_read = BufReader::new(server_read); - let create = client.create_session(SessionConfig::default()); - let respond = async { - let request = read_frame(&mut server_read).await; - assert_eq!(request["method"], "session.create"); - write_frame( - &mut server_write, - &serde_json::to_vec(&json!({ - "jsonrpc": "2.0", - "id": request["id"], - "result": {"sessionId": request["params"]["sessionId"]} - })) - .unwrap(), - ) - .await; - }; - let (session, ()) = tokio::join!(create, respond); - let session = session.unwrap(); - let mut subscriptions: Vec<_> = (0..subscribers).map(|_| session.subscribe()).collect(); - let content = "x".repeat(payload_size); - let frames: Vec<_> = (0..BATCH_SIZE) - .map(|index| { - serde_json::to_vec(&json!({ - "jsonrpc": "2.0", - "method": "session.event", - "params": { - "sessionId": session.id(), - "event": { - "id": index.to_string(), - "timestamp": "2026-01-01T00:00:00Z", - "parentId": "previous-event", - "agentId": "test-agent", - "ephemeral": true, - "type": "assistant.message_delta", - "data": { - "messageId": "message", - "deltaContent": content, - "extra": {"nested": [1, true, null, {"key": "value"}]} - } - } - } - })) - .unwrap() - }) - .collect(); - // A lifecycle barrier exercises the second internal notification consumer - // and ensures its work is included even when there are no session observers. - let mut lifecycle = client.subscribe_lifecycle(); - let barrier = serde_json::to_vec(&json!({ - "jsonrpc": "2.0", - "method": "session.lifecycle", - "params": {"type": "session.updated", "sessionId": session.id()} - })) - .unwrap(); - let session_barriers: Vec<_> = [false, true] - .map(|elicitation| { - serde_json::to_vec(&json!({ - "jsonrpc": "2.0", - "method": "session.event", - "params": { - "sessionId": session.id(), - "event": { - "id": "barrier", - "timestamp": "2026-01-01T00:00:00Z", - "type": "capabilities.changed", - "data": {"ui": {"elicitation": elicitation}} - } - } - })) - .unwrap() - }) - .into(); - let mut timings = Vec::with_capacity(SAMPLES); - let mut allocations = Vec::with_capacity(SAMPLES); - let mut allocated_bytes = Vec::with_capacity(SAMPLES); - - for sample in 0..WARMUP_BATCHES + SAMPLES { - let region = Region::new(ALLOCATOR); - let started = Instant::now(); - tokio::time::timeout(Duration::from_secs(30), async { - let send = async { - for frame in &frames { - write_frame(&mut server_write, frame).await; - } - write_frame(&mut server_write, &session_barriers[sample % 2]).await; - write_frame(&mut server_write, &barrier).await; - }; - let receive = async { - for index in 0..BATCH_SIZE { - for subscription in &mut subscriptions { - let event = subscription.recv().await.unwrap(); - assert_eq!(event.id.parse::().unwrap(), index); - assert_eq!(event.event_type, "assistant.message_delta"); - assert_eq!(event.data["deltaContent"].as_str().unwrap(), content); - assert_eq!(event.data["extra"]["nested"][3]["key"], "value"); - assert_eq!(event.agent_id.as_deref(), Some("test-agent")); - assert_eq!(event.parent_id.as_deref(), Some("previous-event")); - assert_eq!(event.ephemeral, Some(true)); - black_box(event); - } - } - for subscription in &mut subscriptions { - assert_eq!(subscription.recv().await.unwrap().id, "barrier"); - } - while session.capabilities().ui.and_then(|ui| ui.elicitation) - != Some(sample % 2 != 0) - { - tokio::task::yield_now().await; - } - lifecycle.recv().await.unwrap(); - }; - tokio::join!(send, receive); - }) - .await - .expect("notification pipeline stalled"); - let elapsed = started.elapsed().as_nanos() as f64 / BATCH_SIZE as f64; - let stats = region.change(); - if sample >= WARMUP_BATCHES { - timings.push(elapsed); - allocations.push(stats.allocations as f64 / BATCH_SIZE as f64); - allocated_bytes.push(stats.bytes_allocated as f64 / BATCH_SIZE as f64); - } - } - timings.sort_by(f64::total_cmp); - allocations.sort_by(f64::total_cmp); - allocated_bytes.sort_by(f64::total_cmp); - println!( - "{payload_size},{subscribers},{:.0},{:.0},{:.0},{:.1},{:.0}", - timings[0], - timings[SAMPLES / 2], - timings[SAMPLES - 1], - allocations[SAMPLES / 2], - allocated_bytes[SAMPLES / 2], - ); - session.stop_event_loop().await; - drop(session); - drop(client); - drop(server_write); - drop(server_read); - tokio::task::yield_now().await; -} - -fn main() { - println!( - "payload_bytes,subscribers,min_ns_per_event,median_ns_per_event,max_ns_per_event,allocations_per_event,allocated_bytes_per_event" - ); - for payload_size in [128, 4096, 262_144] { - for subscribers in [0, 1, 2] { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - runtime.block_on(measure(payload_size, subscribers)); - } - } -} diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index 83af18379b..f9d196672d 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -3846,42 +3846,19 @@ async fn permission_result_forwards_context_beside_result() { } #[tokio::test] -async fn session_event_notification_reaches_handler() { - let (session, mut server) = create_session_pair().await; - let mut sub = session.subscribe(); - server - .send_event("session.idle", serde_json::json!({})) - .await; - - let event = timeout(TIMEOUT, sub.recv()).await.unwrap().unwrap(); - assert_eq!(event.event_type, "session.idle"); -} - -#[tokio::test] -async fn routed_event_preserves_owned_payload_for_each_subscriber() { +async fn session_event_notification_reaches_subscribers() { let (session, mut server) = create_session_pair().await; let mut first = session.subscribe(); let mut second = session.subscribe(); - let data = serde_json::json!({ - "deltaContent": "content".repeat(32 * 1024), - "extra": {"nested": [1, true, null, {"value": "preserved"}]}, - }); + let data = serde_json::json!({"deltaContent": "hello"}); - // A malformed notification must not prevent subsequent valid delivery. - server - .send_notification( - "session.event", - serde_json::json!({"sessionId": server.session_id, "event": {"data": data}}), - ) - .await; server .send_event("assistant.message_delta", data.clone()) .await; let mut first_event = timeout(TIMEOUT, first.recv()).await.unwrap().unwrap(); assert_eq!(first_event.data, data); - first_event.data["extra"]["nested"][3]["value"] = serde_json::json!("changed"); - first_event.data["deltaContent"] = serde_json::json!("replaced"); + first_event.data["deltaContent"] = serde_json::json!("changed"); let second_event = timeout(TIMEOUT, second.recv()).await.unwrap().unwrap(); assert_eq!(second_event.id, first_event.id); From a3b61f839561fd8786fbc2372838a1e34b638714 Mon Sep 17 00:00:00 2001 From: Mohamed Mansour Date: Thu, 17 Sep 2026 11:05:23 -0700 Subject: [PATCH 3/4] test(rust): cover routing after malformed notifications Send an invalid event before the valid payload in the existing subscriber test, without adding fixtures or helpers. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/tests/session_test.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index f9d196672d..a663f24477 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -3852,6 +3852,12 @@ async fn session_event_notification_reaches_subscribers() { let mut second = session.subscribe(); let data = serde_json::json!({"deltaContent": "hello"}); + server + .send_notification( + "session.event", + serde_json::json!({"sessionId": server.session_id, "event": null}), + ) + .await; server .send_event("assistant.message_delta", data.clone()) .await; From 98e4c4f60dbb654134431e5869db731017715d99 Mon Sep 17 00:00:00 2001 From: Mohamed Mansour Date: Thu, 17 Sep 2026 20:41:20 -0700 Subject: [PATCH 4/4] perf(rust): preserve owned JSON payload containers during decoding Separate payload ownership from metadata deserialization in JSON-RPC messages and session notifications, avoiding recursive Value container reconstruction. Preserve optional/null semantics, envelope validation, diagnostics, and existing subscriber ownership. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/jsonrpc.rs | 88 +++++++++++++++++++++++++++++++++++--- rust/src/router.rs | 13 +++++- rust/tests/session_test.rs | 28 ++++++++++++ 3 files changed, 120 insertions(+), 9 deletions(-) diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 48e6090aed..1a590c555b 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -124,25 +124,43 @@ impl<'de> Deserialize<'de> for JsonRpcMessage { where D: serde::Deserializer<'de>, { - let value = Value::deserialize(deserializer)?; + let mut value = Value::deserialize(deserializer)?; let obj = value - .as_object() + .as_object_mut() .ok_or_else(|| serde::de::Error::custom("expected a JSON object"))?; let has_id = obj.contains_key("id"); let has_method = obj.contains_key("method"); + // Preserve the owned payload instead of rebuilding its JSON containers + // while serde validates the envelope. Optional null payloads remain None. + let payload_key = if has_id && !has_method { + "result" + } else { + "params" + }; + let payload = obj.remove(payload_key).filter(|value| !value.is_null()); + if has_id && has_method { JsonRpcRequest::deserialize(value) - .map(JsonRpcMessage::Request) + .map(|mut request| { + request.params = payload; + JsonRpcMessage::Request(request) + }) .map_err(serde::de::Error::custom) } else if has_id { JsonRpcResponse::deserialize(value) - .map(JsonRpcMessage::Response) + .map(|mut response| { + response.result = payload; + JsonRpcMessage::Response(response) + }) .map_err(serde::de::Error::custom) } else { JsonRpcNotification::deserialize(value) - .map(JsonRpcMessage::Notification) + .map(|mut notification| { + notification.params = payload; + JsonRpcMessage::Notification(notification) + }) .map_err(serde::de::Error::custom) } } @@ -730,8 +748,7 @@ mod tests { #[test] fn deserialize_error_response() { - let json = - r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32600,"message":"Invalid Request"}}"#; + let json = r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32600,"message":"Invalid Request","data":{"nested":[1,{"reason":"invalid"}]}}}"#; let msg: JsonRpcMessage = serde_json::from_str(json).unwrap(); match msg { JsonRpcMessage::Response(r) => { @@ -739,6 +756,10 @@ mod tests { let err = r.error.unwrap(); assert_eq!(err.code, -32600); assert_eq!(err.message, "Invalid Request"); + assert_eq!( + err.data, + Some(serde_json::json!({"nested": [1, {"reason": "invalid"}]})) + ); } other => panic!("expected Response, got {other:?}"), } @@ -750,6 +771,59 @@ mod tests { assert!(result.is_err()); } + #[test] + fn deserialize_preserves_optional_payloads() { + for payload in [ + None, + Some(Value::Null), + Some(serde_json::json!(false)), + Some(serde_json::json!(42)), + Some(serde_json::json!("text")), + Some(serde_json::json!([{"nested": [1, null, true]}])), + Some(serde_json::json!({"rows": [{"content": "result"}]})), + ] { + for mut envelope in [ + serde_json::json!({"jsonrpc": "2.0", "method": "notify"}), + serde_json::json!({"jsonrpc": "2.0", "id": 1, "method": "request"}), + serde_json::json!({"jsonrpc": "2.0", "id": 1}), + ] { + let (payload_key, ignored_key) = if envelope.get("method").is_some() { + ("params", "result") + } else { + ("result", "params") + }; + envelope[ignored_key] = serde_json::json!({"ignored": "opposite payload"}); + if let Some(payload) = &payload { + envelope[payload_key] = payload.clone(); + } + let actual = match serde_json::from_value::(envelope).unwrap() { + JsonRpcMessage::Request(request) => request.params, + JsonRpcMessage::Response(response) => response.result, + JsonRpcMessage::Notification(notification) => notification.params, + }; + assert_eq!(actual, payload.clone().filter(|value| !value.is_null())); + } + } + } + + #[test] + fn deserialize_rejects_invalid_metadata() { + for json in [ + r#"{"jsonrpc":null,"method":"notify","params":{"nested":[1]}}"#, + r#"{"jsonrpc":"2.0","method":42,"params":{"nested":[1]}}"#, + r#"{"jsonrpc":"2.0","id":null,"result":{}}"#, + r#"{"jsonrpc":"2.0","id":"1","result":{}}"#, + r#"{"jsonrpc":"2.0","id":-1,"result":{}}"#, + r#"{"jsonrpc":"2.0","id":1,"method":null,"params":{}}"#, + r#"{"jsonrpc":"2.0","id":1,"result":{},"error":{"code":"bad","message":"error"}}"#, + ] { + assert!( + serde_json::from_str::(json).is_err(), + "{json}" + ); + } + } + #[test] fn request_new_sets_version() { let req = JsonRpcRequest::new(42, "test.method", None); diff --git a/rust/src/router.rs b/rust/src/router.rs index 6adc2b7549..646f999ad3 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -206,7 +206,7 @@ impl SessionRouter { if notification.method != "session.event" { continue; } - let Some(params) = notification.params else { + let Some(mut params) = notification.params else { continue; }; let Some(session_id) = params.get("sessionId").and_then(|v| v.as_str()) @@ -221,8 +221,17 @@ impl SessionRouter { .map(|(id, s)| (id.clone(), s.notifications.clone())) }; if let Some((session_id, sender)) = sender { + // Leave null in the existing slot so serde still rejects + // missing data, without rebuilding the owned payload. + let data = params + .get_mut("event") + .and_then(|event| event.get_mut("data")) + .map(serde_json::Value::take); match serde_json::from_value::(params) { - Ok(event_notification) => { + Ok(mut event_notification) => { + if let Some(data) = data { + event_notification.event.data = data; + } let _ = sender.send(event_notification); } Err(e) => { diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index a663f24477..0310ae380f 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -3872,6 +3872,34 @@ async fn session_event_notification_reaches_subscribers() { assert_eq!(second_event.data, data); } +#[tokio::test] +async fn session_event_notification_preserves_unknown_event_payloads() { + let (session, mut server) = create_session_pair().await; + let mut events = session.subscribe(); + server + .send_notification( + "session.event", + serde_json::json!({ + "sessionId": server.session_id, + "event": {"id": "invalid", "timestamp": "now", "type": "future.event"} + }), + ) + .await; + + for data in [ + Value::Null, + serde_json::json!("text"), + serde_json::json!(false), + serde_json::json!([1, {"nested": [null, true]}]), + serde_json::json!({"result": {"rows": [{"content": "preserved"}]}}), + ] { + server.send_event("future.event", data.clone()).await; + let event = timeout(TIMEOUT, events.recv()).await.unwrap().unwrap(); + assert_eq!(event.event_type, "future.event"); + assert_eq!(event.data, data); + } +} + #[tokio::test] async fn router_routes_to_correct_session() { let (client, mut server_read, mut server_write) = make_client();