From 9994b2c2359afcc9176c0a7044ba14f23707c0af Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 28 Aug 2026 12:08:14 -0600 Subject: [PATCH 1/7] Add reverse RPC timing traces Measure inbound scheduling, hook callback execution, response encoding, writer queue delay, write_all, and flush without exposing callback content or changing RPC behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/Cargo.lock | 66 +++ rust/Cargo.toml | 3 +- rust/src/copilot_request_handler.rs | 3 + rust/src/github_token.rs | 1 + rust/src/hooks.rs | 149 +++++- rust/src/jsonrpc.rs | 751 +++++++++++++++++++++++++++- rust/src/lib.rs | 19 +- rust/src/router.rs | 22 +- rust/src/session.rs | 14 +- 9 files changed, 1010 insertions(+), 18 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 8de6797989..dc7d693485 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -452,6 +452,7 @@ dependencies = [ "tokio-tungstenite", "tokio-util", "tracing", + "tracing-subscriber", "ureq", "uuid", "zip", @@ -764,6 +765,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "leb128fmt" version = "0.1.0" @@ -880,6 +887,15 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1467,6 +1483,15 @@ dependencies = [ "digest", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "1.3.0" @@ -1618,6 +1643,15 @@ dependencies = [ "syn", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -1788,6 +1822,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "nu-ansi-term", + "sharded-slab", + "smallvec", + "thread_local", + "tracing-core", + "tracing-log", ] [[package]] @@ -1885,6 +1945,12 @@ dependencies = [ "getrandom 0.4.2", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 0f18a9b159..b686625375 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -76,7 +76,8 @@ rusqlite = { version = "0.35", features = ["bundled"] } schemars = "1" serial_test = "3" tempfile = "3" -tokio = { version = "1", features = ["rt-multi-thread"] } +tokio = { version = "1", features = ["rt-multi-thread", "test-util"] } +tracing-subscriber = { version = "0.3", features = ["fmt"] } # Integration tests that call test-support-only Client methods (e.g. # `from_streams_with_connection_token`, `from_streams_with_trace_provider`) diff --git a/rust/src/copilot_request_handler.rs b/rust/src/copilot_request_handler.rs index 961ae3876e..d4bf2af1fb 100644 --- a/rust/src/copilot_request_handler.rs +++ b/rust/src/copilot_request_handler.rs @@ -1048,6 +1048,9 @@ impl CopilotRequestDispatcher { } pub(crate) async fn dispatch(self: &Arc, request: JsonRpcRequest) { + let _reverse_rpc = self + .client() + .and_then(|client| client.trace_reverse_request_scheduled(request.id)); match request.method.as_str() { METHOD_HTTP_REQUEST_START => self.handle_start(request).await, METHOD_HTTP_REQUEST_CHUNK => self.handle_chunk(request).await, diff --git a/rust/src/github_token.rs b/rust/src/github_token.rs index c5eaa63adf..456c5a86b9 100644 --- a/rust/src/github_token.rs +++ b/rust/src/github_token.rs @@ -193,6 +193,7 @@ impl GitHubTokenRegistry { return; }; let client = Client::from_inner(inner); + let _reverse_rpc = client.trace_reverse_request_scheduled(request.id); let params = request .params .clone() diff --git a/rust/src/hooks.rs b/rust/src/hooks.rs index 4986d6cb18..db79400c65 100644 --- a/rust/src/hooks.rs +++ b/rust/src/hooks.rs @@ -6,11 +6,11 @@ //! [`Client::create_session`](crate::Client::create_session). use std::path::PathBuf; -use std::time::Instant; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use serde_json::Value; +use tokio::time::Instant; use crate::types::SessionId; @@ -680,11 +680,22 @@ pub trait SessionHooks: Send + Sync + 'static { /// Returns `Ok(Value)` shaped like `{ "output": ... }` on success. /// If no hook is registered ([`HookOutput::None`]), the output is an empty /// object: `{ "output": {} }`. +#[cfg(test)] pub(crate) async fn dispatch_hook( hooks: &dyn SessionHooks, session_id: &SessionId, hook_type: &str, raw_input: Value, +) -> Result { + dispatch_hook_traced(hooks, session_id, hook_type, raw_input, None).await +} + +pub(crate) async fn dispatch_hook_traced( + hooks: &dyn SessionHooks, + session_id: &SessionId, + hook_type: &str, + raw_input: Value, + reverse_rpc_trace: Option<&crate::jsonrpc::ReverseRpcTrace>, ) -> Result { let ctx = HookContext { session_id: session_id.clone(), @@ -743,8 +754,12 @@ pub(crate) async fn dispatch_hook( let dispatch_start = Instant::now(); let output = hooks.on_hook(event).await; + let dispatch_elapsed = dispatch_start.elapsed(); + if let Some(trace) = reverse_rpc_trace { + trace.record_hook_callback(hook_type, dispatch_elapsed); + } tracing::debug!( - elapsed_ms = dispatch_start.elapsed().as_millis(), + elapsed_ms = dispatch_elapsed.as_millis(), session_id = %session_id, hook_type = hook_type, "SessionHooks::on_hook dispatch" @@ -786,7 +801,58 @@ pub(crate) async fn dispatch_hook( #[cfg(test)] mod tests { + use std::io::Write; + use std::sync::Arc; + use std::time::Duration; + + use parking_lot::Mutex; + use tokio::sync::Notify; + use tracing::Instrument; + use tracing_subscriber::Layer; + use tracing_subscriber::fmt::MakeWriter; + use tracing_subscriber::layer::SubscriberExt; + use super::*; + use crate::JsonRpcRequest; + use crate::jsonrpc::ReverseRpcTrace; + + #[derive(Clone, Default)] + struct TraceBuffer(Arc>>); + + impl TraceBuffer { + fn text(&self) -> String { + String::from_utf8(self.0.lock().clone()).unwrap() + } + } + + impl Write for TraceBuffer { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl<'a> MakeWriter<'a> for TraceBuffer { + type Writer = Self; + + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } + } + + async fn wait_for_trace(buffer: &TraceBuffer, needle: &str) { + for _ in 0..20 { + if buffer.text().contains(needle) { + return; + } + tokio::task::yield_now().await; + } + panic!("timing trace did not contain {needle:?}: {}", buffer.text()); + } struct TestHooks; @@ -842,6 +908,85 @@ mod tests { assert_eq!(output["permissionDecisionReason"], "blocked by policy"); } + #[tokio::test(start_paused = true)] + async fn traced_dispatch_measures_only_the_gated_hook_callback() { + const SENTINEL: &str = "PRIVATE_HOOK_SENTINEL_DO_NOT_TRACE"; + + struct GatedHooks { + started: Arc, + release: Arc, + } + + #[async_trait] + impl SessionHooks for GatedHooks { + async fn on_hook(&self, _event: HookEvent) -> HookOutput { + self.started.notify_one(); + self.release.notified().await; + HookOutput::UserPromptSubmitted(UserPromptSubmittedOutput { + modified_prompt: Some(SENTINEL.to_string()), + ..Default::default() + }) + } + } + + let trace_buffer = TraceBuffer::default(); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_writer(trace_buffer.clone()) + .with_ansi(false) + .without_time() + .with_filter(tracing_subscriber::filter::filter_fn(|metadata| { + metadata.target() == "github_copilot_sdk::reverse_rpc_timing" + })), + ); + let _subscriber = tracing::subscriber::set_default(subscriber); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let hooks = Arc::new(GatedHooks { + started: started.clone(), + release: release.clone(), + }); + let request = JsonRpcRequest::new( + 99, + "hooks.invoke", + Some(serde_json::json!({ "sessionId": "session-1" })), + ); + let now = Instant::now(); + let trace = ReverseRpcTrace::for_test(&request, now, now); + + let parent = tracing::error_span!("session_request_handler", session_id = SENTINEL); + let dispatch = tokio::spawn( + async move { + dispatch_hook_traced( + hooks.as_ref(), + &SessionId::new("session-1"), + "userPromptSubmitted", + serde_json::json!({ + "sessionId": SENTINEL, + "timestamp": 1234567890, + "cwd": SENTINEL, + "prompt": SENTINEL + }), + Some(&trace), + ) + .await + } + .instrument(parent), + ); + + started.notified().await; + tokio::time::advance(Duration::from_millis(9)).await; + release.notify_one(); + let output = dispatch.await.unwrap().unwrap(); + assert_eq!(output["output"]["modifiedPrompt"], SENTINEL); + + wait_for_trace(&trace_buffer, "phase=\"hook_callback\"").await; + let traces = trace_buffer.text(); + assert!(traces.contains("phase=\"hook_callback\"")); + assert!(traces.contains("elapsed_us=9000")); + assert!(!traces.contains(SENTINEL)); + } + #[tokio::test] async fn dispatch_pre_tool_use_passthrough() { let hooks = TestHooks; diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 25a405080b..9506c67ce6 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -9,6 +9,7 @@ use serde_json::Value; use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader}; use tokio::sync::{broadcast, mpsc, oneshot}; use tokio::task::JoinHandle; +use tokio::time::Instant as TokioInstant; use tracing::{Instrument, debug, error, warn}; use crate::{Error, ErrorKind, ProtocolErrorKind}; @@ -168,6 +169,15 @@ impl JsonRpcResponse { } const CONTENT_LENGTH_HEADER: &str = "Content-Length: "; +/// Opt-in target for content-free reverse-RPC latency events. +/// +/// `request_forward` measures parsed-request receipt to forwarding, +/// `request_schedule` forwarding to dispatch start (`since_receive_us` is +/// the combined interval), `response_encode` is `serde_json::to_vec`, +/// `writer_queue` is enqueue to dequeue, and `write_all` / `flush` measure +/// the corresponding `AsyncWrite` calls. `hook_callback` is emitted by the +/// hooks dispatcher around `SessionHooks::on_hook` only. +const REVERSE_RPC_TIMING_TARGET: &str = "github_copilot_sdk::reverse_rpc_timing"; /// Rewrites unpaired UTF-16 surrogate escapes to `\uFFFD`. /// @@ -242,6 +252,170 @@ fn repair_lone_surrogates(body: &[u8]) -> Option> { struct WriteCommand { frame: Vec, ack: oneshot::Sender>, + reverse_rpc: Option, + enqueued_at: TokioInstant, +} + +enum ReverseRpcTimingEvent { + Phase { + trace: ReverseRpcTrace, + phase: &'static str, + elapsed_us: u64, + succeeded: bool, + }, + Scheduled { + trace: ReverseRpcTrace, + elapsed_us: u64, + since_receive_us: u64, + }, + HookCallback { + trace: ReverseRpcTrace, + hook_type: String, + elapsed_us: u64, + }, +} + +/// Internal, content-free timing context for one inbound JSON-RPC request. +/// +/// The correlation key is a deterministic digest of the numeric wire ID, +/// RPC method, and optional session ID. The original session ID and params +/// are not retained. +#[derive(Clone)] +pub(crate) struct ReverseRpcTrace { + inner: Arc, +} + +struct ReverseRpcTraceInner { + correlation_key: String, + method: String, + received_at: TokioInstant, + forwarded_at: TokioInstant, + timing_tx: mpsc::UnboundedSender, +} + +impl ReverseRpcTrace { + fn new( + request: &JsonRpcRequest, + received_at: TokioInstant, + forwarded_at: TokioInstant, + timing_tx: mpsc::UnboundedSender, + ) -> Self { + let session_id = request + .params + .as_ref() + .and_then(|params| params.get("sessionId")) + .and_then(Value::as_str); + Self { + inner: Arc::new(ReverseRpcTraceInner { + correlation_key: Self::correlation_key(request.id, &request.method, session_id), + method: request.method.clone(), + received_at, + forwarded_at, + timing_tx, + }), + } + } + + #[cfg(test)] + pub(crate) fn for_test( + request: &JsonRpcRequest, + received_at: TokioInstant, + forwarded_at: TokioInstant, + ) -> Self { + let (timing_tx, timing_rx) = mpsc::unbounded_channel(); + tokio::spawn(JsonRpcClient::timing_loop(timing_rx)); + Self::new(request, received_at, forwarded_at, timing_tx) + } + + fn correlation_key(request_id: u64, method: &str, session_id: Option<&str>) -> String { + let mut hash = 0xcbf29ce484222325_u64; + for byte in session_id + .unwrap_or("") + .bytes() + .chain([0xff]) + .chain(method.bytes()) + .chain([0xfe]) + .chain(request_id.to_le_bytes()) + { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x100000001b3); + } + format!("rrpc-{hash:016x}") + } + + fn elapsed_us(duration: std::time::Duration) -> u64 { + u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) + } + + fn record_forwarded(&self, succeeded: bool) { + self.record_phase( + "request_forward", + self.inner + .forwarded_at + .duration_since(self.inner.received_at), + succeeded, + ); + } + + fn record_scheduled(&self, scheduled_at: TokioInstant) { + let _ = self.inner.timing_tx.send(ReverseRpcTimingEvent::Scheduled { + trace: self.clone(), + elapsed_us: Self::elapsed_us(scheduled_at.duration_since(self.inner.forwarded_at)), + since_receive_us: Self::elapsed_us(scheduled_at.duration_since(self.inner.received_at)), + }); + } + + pub(crate) fn record_hook_callback(&self, hook_type: &str, elapsed: std::time::Duration) { + let _ = self + .inner + .timing_tx + .send(ReverseRpcTimingEvent::HookCallback { + trace: self.clone(), + hook_type: hook_type.to_string(), + elapsed_us: Self::elapsed_us(elapsed), + }); + } + + fn record_phase(&self, phase: &'static str, elapsed: std::time::Duration, succeeded: bool) { + let _ = self.inner.timing_tx.send(ReverseRpcTimingEvent::Phase { + trace: self.clone(), + phase, + elapsed_us: Self::elapsed_us(elapsed), + succeeded, + }); + } +} + +pub(crate) struct ReverseRpcDispatchGuard { + reverse_requests: Arc>>, + request_id: u64, + trace: ReverseRpcTrace, +} + +impl ReverseRpcDispatchGuard { + pub(crate) fn trace(&self) -> &ReverseRpcTrace { + &self.trace + } +} + +impl Drop for ReverseRpcDispatchGuard { + fn drop(&mut self) { + remove_reverse_request_if_same(&self.reverse_requests, self.request_id, &self.trace); + } +} + +fn remove_reverse_request_if_same( + reverse_requests: &RwLock>, + request_id: u64, + trace: &ReverseRpcTrace, +) { + let mut reverse_requests = reverse_requests.write(); + if reverse_requests + .get(&request_id) + .is_some_and(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) + { + reverse_requests.remove(&request_id); + } } /// Low-level JSON-RPC 2.0 client over Content-Length-framed streams. @@ -264,10 +438,12 @@ pub struct JsonRpcClient { /// natural request/response back-pressure of the wire. write_tx: mpsc::UnboundedSender, pending_requests: Arc>>, + reverse_requests: Arc>>, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, read_task: Mutex>>, write_task: Mutex>>, + timing_task: Mutex>>, } impl JsonRpcClient { @@ -282,23 +458,54 @@ impl JsonRpcClient { reader: impl AsyncRead + Unpin + Send + 'static, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, + ) -> Self { + Self::new_inner(writer, reader, notification_tx, request_tx, false) + } + + pub(crate) fn new_with_reverse_rpc_timing( + writer: impl AsyncWrite + Unpin + Send + 'static, + reader: impl AsyncRead + Unpin + Send + 'static, + notification_tx: broadcast::Sender, + request_tx: mpsc::UnboundedSender, + ) -> Self { + Self::new_inner(writer, reader, notification_tx, request_tx, true) + } + + fn new_inner( + writer: impl AsyncWrite + Unpin + Send + 'static, + reader: impl AsyncRead + Unpin + Send + 'static, + notification_tx: broadcast::Sender, + request_tx: mpsc::UnboundedSender, + trace_reverse_rpc: bool, ) -> Self { let (write_tx, write_rx) = mpsc::unbounded_channel::(); let writer_span = tracing::error_span!("jsonrpc_write_loop"); let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span)); + let (timing_tx, timing_task) = if trace_reverse_rpc { + let (timing_tx, timing_rx) = mpsc::unbounded_channel::(); + ( + Some(timing_tx), + Some(tokio::spawn(Self::timing_loop(timing_rx))), + ) + } else { + (None, None) + }; let client = Self { request_id: AtomicU64::new(1), write_tx, pending_requests: Arc::new(RwLock::new(HashMap::new())), + reverse_requests: Arc::new(RwLock::new(HashMap::new())), notification_tx, request_tx, read_task: Mutex::new(None), write_task: Mutex::new(Some(write_task)), + timing_task: Mutex::new(timing_task), }; let pending_requests = client.pending_requests.clone(); + let reverse_requests = client.reverse_requests.clone(); let notification_tx_clone = client.notification_tx.clone(); let request_tx_clone = client.request_tx.clone(); let reader_span = tracing::error_span!("jsonrpc_read_loop"); @@ -308,8 +515,10 @@ impl JsonRpcClient { Self::read_loop( reader, pending_requests, + reverse_requests, notification_tx_clone, request_tx_clone, + timing_tx, ) .await; } @@ -327,7 +536,69 @@ impl JsonRpcClient { if let Some(task) = self.write_task.lock().take() { task.abort(); } + if let Some(task) = self.timing_task.lock().take() { + task.abort(); + } self.pending_requests.write().clear(); + self.reverse_requests.write().clear(); + } + + async fn timing_loop(mut rx: mpsc::UnboundedReceiver) { + while let Some(event) = rx.recv().await { + match event { + ReverseRpcTimingEvent::Phase { + trace, + phase, + elapsed_us, + succeeded, + } => { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + correlation_key = %trace.inner.correlation_key, + rpc_method = %trace.inner.method, + phase, + elapsed_us, + status = if succeeded { "succeeded" } else { "failed" }, + "reverse JSON-RPC timing" + ); + } + ReverseRpcTimingEvent::Scheduled { + trace, + elapsed_us, + since_receive_us, + } => { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + correlation_key = %trace.inner.correlation_key, + rpc_method = %trace.inner.method, + phase = "request_schedule", + elapsed_us, + since_receive_us, + status = "succeeded", + "reverse JSON-RPC timing" + ); + } + ReverseRpcTimingEvent::HookCallback { + trace, + hook_type, + elapsed_us, + } => { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + correlation_key = %trace.inner.correlation_key, + rpc_method = %trace.inner.method, + hook_type, + phase = "hook_callback", + elapsed_us, + status = "succeeded", + "reverse JSON-RPC timing" + ); + } + } + } } /// Writer-actor task. Owns the `AsyncWrite`, drains the command queue, @@ -346,26 +617,53 @@ impl JsonRpcClient { mut writer: impl AsyncWrite + Unpin + Send + 'static, mut rx: mpsc::UnboundedReceiver, ) { - while let Some(WriteCommand { frame, ack }) = rx.recv().await { - let result = async { - writer.write_all(&frame).await?; - writer.flush().await?; - Ok::<_, std::io::Error>(()) - } - .await; + while let Some(WriteCommand { + frame, + ack, + reverse_rpc, + enqueued_at, + }) = rx.recv().await + { + let queue_elapsed = enqueued_at.elapsed(); + + let write_start = TokioInstant::now(); + let write_result = writer.write_all(&frame).await; + let write_elapsed = write_start.elapsed(); + let write_succeeded = write_result.is_ok(); + + let (result, flush_timing) = match write_result { + Ok(()) => { + let flush_start = TokioInstant::now(); + let flush_result = writer.flush().await; + let flush_elapsed = flush_start.elapsed(); + let flush_succeeded = flush_result.is_ok(); + (flush_result, Some((flush_elapsed, flush_succeeded))) + } + Err(error) => (Err(error), None), + }; // Caller may have dropped the ack receiver (e.g. their // `await` was cancelled); that's fine — we still completed // the write, which was the whole point. let _ = ack.send(result); + + if let Some(trace) = &reverse_rpc { + trace.record_phase("writer_queue", queue_elapsed, true); + trace.record_phase("write_all", write_elapsed, write_succeeded); + if let Some((flush_elapsed, flush_succeeded)) = flush_timing { + trace.record_phase("flush", flush_elapsed, flush_succeeded); + } + } } } async fn read_loop( reader: impl AsyncRead + Unpin + Send, pending_requests: Arc>>, + reverse_requests: Arc>>, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, + timing_tx: Option>, ) { let mut reader = BufReader::new(reader); @@ -430,7 +728,24 @@ impl JsonRpcClient { let _ = notification_tx.send(notification); } JsonRpcMessage::Request(request) => { - if request_tx.send(request).is_err() { + let request_id = request.id; + let trace = timing_tx.as_ref().map(|timing_tx| { + ReverseRpcTrace::new( + &request, + TokioInstant::now(), + TokioInstant::now(), + timing_tx.clone(), + ) + }); + if let Some(trace) = &trace { + reverse_requests.write().insert(request_id, trace.clone()); + } + let forwarded = request_tx.send(request).is_ok(); + if let Some(trace) = &trace { + trace.record_forwarded(forwarded); + } + if !forwarded { + reverse_requests.write().remove(&request_id); warn!("failed to forward JSON-RPC request, channel closed"); } } @@ -455,6 +770,7 @@ impl JsonRpcClient { ); pending.clear(); } + reverse_requests.write().clear(); } async fn read_message( @@ -639,7 +955,29 @@ impl JsonRpcClient { /// drops the ack receiver; the actor still completes the frame and /// flushes. A partial frame can never appear on the wire. pub async fn write(&self, message: &T) -> Result<(), Error> { - let body = serde_json::to_vec(message)?; + self.write_frame(message, None).await + } + + pub(crate) async fn write_response(&self, response: &JsonRpcResponse) -> Result<(), Error> { + let trace = self.reverse_requests.read().get(&response.id).cloned(); + let result = self.write_frame(response, trace.clone()).await; + if let Some(trace) = &trace { + remove_reverse_request_if_same(&self.reverse_requests, response.id, trace); + } + result + } + + async fn write_frame( + &self, + message: &T, + reverse_rpc: Option, + ) -> Result<(), Error> { + let encode_start = TokioInstant::now(); + let encoded = serde_json::to_vec(message); + if let Some(trace) = &reverse_rpc { + trace.record_phase("response_encode", encode_start.elapsed(), encoded.is_ok()); + } + let body = encoded?; let mut frame = Vec::with_capacity(CONTENT_LENGTH_HEADER.len() + 16 + body.len() + 4); frame.extend_from_slice(CONTENT_LENGTH_HEADER.as_bytes()); frame.extend_from_slice(body.len().to_string().as_bytes()); @@ -647,8 +985,14 @@ impl JsonRpcClient { frame.extend_from_slice(&body); let (ack_tx, ack_rx) = oneshot::channel(); + let enqueued_at = TokioInstant::now(); self.write_tx - .send(WriteCommand { frame, ack: ack_tx }) + .send(WriteCommand { + frame, + ack: ack_tx, + reverse_rpc, + enqueued_at, + }) .map_err(|_| { Error::from(std::io::Error::new( std::io::ErrorKind::BrokenPipe, @@ -665,6 +1009,25 @@ impl JsonRpcClient { ))), } } + + pub(crate) fn trace_reverse_request_scheduled( + &self, + request_id: u64, + ) -> Option { + let trace = self.reverse_requests.read().get(&request_id).cloned(); + if let Some(trace) = &trace { + trace.record_scheduled(TokioInstant::now()); + } + trace.map(|trace| ReverseRpcDispatchGuard { + reverse_requests: self.reverse_requests.clone(), + request_id, + trace, + }) + } + + pub(crate) fn abandon_reverse_request(&self, request_id: u64) { + self.reverse_requests.write().remove(&request_id); + } } /// RAII guard that removes a pending-request entry from the map if the @@ -692,8 +1055,175 @@ impl Drop for PendingGuard<'_> { #[cfg(test)] mod tests { + use std::collections::VecDeque; + use std::future::Future; + use std::io::Write; + use std::pin::Pin; + use std::sync::Arc; + use std::task::{Context, Poll}; + use std::time::Duration; + + use parking_lot::Mutex; + use tokio::io::{AsyncWrite, AsyncWriteExt}; + use tokio::time::Sleep; + use tracing_subscriber::Layer; + use tracing_subscriber::fmt::MakeWriter; + use tracing_subscriber::layer::SubscriberExt; + use super::*; + #[derive(Clone, Default)] + struct TraceBuffer(Arc>>); + + impl TraceBuffer { + fn text(&self) -> String { + String::from_utf8(self.0.lock().clone()).unwrap() + } + } + + impl Write for TraceBuffer { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl<'a> MakeWriter<'a> for TraceBuffer { + type Writer = Self; + + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } + } + + struct DelayedWriter { + write_delays: VecDeque, + flush_delays: VecDeque, + write_sleep: Option>>, + flush_sleep: Option>>, + started_tx: mpsc::UnboundedSender<&'static str>, + } + + impl DelayedWriter { + fn new( + write_delays: [Duration; 2], + flush_delays: [Duration; 2], + ) -> (Self, mpsc::UnboundedReceiver<&'static str>) { + let (started_tx, started_rx) = mpsc::unbounded_channel(); + ( + Self { + write_delays: write_delays.into(), + flush_delays: flush_delays.into(), + write_sleep: None, + flush_sleep: None, + started_tx, + }, + started_rx, + ) + } + + fn poll_delay( + operation: &'static str, + delay: &mut VecDeque, + sleep: &mut Option>>, + started_tx: &mpsc::UnboundedSender<&'static str>, + cx: &mut Context<'_>, + ) -> Poll<()> { + if sleep.is_none() { + let duration = delay.pop_front().unwrap_or_default(); + let _ = started_tx.send(operation); + if duration.is_zero() { + return Poll::Ready(()); + } + *sleep = Some(Box::pin(tokio::time::sleep(duration))); + } + + match sleep + .as_mut() + .expect("delay sleep must exist") + .as_mut() + .poll(cx) + { + Poll::Ready(()) => { + *sleep = None; + Poll::Ready(()) + } + Poll::Pending => Poll::Pending, + } + } + } + + impl AsyncWrite for DelayedWriter { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let Self { + write_delays, + write_sleep, + started_tx, + .. + } = self.as_mut().get_mut(); + match Self::poll_delay("write", write_delays, write_sleep, started_tx, cx) { + Poll::Ready(()) => Poll::Ready(Ok(buf.len())), + Poll::Pending => Poll::Pending, + } + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let Self { + flush_delays, + flush_sleep, + started_tx, + .. + } = self.as_mut().get_mut(); + match Self::poll_delay("flush", flush_delays, flush_sleep, started_tx, cx) { + Poll::Ready(()) => Poll::Ready(Ok(())), + Poll::Pending => Poll::Pending, + } + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + fn trace_subscriber(buffer: TraceBuffer) -> impl tracing::Subscriber { + tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_writer(buffer) + .with_ansi(false) + .without_time() + .with_filter(tracing_subscriber::filter::filter_fn(|metadata| { + metadata.target() == REVERSE_RPC_TIMING_TARGET + })), + ) + } + + async fn wait_for_trace(buffer: &TraceBuffer, needle: &str) { + for _ in 0..20 { + if buffer.text().contains(needle) { + return; + } + tokio::task::yield_now().await; + } + panic!("timing trace did not contain {needle:?}: {}", buffer.text()); + } + + fn frame(message: &impl Serialize) -> Vec { + let body = serde_json::to_vec(message).unwrap(); + format!("Content-Length: {}\r\n\r\n", body.len()) + .into_bytes() + .into_iter() + .chain(body) + .collect() + } + #[test] fn deserialize_notification() { let json = r#"{"jsonrpc":"2.0","method":"session.event","params":{"id":"e1"}}"#; @@ -780,4 +1310,205 @@ mod tests { let json = serde_json::to_string(&r).unwrap(); assert!(!json.contains("error")); } + + #[test] + fn reverse_request_correlation_is_stable_and_opaque() { + let request = JsonRpcRequest::new( + 17, + "hooks.invoke", + Some(serde_json::json!({ "sessionId": "private-session-id" })), + ); + let same = ReverseRpcTrace::correlation_key( + request.id, + &request.method, + Some("private-session-id"), + ); + let repeated = ReverseRpcTrace::correlation_key( + request.id, + &request.method, + Some("private-session-id"), + ); + let different_session = + ReverseRpcTrace::correlation_key(request.id, &request.method, Some("other-session")); + + assert_eq!(same, repeated); + assert_ne!(same, different_session); + assert!(same.starts_with("rrpc-")); + assert!(!same.contains("private-session-id")); + } + + #[test] + fn reverse_request_guard_only_removes_its_own_generation() { + let (timing_tx, _timing_rx) = mpsc::unbounded_channel(); + let request = JsonRpcRequest::new(17, "hooks.invoke", None); + let now = TokioInstant::now(); + let first = ReverseRpcTrace::new(&request, now, now, timing_tx.clone()); + let second = ReverseRpcTrace::new(&request, now, now, timing_tx); + let reverse_requests = Arc::new(RwLock::new(HashMap::new())); + reverse_requests.write().insert(request.id, first.clone()); + let guard = ReverseRpcDispatchGuard { + reverse_requests: reverse_requests.clone(), + request_id: request.id, + trace: first, + }; + + reverse_requests.write().insert(request.id, second.clone()); + drop(guard); + + let retained = reverse_requests + .read() + .get(&request.id) + .cloned() + .expect("new request generation should remain tracked"); + assert!(Arc::ptr_eq(&retained.inner, &second.inner)); + } + + #[tokio::test] + async fn public_client_does_not_retain_reverse_request_timing_state() { + let (mut server, reader) = tokio::io::duplex(4096); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, mut request_rx) = mpsc::unbounded_channel(); + let client = JsonRpcClient::new(tokio::io::sink(), reader, notification_tx, request_tx); + let request = JsonRpcRequest::new(23, "consumer.request", None); + + server.write_all(&frame(&request)).await.unwrap(); + let forwarded = request_rx.recv().await.unwrap(); + + assert_eq!(forwarded.id, request.id); + assert!(client.reverse_requests.read().is_empty()); + client.force_close(); + } + + #[tokio::test(start_paused = true)] + async fn reverse_request_timing_tracks_gated_scheduling_without_content() { + const SENTINEL: &str = "PRIVATE_SENTINEL_DO_NOT_TRACE"; + + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (mut server, reader) = tokio::io::duplex(4096); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, mut request_rx) = mpsc::unbounded_channel(); + let client = JsonRpcClient::new_with_reverse_rpc_timing( + tokio::io::sink(), + reader, + notification_tx, + request_tx, + ); + let request = JsonRpcRequest::new( + 41, + "hooks.invoke", + Some(serde_json::json!({ + "sessionId": SENTINEL, + "hookType": "userPromptSubmitted", + "input": { + "prompt": SENTINEL, + "cwd": SENTINEL, + "toolArgs": { "secret": SENTINEL } + } + })), + ); + + server.write_all(&frame(&request)).await.unwrap(); + let forwarded = request_rx.recv().await.unwrap(); + tokio::time::advance(Duration::from_millis(13)).await; + + let trace = client + .trace_reverse_request_scheduled(forwarded.id) + .expect("reverse request timing should be tracked"); + trace + .trace() + .record_hook_callback("userPromptSubmitted", Duration::from_millis(3)); + client + .write_response(&JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: forwarded.id, + result: Some(serde_json::json!({ "output": SENTINEL })), + error: None, + }) + .await + .unwrap(); + + wait_for_trace(&trace_buffer, "phase=\"flush\"").await; + let output = trace_buffer.text(); + assert!(output.contains("github_copilot_sdk::reverse_rpc_timing")); + assert!(output.contains("phase=\"request_forward\"")); + assert!(output.contains("phase=\"request_schedule\"")); + assert!(output.contains("elapsed_us=13000")); + assert!(output.contains("phase=\"hook_callback\"")); + assert!(output.contains("phase=\"response_encode\"")); + assert!(output.contains("phase=\"writer_queue\"")); + assert!(output.contains("phase=\"write_all\"")); + assert!(output.contains("phase=\"flush\"")); + assert!(output.contains("correlation_key=rrpc-")); + assert!(!output.contains(SENTINEL)); + + client.force_close(); + } + + #[tokio::test(start_paused = true)] + async fn reverse_response_timing_distinguishes_writer_queue_write_and_flush() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (writer, mut started_rx) = DelayedWriter::new( + [Duration::from_millis(20), Duration::from_millis(7)], + [Duration::ZERO, Duration::from_millis(11)], + ); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, _request_rx) = mpsc::unbounded_channel(); + let client = JsonRpcClient::new(writer, tokio::io::empty(), notification_tx, request_tx); + + let (first_ack_tx, first_ack_rx) = oneshot::channel(); + client + .write_tx + .send(WriteCommand { + frame: frame(&serde_json::json!({})), + ack: first_ack_tx, + reverse_rpc: None, + enqueued_at: TokioInstant::now(), + }) + .unwrap(); + assert_eq!(started_rx.recv().await, Some("write")); + tokio::time::advance(Duration::from_millis(5)).await; + + let request = JsonRpcRequest::new( + 7, + "hooks.invoke", + Some(serde_json::json!({ "sessionId": "session" })), + ); + let now = TokioInstant::now(); + let trace = ReverseRpcTrace::for_test(&request, now, now); + let (second_ack_tx, second_ack_rx) = oneshot::channel(); + client + .write_tx + .send(WriteCommand { + frame: frame(&JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: 7, + result: Some(serde_json::json!({})), + error: None, + }), + ack: second_ack_tx, + reverse_rpc: Some(trace), + enqueued_at: TokioInstant::now(), + }) + .unwrap(); + + tokio::time::advance(Duration::from_millis(15)).await; + assert_eq!(started_rx.recv().await, Some("flush")); + assert_eq!(started_rx.recv().await, Some("write")); + tokio::time::advance(Duration::from_millis(7)).await; + assert_eq!(started_rx.recv().await, Some("flush")); + tokio::time::advance(Duration::from_millis(11)).await; + + first_ack_rx.await.unwrap().unwrap(); + second_ack_rx.await.unwrap().unwrap(); + + wait_for_trace(&trace_buffer, "phase=\"flush\"").await; + let output = trace_buffer.text(); + assert!(output.contains("phase=\"writer_queue\" elapsed_us=15000")); + assert!(output.contains("phase=\"write_all\" elapsed_us=7000")); + assert!(output.contains("phase=\"flush\" elapsed_us=11000")); + + client.force_close(); + } } diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 4a9f73ca4c..f3383e4bc4 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1431,6 +1431,7 @@ impl Client { Some(dispatcher.clone()), client.inner.on_github_telemetry.clone(), client.inner.github_token_registry.clone(), + Arc::downgrade(&client.inner), ); client.rpc().llm_inference().set_provider().await?; let llm_inference_elapsed = llm_inference_start.elapsed(); @@ -1601,7 +1602,7 @@ impl Client { let setup_start = Instant::now(); let (request_tx, request_rx) = mpsc::unbounded_channel::(); let (notification_broadcast_tx, _) = broadcast::channel::(1024); - let rpc = JsonRpcClient::new( + let rpc = JsonRpcClient::new_with_reverse_rpc_timing( writer, reader, notification_broadcast_tx.clone(), @@ -2025,7 +2026,18 @@ impl Client { /// Send a JSON-RPC response back to the CLI (e.g. for permission or tool call requests). pub(crate) async fn send_response(&self, response: &JsonRpcResponse) -> Result<()> { - self.inner.rpc.write(response).await + self.inner.rpc.write_response(response).await + } + + pub(crate) fn trace_reverse_request_scheduled( + &self, + request_id: u64, + ) -> Option { + self.inner.rpc.trace_reverse_request_scheduled(request_id) + } + + pub(crate) fn abandon_reverse_request(&self, request_id: u64) { + self.inner.rpc.abandon_reverse_request(request_id); } /// Reconstruct a [`Client`] handle from a shared inner pointer. @@ -2058,6 +2070,7 @@ impl Client { self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), self.inner.github_token_registry.clone(), + Arc::downgrade(&self.inner), ); self.inner.router.register(session_id) } @@ -2077,6 +2090,7 @@ impl Client { self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), self.inner.github_token_registry.clone(), + Arc::downgrade(&self.inner), ); let id = self.inner.github_token_registry.register(provider); github_token::GitHubTokenRegistration::new(self.inner.github_token_registry.clone(), id) @@ -2296,6 +2310,7 @@ impl Client { self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), self.inner.github_token_registry.clone(), + Arc::downgrade(&self.inner), ); } diff --git a/rust/src/router.rs b/rust/src/router.rs index 1dec9d16f3..8c9fb922c1 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -1,5 +1,5 @@ use std::collections::HashMap; -use std::sync::Arc; +use std::sync::{Arc, Weak}; use parking_lot::Mutex; use tokio::sync::{broadcast, mpsc}; @@ -7,6 +7,7 @@ use tracing::warn; use crate::jsonrpc::{JsonRpcNotification, JsonRpcRequest}; use crate::types::{SessionEventNotification, SessionId}; +use crate::{Client, ClientInner}; /// Per-session channels created by the router during session registration. pub(crate) struct SessionChannels { @@ -88,6 +89,7 @@ impl SessionRouter { llm_inference: Option>, github_telemetry: Option, github_token_registry: Arc, + client: Weak, ) { let mut started = self.started.lock(); if *started { @@ -192,10 +194,14 @@ impl SessionRouter { if let Some(dispatcher) = &llm_inference { dispatcher.dispatch(request).await; } else { + let request_id = request.id; warn!( method = %request.method, "llmInference request with no provider registered" ); + if let Some(inner) = client.upgrade() { + Client::from_inner(inner).abandon_reverse_request(request_id); + } } continue; } @@ -212,19 +218,31 @@ impl SessionRouter { guard.get(sid).map(|s| s.requests.clone()) }; if let Some(sender) = sender { - let _ = sender.send(request); + if let Err(error) = sender.send(request) + && let Some(inner) = client.upgrade() + { + Client::from_inner(inner).abandon_reverse_request(error.0.id); + } } else { + let request_id = request.id; warn!( session_id = sid, method = %request.method, "request for unregistered session" ); + if let Some(inner) = client.upgrade() { + Client::from_inner(inner).abandon_reverse_request(request_id); + } } } else { + let request_id = request.id; warn!( method = %request.method, "request missing sessionId" ); + if let Some(inner) = client.upgrade() { + Client::from_inner(inner).abandon_reverse_request(request_id); + } } } }); diff --git a/rust/src/session.rs b/rust/src/session.rs index b9d2173055..398a097ddd 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1609,6 +1609,9 @@ fn spawn_event_loop( .tx .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into())); } + while let Ok(request) = requests.try_recv() { + client.abandon_reverse_request(request.id); + } } .instrument(span), ) @@ -2350,6 +2353,7 @@ async fn handle_request( ) { let sid = session_id.clone(); let client = ctx.client; + let reverse_rpc_trace = client.trace_reverse_request_scheduled(request.id); let handlers = ctx.handlers; let hooks = ctx.hooks; let transforms = ctx.transforms; @@ -2385,7 +2389,15 @@ async fn handle_request( .unwrap_or(Value::Object(Default::default())); let rpc_result = if let Some(hooks) = hooks { - match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await { + match crate::hooks::dispatch_hook_traced( + hooks, + &sid, + hook_type, + input, + reverse_rpc_trace.as_ref().map(|guard| guard.trace()), + ) + .await + { Ok(output) => output, Err(e) => { warn!(error = %e, hook_type = hook_type, "hook dispatch failed"); From 91a413151efdb9e64f963e476563493b1d41d95a Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 28 Aug 2026 12:26:04 -0600 Subject: [PATCH 2/7] Harden reverse RPC correlation timing Capture the forwarding boundary explicitly and key request-derived correlation hashes per client to prevent cross-run linkability or offline session ID guessing. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/jsonrpc.rs | 148 ++++++++++++++++++++++++++++++++------------ 1 file changed, 110 insertions(+), 38 deletions(-) diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 9506c67ce6..816dd52425 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -1,5 +1,7 @@ -use std::collections::HashMap; +use std::collections::{HashMap, hash_map::RandomState}; +use std::hash::{BuildHasher, Hash, Hasher}; use std::sync::Arc; +use std::sync::OnceLock; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; @@ -289,7 +291,7 @@ struct ReverseRpcTraceInner { correlation_key: String, method: String, received_at: TokioInstant, - forwarded_at: TokioInstant, + forwarded_at: OnceLock, timing_tx: mpsc::UnboundedSender, } @@ -297,7 +299,7 @@ impl ReverseRpcTrace { fn new( request: &JsonRpcRequest, received_at: TokioInstant, - forwarded_at: TokioInstant, + correlation_hasher: &RandomState, timing_tx: mpsc::UnboundedSender, ) -> Self { let session_id = request @@ -307,10 +309,15 @@ impl ReverseRpcTrace { .and_then(Value::as_str); Self { inner: Arc::new(ReverseRpcTraceInner { - correlation_key: Self::correlation_key(request.id, &request.method, session_id), + correlation_key: Self::correlation_key( + correlation_hasher, + request.id, + &request.method, + session_id, + ), method: request.method.clone(), received_at, - forwarded_at, + forwarded_at: OnceLock::new(), timing_tx, }), } @@ -324,43 +331,59 @@ impl ReverseRpcTrace { ) -> Self { let (timing_tx, timing_rx) = mpsc::unbounded_channel(); tokio::spawn(JsonRpcClient::timing_loop(timing_rx)); - Self::new(request, received_at, forwarded_at, timing_tx) + let trace = Self::new(request, received_at, &RandomState::new(), timing_tx); + trace.mark_forwarding(forwarded_at); + trace } - fn correlation_key(request_id: u64, method: &str, session_id: Option<&str>) -> String { - let mut hash = 0xcbf29ce484222325_u64; - for byte in session_id - .unwrap_or("") - .bytes() - .chain([0xff]) - .chain(method.bytes()) - .chain([0xfe]) - .chain(request_id.to_le_bytes()) - { - hash ^= u64::from(byte); - hash = hash.wrapping_mul(0x100000001b3); - } - format!("rrpc-{hash:016x}") + fn correlation_key( + correlation_hasher: &RandomState, + request_id: u64, + method: &str, + session_id: Option<&str>, + ) -> String { + // A per-client keyed hash keeps the request-derived key stable for all + // phases without making custom session IDs guessable from trace output. + let mut hasher = correlation_hasher.build_hasher(); + session_id.unwrap_or("").hash(&mut hasher); + method.hash(&mut hasher); + request_id.hash(&mut hasher); + format!("rrpc-{:016x}", hasher.finish()) } fn elapsed_us(duration: std::time::Duration) -> u64 { u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) } + fn mark_forwarding(&self, forwarded_at: TokioInstant) { + self.inner + .forwarded_at + .set(forwarded_at) + .expect("forwarding timestamp must be recorded exactly once"); + } + fn record_forwarded(&self, succeeded: bool) { + let forwarded_at = self + .inner + .forwarded_at + .get() + .expect("forwarding timestamp must be set before forwarding"); self.record_phase( "request_forward", - self.inner - .forwarded_at - .duration_since(self.inner.received_at), + forwarded_at.duration_since(self.inner.received_at), succeeded, ); } fn record_scheduled(&self, scheduled_at: TokioInstant) { + let forwarded_at = self + .inner + .forwarded_at + .get() + .expect("forwarding timestamp must be set before scheduling"); let _ = self.inner.timing_tx.send(ReverseRpcTimingEvent::Scheduled { trace: self.clone(), - elapsed_us: Self::elapsed_us(scheduled_at.duration_since(self.inner.forwarded_at)), + elapsed_us: Self::elapsed_us(scheduled_at.duration_since(*forwarded_at)), since_receive_us: Self::elapsed_us(scheduled_at.duration_since(self.inner.received_at)), }); } @@ -482,14 +505,15 @@ impl JsonRpcClient { let writer_span = tracing::error_span!("jsonrpc_write_loop"); let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span)); - let (timing_tx, timing_task) = if trace_reverse_rpc { + let (timing_tx, timing_task, correlation_hasher) = if trace_reverse_rpc { let (timing_tx, timing_rx) = mpsc::unbounded_channel::(); ( Some(timing_tx), Some(tokio::spawn(Self::timing_loop(timing_rx))), + Some(RandomState::new()), ) } else { - (None, None) + (None, None, None) }; let client = Self { @@ -519,6 +543,7 @@ impl JsonRpcClient { notification_tx_clone, request_tx_clone, timing_tx, + correlation_hasher, ) .await; } @@ -664,6 +689,7 @@ impl JsonRpcClient { notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, timing_tx: Option>, + correlation_hasher: Option, ) { let mut reader = BufReader::new(reader); @@ -729,16 +755,19 @@ impl JsonRpcClient { } JsonRpcMessage::Request(request) => { let request_id = request.id; - let trace = timing_tx.as_ref().map(|timing_tx| { - ReverseRpcTrace::new( - &request, - TokioInstant::now(), - TokioInstant::now(), - timing_tx.clone(), - ) - }); + let trace = timing_tx.as_ref().zip(correlation_hasher.as_ref()).map( + |(timing_tx, correlation_hasher)| { + ReverseRpcTrace::new( + &request, + TokioInstant::now(), + correlation_hasher, + timing_tx.clone(), + ) + }, + ); if let Some(trace) = &trace { reverse_requests.write().insert(request_id, trace.clone()); + trace.mark_forwarding(TokioInstant::now()); } let forwarded = request_tx.send(request).is_ok(); if let Some(trace) = &trace { @@ -1318,18 +1347,25 @@ mod tests { "hooks.invoke", Some(serde_json::json!({ "sessionId": "private-session-id" })), ); + let correlation_hasher = RandomState::new(); let same = ReverseRpcTrace::correlation_key( + &correlation_hasher, request.id, &request.method, Some("private-session-id"), ); let repeated = ReverseRpcTrace::correlation_key( + &correlation_hasher, request.id, &request.method, Some("private-session-id"), ); - let different_session = - ReverseRpcTrace::correlation_key(request.id, &request.method, Some("other-session")); + let different_session = ReverseRpcTrace::correlation_key( + &correlation_hasher, + request.id, + &request.method, + Some("other-session"), + ); assert_eq!(same, repeated); assert_ne!(same, different_session); @@ -1342,8 +1378,9 @@ mod tests { let (timing_tx, _timing_rx) = mpsc::unbounded_channel(); let request = JsonRpcRequest::new(17, "hooks.invoke", None); let now = TokioInstant::now(); - let first = ReverseRpcTrace::new(&request, now, now, timing_tx.clone()); - let second = ReverseRpcTrace::new(&request, now, now, timing_tx); + let correlation_hasher = RandomState::new(); + let first = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing_tx.clone()); + let second = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing_tx); let reverse_requests = Arc::new(RwLock::new(HashMap::new())); reverse_requests.write().insert(request.id, first.clone()); let guard = ReverseRpcDispatchGuard { @@ -1363,6 +1400,41 @@ mod tests { assert!(Arc::ptr_eq(&retained.inner, &second.inner)); } + #[tokio::test(start_paused = true)] + async fn reverse_request_timing_uses_the_forwarding_boundary() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (timing_tx, timing_rx) = mpsc::unbounded_channel(); + tokio::spawn(JsonRpcClient::timing_loop(timing_rx)); + let request = JsonRpcRequest::new( + 29, + "hooks.invoke", + Some(serde_json::json!({ "sessionId": "session" })), + ); + let received_at = TokioInstant::now(); + let trace = ReverseRpcTrace::new(&request, received_at, &RandomState::new(), timing_tx); + + tokio::time::advance(Duration::from_millis(5)).await; + trace.mark_forwarding(TokioInstant::now()); + trace.record_forwarded(true); + tokio::time::advance(Duration::from_millis(7)).await; + trace.record_scheduled(TokioInstant::now()); + + wait_for_trace(&trace_buffer, "phase=\"request_schedule\"").await; + let output = trace_buffer.text(); + let forward = output + .lines() + .find(|line| line.contains("phase=\"request_forward\"")) + .expect("request_forward timing should be emitted"); + let schedule = output + .lines() + .find(|line| line.contains("phase=\"request_schedule\"")) + .expect("request_schedule timing should be emitted"); + assert!(forward.contains("elapsed_us=5000")); + assert!(schedule.contains("elapsed_us=7000")); + assert!(schedule.contains("since_receive_us=12000")); + } + #[tokio::test] async fn public_client_does_not_retain_reverse_request_timing_state() { let (mut server, reader) = tokio::io::duplex(4096); From 737f5f899fcf957ec762958df66322b723d8226c Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 28 Aug 2026 14:04:55 -0600 Subject: [PATCH 3/7] Gate and bound reverse RPC timing Make reverse-RPC diagnostics opt-in at the tracing target, preserve reconstructible phase placement, and bound timing delivery without affecting response behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/hooks.rs | 3 +- rust/src/jsonrpc.rs | 960 ++++++++++++++++++++++++++++++++++++++------ 2 files changed, 832 insertions(+), 131 deletions(-) diff --git a/rust/src/hooks.rs b/rust/src/hooks.rs index db79400c65..2574ad6dc5 100644 --- a/rust/src/hooks.rs +++ b/rust/src/hooks.rs @@ -756,7 +756,7 @@ pub(crate) async fn dispatch_hook_traced( let output = hooks.on_hook(event).await; let dispatch_elapsed = dispatch_start.elapsed(); if let Some(trace) = reverse_rpc_trace { - trace.record_hook_callback(hook_type, dispatch_elapsed); + trace.record_hook_callback(hook_type, dispatch_start, dispatch_elapsed); } tracing::debug!( elapsed_ms = dispatch_elapsed.as_millis(), @@ -983,6 +983,7 @@ mod tests { wait_for_trace(&trace_buffer, "phase=\"hook_callback\"").await; let traces = trace_buffer.text(); assert!(traces.contains("phase=\"hook_callback\"")); + assert!(traces.contains("start_offset_us=0")); assert!(traces.contains("elapsed_us=9000")); assert!(!traces.contains(SENTINEL)); } diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 816dd52425..79dacc2bee 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -1,4 +1,4 @@ -use std::collections::{HashMap, hash_map::RandomState}; +use std::collections::{HashMap, VecDeque, hash_map::RandomState}; use std::hash::{BuildHasher, Hash, Hasher}; use std::sync::Arc; use std::sync::OnceLock; @@ -178,8 +178,12 @@ const CONTENT_LENGTH_HEADER: &str = "Content-Length: "; /// the combined interval), `response_encode` is `serde_json::to_vec`, /// `writer_queue` is enqueue to dequeue, and `write_all` / `flush` measure /// the corresponding `AsyncWrite` calls. `hook_callback` is emitted by the -/// hooks dispatcher around `SessionHooks::on_hook` only. +/// hooks dispatcher around `SessionHooks::on_hook` only. Every request phase +/// includes its start offset from request receipt, and `request_complete` +/// records total elapsed time. Collection is enabled only when this target +/// has an active DEBUG subscriber. const REVERSE_RPC_TIMING_TARGET: &str = "github_copilot_sdk::reverse_rpc_timing"; +const REVERSE_RPC_TIMING_CAPACITY: usize = 256; /// Rewrites unpaired UTF-16 surrogate escapes to `\uFFFD`. /// @@ -254,27 +258,91 @@ fn repair_lone_surrogates(body: &[u8]) -> Option> { struct WriteCommand { frame: Vec, ack: oneshot::Sender>, - reverse_rpc: Option, + reverse_rpc: Option, enqueued_at: TokioInstant, } +struct ReverseRpcWriteTrace { + trace: ReverseRpcTrace, + completed: bool, +} + +impl ReverseRpcWriteTrace { + fn new(trace: ReverseRpcTrace) -> Self { + Self { + trace, + completed: false, + } + } + + fn record_complete(&mut self, completed_at: TokioInstant, succeeded: bool) { + self.completed = self.trace.record_complete(completed_at, succeeded); + } +} + +impl Drop for ReverseRpcWriteTrace { + fn drop(&mut self) { + if !self.completed { + let _ = self.trace.record_complete(TokioInstant::now(), false); + } + } +} + enum ReverseRpcTimingEvent { Phase { trace: ReverseRpcTrace, phase: &'static str, + start_offset_us: u64, elapsed_us: u64, succeeded: bool, }, Scheduled { trace: ReverseRpcTrace, + start_offset_us: u64, elapsed_us: u64, since_receive_us: u64, }, HookCallback { trace: ReverseRpcTrace, hook_type: String, + start_offset_us: u64, elapsed_us: u64, }, + Complete { + trace: ReverseRpcTrace, + required_phase_records: u64, + elapsed_us: u64, + succeeded: bool, + }, +} + +#[derive(Clone)] +struct ReverseRpcTimingEmitter { + phase_tx: mpsc::Sender, + terminal_tx: mpsc::Sender, + dropped_records: Arc, +} + +impl ReverseRpcTimingEmitter { + fn emit(&self, event: ReverseRpcTimingEvent) -> bool { + let result = if matches!(&event, ReverseRpcTimingEvent::Complete { .. }) { + self.terminal_tx.try_send(event) + } else { + self.phase_tx.try_send(event) + }; + match result { + Ok(()) => true, + Err(mpsc::error::TrySendError::Full(_)) => { + let _ = self.dropped_records.fetch_update( + Ordering::Relaxed, + Ordering::Relaxed, + |count| Some(count.saturating_add(1)), + ); + false + } + Err(mpsc::error::TrySendError::Closed(_)) => false, + } + } } /// Internal, content-free timing context for one inbound JSON-RPC request. @@ -292,7 +360,24 @@ struct ReverseRpcTraceInner { method: String, received_at: TokioInstant, forwarded_at: OnceLock, - timing_tx: mpsc::UnboundedSender, + timing: ReverseRpcTimingEmitter, + timing_state: Mutex, + emitted_phase_records: AtomicU64, +} + +#[derive(Clone, Copy)] +struct PendingReverseRpcCompletion { + required_phase_records: u64, + elapsed_us: u64, + succeeded: bool, +} + +#[derive(Default)] +struct ReverseRpcTimingState { + accepted_phase_records: u64, + pending_completion: Option, + completion_recorded: bool, + writer_owns_completion: bool, } impl ReverseRpcTrace { @@ -300,7 +385,7 @@ impl ReverseRpcTrace { request: &JsonRpcRequest, received_at: TokioInstant, correlation_hasher: &RandomState, - timing_tx: mpsc::UnboundedSender, + timing: ReverseRpcTimingEmitter, ) -> Self { let session_id = request .params @@ -318,7 +403,9 @@ impl ReverseRpcTrace { method: request.method.clone(), received_at, forwarded_at: OnceLock::new(), - timing_tx, + timing, + timing_state: Mutex::new(ReverseRpcTimingState::default()), + emitted_phase_records: AtomicU64::new(0), }), } } @@ -329,9 +416,24 @@ impl ReverseRpcTrace { received_at: TokioInstant, forwarded_at: TokioInstant, ) -> Self { - let (timing_tx, timing_rx) = mpsc::unbounded_channel(); - tokio::spawn(JsonRpcClient::timing_loop(timing_rx)); - let trace = Self::new(request, received_at, &RandomState::new(), timing_tx); + let (phase_tx, phase_rx) = mpsc::channel(REVERSE_RPC_TIMING_CAPACITY); + let (terminal_tx, terminal_rx) = mpsc::channel(REVERSE_RPC_TIMING_CAPACITY); + let dropped_records = Arc::new(AtomicU64::new(0)); + tokio::spawn(JsonRpcClient::timing_loop( + phase_rx, + terminal_rx, + dropped_records.clone(), + )); + let trace = Self::new( + request, + received_at, + &RandomState::new(), + ReverseRpcTimingEmitter { + phase_tx, + terminal_tx, + dropped_records, + }, + ); trace.mark_forwarding(forwarded_at); trace } @@ -355,6 +457,10 @@ impl ReverseRpcTrace { u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) } + fn start_offset_us(&self, started_at: TokioInstant) -> u64 { + Self::elapsed_us(started_at.duration_since(self.inner.received_at)) + } + fn mark_forwarding(&self, forwarded_at: TokioInstant) { self.inner .forwarded_at @@ -370,6 +476,7 @@ impl ReverseRpcTrace { .expect("forwarding timestamp must be set before forwarding"); self.record_phase( "request_forward", + self.inner.received_at, forwarded_at.duration_since(self.inner.received_at), succeeded, ); @@ -381,32 +488,102 @@ impl ReverseRpcTrace { .forwarded_at .get() .expect("forwarding timestamp must be set before scheduling"); - let _ = self.inner.timing_tx.send(ReverseRpcTimingEvent::Scheduled { + self.emit_phase(ReverseRpcTimingEvent::Scheduled { trace: self.clone(), + start_offset_us: self.start_offset_us(*forwarded_at), elapsed_us: Self::elapsed_us(scheduled_at.duration_since(*forwarded_at)), since_receive_us: Self::elapsed_us(scheduled_at.duration_since(self.inner.received_at)), }); } - pub(crate) fn record_hook_callback(&self, hook_type: &str, elapsed: std::time::Duration) { - let _ = self - .inner - .timing_tx - .send(ReverseRpcTimingEvent::HookCallback { - trace: self.clone(), - hook_type: hook_type.to_string(), - elapsed_us: Self::elapsed_us(elapsed), - }); + pub(crate) fn record_hook_callback( + &self, + hook_type: &str, + started_at: TokioInstant, + elapsed: std::time::Duration, + ) { + self.emit_phase(ReverseRpcTimingEvent::HookCallback { + trace: self.clone(), + hook_type: hook_type.to_string(), + start_offset_us: self.start_offset_us(started_at), + elapsed_us: Self::elapsed_us(elapsed), + }); } - fn record_phase(&self, phase: &'static str, elapsed: std::time::Duration, succeeded: bool) { - let _ = self.inner.timing_tx.send(ReverseRpcTimingEvent::Phase { + fn record_phase( + &self, + phase: &'static str, + started_at: TokioInstant, + elapsed: std::time::Duration, + succeeded: bool, + ) { + self.emit_phase(ReverseRpcTimingEvent::Phase { trace: self.clone(), phase, + start_offset_us: self.start_offset_us(started_at), elapsed_us: Self::elapsed_us(elapsed), succeeded, }); } + + fn emit_phase(&self, event: ReverseRpcTimingEvent) { + let mut state = self.inner.timing_state.lock(); + if state.pending_completion.is_some() || state.completion_recorded { + return; + } + if self.inner.timing.emit(event) { + state.accepted_phase_records = state.accepted_phase_records.saturating_add(1); + } + } + + fn record_complete(&self, completed_at: TokioInstant, succeeded: bool) -> bool { + let mut state = self.inner.timing_state.lock(); + self.record_complete_with_state(&mut state, completed_at, succeeded) + } + + fn record_abandoned(&self, completed_at: TokioInstant) { + let mut state = self.inner.timing_state.lock(); + if !state.writer_owns_completion { + let _ = self.record_complete_with_state(&mut state, completed_at, false); + } + } + + fn record_complete_with_state( + &self, + state: &mut ReverseRpcTimingState, + completed_at: TokioInstant, + succeeded: bool, + ) -> bool { + if state.completion_recorded { + return true; + } + if state.pending_completion.is_none() { + state.pending_completion = Some(PendingReverseRpcCompletion { + required_phase_records: state.accepted_phase_records, + elapsed_us: Self::elapsed_us(completed_at.duration_since(self.inner.received_at)), + succeeded, + }); + } + let completion = state + .pending_completion + .expect("pending completion must be initialized"); + if self.inner.timing.emit(ReverseRpcTimingEvent::Complete { + trace: self.clone(), + required_phase_records: completion.required_phase_records, + elapsed_us: completion.elapsed_us, + succeeded: completion.succeeded, + }) { + state.pending_completion = None; + state.completion_recorded = true; + true + } else { + false + } + } + + fn transfer_completion_to_writer(&self) { + self.inner.timing_state.lock().writer_owns_completion = true; + } } pub(crate) struct ReverseRpcDispatchGuard { @@ -423,7 +600,9 @@ impl ReverseRpcDispatchGuard { impl Drop for ReverseRpcDispatchGuard { fn drop(&mut self) { - remove_reverse_request_if_same(&self.reverse_requests, self.request_id, &self.trace); + if remove_reverse_request_if_same(&self.reverse_requests, self.request_id, &self.trace) { + self.trace.record_abandoned(TokioInstant::now()); + } } } @@ -431,13 +610,16 @@ fn remove_reverse_request_if_same( reverse_requests: &RwLock>, request_id: u64, trace: &ReverseRpcTrace, -) { +) -> bool { let mut reverse_requests = reverse_requests.write(); if reverse_requests .get(&request_id) .is_some_and(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) { reverse_requests.remove(&request_id); + true + } else { + false } } @@ -505,11 +687,23 @@ impl JsonRpcClient { let writer_span = tracing::error_span!("jsonrpc_write_loop"); let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span)); - let (timing_tx, timing_task, correlation_hasher) = if trace_reverse_rpc { - let (timing_tx, timing_rx) = mpsc::unbounded_channel::(); + let (timing, timing_task, correlation_hasher) = if trace_reverse_rpc { + let (phase_tx, phase_rx) = + mpsc::channel::(REVERSE_RPC_TIMING_CAPACITY); + let (terminal_tx, terminal_rx) = + mpsc::channel::(REVERSE_RPC_TIMING_CAPACITY); + let dropped_records = Arc::new(AtomicU64::new(0)); ( - Some(timing_tx), - Some(tokio::spawn(Self::timing_loop(timing_rx))), + Some(ReverseRpcTimingEmitter { + phase_tx, + terminal_tx, + dropped_records: dropped_records.clone(), + }), + Some(tokio::spawn(Self::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + ))), Some(RandomState::new()), ) } else { @@ -542,7 +736,7 @@ impl JsonRpcClient { reverse_requests, notification_tx_clone, request_tx_clone, - timing_tx, + timing, correlation_hasher, ) .await; @@ -561,71 +755,201 @@ impl JsonRpcClient { if let Some(task) = self.write_task.lock().take() { task.abort(); } - if let Some(task) = self.timing_task.lock().take() { - task.abort(); - } self.pending_requests.write().clear(); - self.reverse_requests.write().clear(); + let abandoned = self + .reverse_requests + .write() + .drain() + .map(|(_, trace)| trace) + .collect::>(); + for trace in abandoned { + trace.record_abandoned(TokioInstant::now()); + } + // Detach the timing task so it can drain the bounded queue. The + // aborted read/write tasks drop the remaining senders, so it exits + // once those final diagnostics are emitted. + let _ = self.timing_task.lock().take(); } - async fn timing_loop(mut rx: mpsc::UnboundedReceiver) { - while let Some(event) = rx.recv().await { - match event { - ReverseRpcTimingEvent::Phase { + async fn timing_loop( + mut phase_rx: mpsc::Receiver, + mut terminal_rx: mpsc::Receiver, + dropped_records: Arc, + ) { + let mut phase_closed = false; + let mut terminal_closed = false; + let mut pending_terminals = VecDeque::new(); + while !phase_closed || !terminal_closed { + Self::record_dropped_timing_records(&dropped_records); + tokio::select! { + biased; + event = terminal_rx.recv(), if !terminal_closed => { + if let Some(event) = event { + Self::record_or_defer_timing_event(event, &mut pending_terminals); + } else { + terminal_closed = true; + } + } + event = phase_rx.recv(), if !phase_closed => { + if let Some(event) = event { + Self::record_or_defer_timing_event(event, &mut pending_terminals); + } else { + phase_closed = true; + } + } + } + Self::record_dropped_timing_records(&dropped_records); + } + Self::record_dropped_timing_records(&dropped_records); + } + + fn record_or_defer_timing_event( + event: ReverseRpcTimingEvent, + pending_terminals: &mut VecDeque, + ) { + if let ReverseRpcTimingEvent::Complete { + trace, + required_phase_records, + .. + } = &event + && trace.inner.emitted_phase_records.load(Ordering::Acquire) < *required_phase_records + { + pending_terminals.push_back(event); + return; + } + + let phase_trace = match &event { + ReverseRpcTimingEvent::Complete { .. } => None, + ReverseRpcTimingEvent::Phase { trace, .. } + | ReverseRpcTimingEvent::Scheduled { trace, .. } + | ReverseRpcTimingEvent::HookCallback { trace, .. } => Some(trace.clone()), + }; + Self::record_timing_event(event); + if let Some(trace) = phase_trace { + trace + .inner + .emitted_phase_records + .fetch_add(1, Ordering::Release); + } + + let mut index = 0; + while index < pending_terminals.len() { + let ready = match &pending_terminals[index] { + ReverseRpcTimingEvent::Complete { trace, - phase, - elapsed_us, - succeeded, + required_phase_records, + .. } => { - debug!( - target: REVERSE_RPC_TIMING_TARGET, - parent: None, - correlation_key = %trace.inner.correlation_key, - rpc_method = %trace.inner.method, - phase, - elapsed_us, - status = if succeeded { "succeeded" } else { "failed" }, - "reverse JSON-RPC timing" - ); + trace.inner.emitted_phase_records.load(Ordering::Acquire) + >= *required_phase_records } - ReverseRpcTimingEvent::Scheduled { - trace, + _ => unreachable!("only terminal records are deferred"), + }; + if ready { + let terminal = pending_terminals + .remove(index) + .expect("pending terminal index should exist"); + Self::record_timing_event(terminal); + } else { + index += 1; + } + } + } + + fn record_timing_event(event: ReverseRpcTimingEvent) { + match event { + ReverseRpcTimingEvent::Phase { + trace, + phase, + start_offset_us, + elapsed_us, + succeeded, + } => { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + correlation_key = %trace.inner.correlation_key, + rpc_method = %trace.inner.method, + phase, + start_offset_us, + elapsed_us, + status = if succeeded { "succeeded" } else { "failed" }, + "reverse JSON-RPC timing" + ); + } + ReverseRpcTimingEvent::Scheduled { + trace, + start_offset_us, + elapsed_us, + since_receive_us, + } => { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + correlation_key = %trace.inner.correlation_key, + rpc_method = %trace.inner.method, + phase = "request_schedule", + start_offset_us, elapsed_us, since_receive_us, - } => { - debug!( - target: REVERSE_RPC_TIMING_TARGET, - parent: None, - correlation_key = %trace.inner.correlation_key, - rpc_method = %trace.inner.method, - phase = "request_schedule", - elapsed_us, - since_receive_us, - status = "succeeded", - "reverse JSON-RPC timing" - ); - } - ReverseRpcTimingEvent::HookCallback { - trace, + status = "succeeded", + "reverse JSON-RPC timing" + ); + } + ReverseRpcTimingEvent::HookCallback { + trace, + hook_type, + start_offset_us, + elapsed_us, + } => { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + correlation_key = %trace.inner.correlation_key, + rpc_method = %trace.inner.method, hook_type, + phase = "hook_callback", + start_offset_us, elapsed_us, - } => { - debug!( - target: REVERSE_RPC_TIMING_TARGET, - parent: None, - correlation_key = %trace.inner.correlation_key, - rpc_method = %trace.inner.method, - hook_type, - phase = "hook_callback", - elapsed_us, - status = "succeeded", - "reverse JSON-RPC timing" - ); - } + status = "succeeded", + "reverse JSON-RPC timing" + ); + } + ReverseRpcTimingEvent::Complete { + trace, + required_phase_records: _, + elapsed_us, + succeeded, + } => { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + correlation_key = %trace.inner.correlation_key, + rpc_method = %trace.inner.method, + phase = "request_complete", + start_offset_us = 0_u64, + elapsed_us, + status = if succeeded { "succeeded" } else { "failed" }, + "reverse JSON-RPC timing" + ); } } } + fn record_dropped_timing_records(dropped_records: &AtomicU64) { + let dropped_records = dropped_records.swap(0, Ordering::Relaxed); + if dropped_records > 0 { + debug!( + target: REVERSE_RPC_TIMING_TARGET, + parent: None, + phase = "records_dropped", + dropped_records, + status = "dropped", + "reverse JSON-RPC timing records dropped" + ); + } + } + /// Writer-actor task. Owns the `AsyncWrite`, drains the command queue, /// and writes each frame atomically (header + body + flush) before /// signaling the ack. @@ -645,7 +969,7 @@ impl JsonRpcClient { while let Some(WriteCommand { frame, ack, - reverse_rpc, + mut reverse_rpc, enqueued_at, }) = rx.recv().await { @@ -662,22 +986,29 @@ impl JsonRpcClient { let flush_result = writer.flush().await; let flush_elapsed = flush_start.elapsed(); let flush_succeeded = flush_result.is_ok(); - (flush_result, Some((flush_elapsed, flush_succeeded))) + ( + flush_result, + Some((flush_start, flush_elapsed, flush_succeeded)), + ) } Err(error) => (Err(error), None), }; + let completed_at = TokioInstant::now(); + let succeeded = result.is_ok(); // Caller may have dropped the ack receiver (e.g. their // `await` was cancelled); that's fine — we still completed // the write, which was the whole point. let _ = ack.send(result); - if let Some(trace) = &reverse_rpc { - trace.record_phase("writer_queue", queue_elapsed, true); - trace.record_phase("write_all", write_elapsed, write_succeeded); - if let Some((flush_elapsed, flush_succeeded)) = flush_timing { - trace.record_phase("flush", flush_elapsed, flush_succeeded); + if let Some(write_trace) = &mut reverse_rpc { + let trace = &write_trace.trace; + trace.record_phase("writer_queue", enqueued_at, queue_elapsed, true); + trace.record_phase("write_all", write_start, write_elapsed, write_succeeded); + if let Some((flush_start, flush_elapsed, flush_succeeded)) = flush_timing { + trace.record_phase("flush", flush_start, flush_elapsed, flush_succeeded); } + write_trace.record_complete(completed_at, succeeded); } } } @@ -688,7 +1019,7 @@ impl JsonRpcClient { reverse_requests: Arc>>, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, - timing_tx: Option>, + timing: Option, correlation_hasher: Option, ) { let mut reader = BufReader::new(reader); @@ -755,16 +1086,23 @@ impl JsonRpcClient { } JsonRpcMessage::Request(request) => { let request_id = request.id; - let trace = timing_tx.as_ref().zip(correlation_hasher.as_ref()).map( - |(timing_tx, correlation_hasher)| { - ReverseRpcTrace::new( - &request, - TokioInstant::now(), - correlation_hasher, - timing_tx.clone(), - ) - }, - ); + let trace = if tracing::enabled!( + target: REVERSE_RPC_TIMING_TARGET, + tracing::Level::DEBUG + ) { + timing.as_ref().zip(correlation_hasher.as_ref()).map( + |(timing, correlation_hasher)| { + ReverseRpcTrace::new( + &request, + TokioInstant::now(), + correlation_hasher, + timing.clone(), + ) + }, + ) + } else { + None + }; if let Some(trace) = &trace { reverse_requests.write().insert(request_id, trace.clone()); trace.mark_forwarding(TokioInstant::now()); @@ -774,7 +1112,9 @@ impl JsonRpcClient { trace.record_forwarded(forwarded); } if !forwarded { - reverse_requests.write().remove(&request_id); + if let Some(trace) = reverse_requests.write().remove(&request_id) { + trace.record_abandoned(TokioInstant::now()); + } warn!("failed to forward JSON-RPC request, channel closed"); } } @@ -799,7 +1139,14 @@ impl JsonRpcClient { ); pending.clear(); } - reverse_requests.write().clear(); + let abandoned = reverse_requests + .write() + .drain() + .map(|(_, trace)| trace) + .collect::>(); + for trace in abandoned { + trace.record_abandoned(TokioInstant::now()); + } } async fn read_message( @@ -991,7 +1338,7 @@ impl JsonRpcClient { let trace = self.reverse_requests.read().get(&response.id).cloned(); let result = self.write_frame(response, trace.clone()).await; if let Some(trace) = &trace { - remove_reverse_request_if_same(&self.reverse_requests, response.id, trace); + let _ = remove_reverse_request_if_same(&self.reverse_requests, response.id, trace); } result } @@ -1004,7 +1351,15 @@ impl JsonRpcClient { let encode_start = TokioInstant::now(); let encoded = serde_json::to_vec(message); if let Some(trace) = &reverse_rpc { - trace.record_phase("response_encode", encode_start.elapsed(), encoded.is_ok()); + trace.record_phase( + "response_encode", + encode_start, + encode_start.elapsed(), + encoded.is_ok(), + ); + if encoded.is_err() { + trace.record_complete(TokioInstant::now(), false); + } } let body = encoded?; let mut frame = Vec::with_capacity(CONTENT_LENGTH_HEADER.len() + 16 + body.len() + 4); @@ -1015,27 +1370,41 @@ impl JsonRpcClient { let (ack_tx, ack_rx) = oneshot::channel(); let enqueued_at = TokioInstant::now(); - self.write_tx + let response_trace = reverse_rpc.clone(); + if let Some(trace) = &response_trace { + trace.transfer_completion_to_writer(); + } + if self + .write_tx .send(WriteCommand { frame, ack: ack_tx, - reverse_rpc, + reverse_rpc: reverse_rpc.map(ReverseRpcWriteTrace::new), enqueued_at, }) - .map_err(|_| { - Error::from(std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "writer actor has shut down", - )) - })?; + .is_err() + { + if let Some(trace) = &response_trace { + trace.record_complete(TokioInstant::now(), false); + } + return Err(Error::from(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "writer actor has shut down", + ))); + } match ack_rx.await { Ok(Ok(())) => Ok(()), Ok(Err(e)) => Err(Error::from(e)), - Err(_) => Err(Error::from(std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "writer actor dropped ack without responding", - ))), + Err(_) => { + if let Some(trace) = &response_trace { + trace.record_complete(TokioInstant::now(), false); + } + Err(Error::from(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "writer actor dropped ack without responding", + ))) + } } } @@ -1055,7 +1424,9 @@ impl JsonRpcClient { } pub(crate) fn abandon_reverse_request(&self, request_id: u64) { - self.reverse_requests.write().remove(&request_id); + if let Some(trace) = self.reverse_requests.write().remove(&request_id) { + trace.record_abandoned(TokioInstant::now()); + } } } @@ -1086,7 +1457,7 @@ impl Drop for PendingGuard<'_> { mod tests { use std::collections::VecDeque; use std::future::Future; - use std::io::Write; + use std::io::{self, Write}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -1222,6 +1593,36 @@ mod tests { } } + #[derive(Clone)] + struct BlockingTraceWriter { + entered_tx: std::sync::mpsc::Sender<()>, + release: Arc<(std::sync::Mutex, std::sync::Condvar)>, + } + + impl Write for BlockingTraceWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + let _ = self.entered_tx.send(()); + let (released, condvar) = &*self.release; + let mut released = released.lock().unwrap(); + while !*released { + released = condvar.wait(released).unwrap(); + } + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + impl<'a> MakeWriter<'a> for BlockingTraceWriter { + type Writer = Self; + + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } + } + fn trace_subscriber(buffer: TraceBuffer) -> impl tracing::Subscriber { tracing_subscriber::registry().with( tracing_subscriber::fmt::layer() @@ -1234,6 +1635,29 @@ mod tests { ) } + fn timing_channel( + capacity: usize, + ) -> ( + ReverseRpcTimingEmitter, + mpsc::Receiver, + mpsc::Receiver, + Arc, + ) { + let (phase_tx, phase_rx) = mpsc::channel(capacity); + let (terminal_tx, terminal_rx) = mpsc::channel(capacity); + let dropped_records = Arc::new(AtomicU64::new(0)); + ( + ReverseRpcTimingEmitter { + phase_tx, + terminal_tx, + dropped_records: dropped_records.clone(), + }, + phase_rx, + terminal_rx, + dropped_records, + ) + } + async fn wait_for_trace(buffer: &TraceBuffer, needle: &str) { for _ in 0..20 { if buffer.text().contains(needle) { @@ -1375,12 +1799,12 @@ mod tests { #[test] fn reverse_request_guard_only_removes_its_own_generation() { - let (timing_tx, _timing_rx) = mpsc::unbounded_channel(); + let (timing, _phase_rx, _terminal_rx, _dropped_records) = timing_channel(1); let request = JsonRpcRequest::new(17, "hooks.invoke", None); let now = TokioInstant::now(); let correlation_hasher = RandomState::new(); - let first = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing_tx.clone()); - let second = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing_tx); + let first = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing.clone()); + let second = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing); let reverse_requests = Arc::new(RwLock::new(HashMap::new())); reverse_requests.write().insert(request.id, first.clone()); let guard = ReverseRpcDispatchGuard { @@ -1404,15 +1828,20 @@ mod tests { async fn reverse_request_timing_uses_the_forwarding_boundary() { let trace_buffer = TraceBuffer::default(); let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); - let (timing_tx, timing_rx) = mpsc::unbounded_channel(); - tokio::spawn(JsonRpcClient::timing_loop(timing_rx)); + let (timing, phase_rx, terminal_rx, dropped_records) = + timing_channel(REVERSE_RPC_TIMING_CAPACITY); + tokio::spawn(JsonRpcClient::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + )); let request = JsonRpcRequest::new( 29, "hooks.invoke", Some(serde_json::json!({ "sessionId": "session" })), ); let received_at = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, received_at, &RandomState::new(), timing_tx); + let trace = ReverseRpcTrace::new(&request, received_at, &RandomState::new(), timing); tokio::time::advance(Duration::from_millis(5)).await; trace.mark_forwarding(TokioInstant::now()); @@ -1431,7 +1860,9 @@ mod tests { .find(|line| line.contains("phase=\"request_schedule\"")) .expect("request_schedule timing should be emitted"); assert!(forward.contains("elapsed_us=5000")); + assert!(forward.contains("start_offset_us=0")); assert!(schedule.contains("elapsed_us=7000")); + assert!(schedule.contains("start_offset_us=5000")); assert!(schedule.contains("since_receive_us=12000")); } @@ -1451,6 +1882,73 @@ mod tests { client.force_close(); } + #[tokio::test] + async fn disabled_timing_target_does_not_allocate_reverse_request_state() { + let _subscriber = + tracing::subscriber::set_default(tracing::subscriber::NoSubscriber::default()); + let (mut server, reader) = tokio::io::duplex(4096); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, mut request_rx) = mpsc::unbounded_channel(); + let client = JsonRpcClient::new_with_reverse_rpc_timing( + tokio::io::sink(), + reader, + notification_tx, + request_tx, + ); + let request = JsonRpcRequest::new( + 29, + "hooks.invoke", + Some(serde_json::json!({ "sessionId": "session" })), + ); + + server.write_all(&frame(&request)).await.unwrap(); + let forwarded = request_rx.recv().await.unwrap(); + + assert_eq!(forwarded.id, request.id); + assert!(client.reverse_requests.read().is_empty()); + client.force_close(); + } + + #[tokio::test] + async fn saturated_timing_queue_drops_records_and_reports_the_count() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); + let request = JsonRpcRequest::new(37, "hooks.invoke", None); + let now = TokioInstant::now(); + let trace = ReverseRpcTrace::new(&request, now, &RandomState::new(), timing); + trace.mark_forwarding(now); + + trace.record_phase("first", now, Duration::ZERO, true); + trace.record_phase("second", now, Duration::ZERO, true); + trace.record_complete(now, true); + assert_eq!( + trace.inner.timing.dropped_records.load(Ordering::Relaxed), + 1 + ); + drop(trace); + tokio::spawn(JsonRpcClient::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + )); + + wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; + let output = trace_buffer.text(); + let dropped = output + .lines() + .find(|line| line.contains("phase=\"records_dropped\"")) + .expect("saturation diagnostic should be emitted"); + assert!(dropped.contains("dropped_records=1")); + assert!(output.contains("phase=\"first\"")); + assert!(!output.contains("phase=\"second\"")); + assert!(output.contains("phase=\"request_complete\"")); + assert!( + output.find("phase=\"first\"").unwrap() + < output.find("phase=\"request_complete\"").unwrap() + ); + } + #[tokio::test(start_paused = true)] async fn reverse_request_timing_tracks_gated_scheduling_without_content() { const SENTINEL: &str = "PRIVATE_SENTINEL_DO_NOT_TRACE"; @@ -1487,9 +1985,13 @@ mod tests { let trace = client .trace_reverse_request_scheduled(forwarded.id) .expect("reverse request timing should be tracked"); - trace - .trace() - .record_hook_callback("userPromptSubmitted", Duration::from_millis(3)); + let callback_start = TokioInstant::now(); + tokio::time::advance(Duration::from_millis(3)).await; + trace.trace().record_hook_callback( + "userPromptSubmitted", + callback_start, + callback_start.elapsed(), + ); client .write_response(&JsonRpcResponse { jsonrpc: "2.0".to_string(), @@ -1500,7 +2002,7 @@ mod tests { .await .unwrap(); - wait_for_trace(&trace_buffer, "phase=\"flush\"").await; + wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; let output = trace_buffer.text(); assert!(output.contains("github_copilot_sdk::reverse_rpc_timing")); assert!(output.contains("phase=\"request_forward\"")); @@ -1511,6 +2013,8 @@ mod tests { assert!(output.contains("phase=\"writer_queue\"")); assert!(output.contains("phase=\"write_all\"")); assert!(output.contains("phase=\"flush\"")); + assert!(output.contains("phase=\"request_complete\"")); + assert!(output.contains("start_offset_us=")); assert!(output.contains("correlation_key=rrpc-")); assert!(!output.contains(SENTINEL)); @@ -1560,7 +2064,7 @@ mod tests { error: None, }), ack: second_ack_tx, - reverse_rpc: Some(trace), + reverse_rpc: Some(ReverseRpcWriteTrace::new(trace)), enqueued_at: TokioInstant::now(), }) .unwrap(); @@ -1575,12 +2079,208 @@ mod tests { first_ack_rx.await.unwrap().unwrap(); second_ack_rx.await.unwrap().unwrap(); - wait_for_trace(&trace_buffer, "phase=\"flush\"").await; + wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; let output = trace_buffer.text(); - assert!(output.contains("phase=\"writer_queue\" elapsed_us=15000")); - assert!(output.contains("phase=\"write_all\" elapsed_us=7000")); - assert!(output.contains("phase=\"flush\" elapsed_us=11000")); + assert!(output.contains("phase=\"writer_queue\" start_offset_us=0 elapsed_us=15000")); + assert!(output.contains("phase=\"write_all\" start_offset_us=15000 elapsed_us=7000")); + assert!(output.contains("phase=\"flush\" start_offset_us=22000 elapsed_us=11000")); + assert!(output.contains("phase=\"request_complete\" start_offset_us=0 elapsed_us=33000")); + let writer_queue = output.find("phase=\"writer_queue\"").unwrap(); + let write_all = output.find("phase=\"write_all\"").unwrap(); + let flush = output.find("phase=\"flush\"").unwrap(); + let complete = output.find("phase=\"request_complete\"").unwrap(); + assert!(writer_queue < write_all); + assert!(write_all < flush); + assert!(flush < complete); + + client.force_close(); + } + + #[tokio::test(start_paused = true)] + async fn cancelled_response_keeps_the_writer_terminal_outcome() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (writer, mut started_rx) = DelayedWriter::new( + [Duration::from_millis(10), Duration::ZERO], + [Duration::ZERO, Duration::ZERO], + ); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, _request_rx) = mpsc::unbounded_channel(); + let (_server_guard, reader) = tokio::io::duplex(64); + let client = Arc::new(JsonRpcClient::new( + writer, + reader, + notification_tx, + request_tx, + )); + let request = JsonRpcRequest::new(53, "hooks.invoke", None); + let now = TokioInstant::now(); + let trace = ReverseRpcTrace::for_test(&request, now, now); + client + .reverse_requests + .write() + .insert(request.id, trace.clone()); + let dispatch_guard = ReverseRpcDispatchGuard { + reverse_requests: client.reverse_requests.clone(), + request_id: request.id, + trace, + }; + let response_task = tokio::spawn({ + let client = client.clone(); + async move { + client + .write_response(&JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: request.id, + result: Some(serde_json::json!({})), + error: None, + }) + .await + } + }); + + assert_eq!(started_rx.recv().await, Some("write")); + response_task.abort(); + let _ = response_task.await; + drop(dispatch_guard); + tokio::time::advance(Duration::from_millis(10)).await; + assert_eq!(started_rx.recv().await, Some("flush")); + + wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; + let complete = trace_buffer + .text() + .lines() + .filter(|line| line.contains("phase=\"request_complete\"")) + .map(str::to_owned) + .collect::>(); + assert_eq!(complete.len(), 1); + assert!(complete[0].contains("status=\"succeeded\"")); + assert!(client.reverse_requests.read().is_empty()); + + client.force_close(); + } + + #[tokio::test(start_paused = true)] + async fn force_close_emits_one_failed_terminal_record() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (writer, mut started_rx) = DelayedWriter::new( + [Duration::from_secs(60), Duration::ZERO], + [Duration::ZERO, Duration::ZERO], + ); + let (mut server, reader) = tokio::io::duplex(4096); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, mut request_rx) = mpsc::unbounded_channel(); + let client = Arc::new(JsonRpcClient::new_with_reverse_rpc_timing( + writer, + reader, + notification_tx, + request_tx, + )); + let request = JsonRpcRequest::new(59, "hooks.invoke", None); + server.write_all(&frame(&request)).await.unwrap(); + let forwarded = request_rx.recv().await.unwrap(); + let dispatch_guard = client + .trace_reverse_request_scheduled(forwarded.id) + .expect("enabled timing target should track the request"); + let response_task = tokio::spawn({ + let client = client.clone(); + async move { + client + .write_response(&JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: forwarded.id, + result: Some(serde_json::json!({})), + error: None, + }) + .await + } + }); + + assert_eq!(started_rx.recv().await, Some("write")); + client.force_close(); + assert!(response_task.await.unwrap().is_err()); + drop(dispatch_guard); + + wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; + let complete = trace_buffer + .text() + .lines() + .filter(|line| line.contains("phase=\"request_complete\"")) + .map(str::to_owned) + .collect::>(); + assert_eq!(complete.len(), 1); + assert!(complete[0].contains("status=\"failed\"")); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn slow_timing_subscriber_does_not_delay_response_ack() { + let (entered_tx, entered_rx) = std::sync::mpsc::channel(); + let release = Arc::new((std::sync::Mutex::new(false), std::sync::Condvar::new())); + let blocking_writer = BlockingTraceWriter { + entered_tx, + release: release.clone(), + }; + let (timing, phase_rx, terminal_rx, dropped_records) = + timing_channel(REVERSE_RPC_TIMING_CAPACITY); + let timing_thread = std::thread::spawn(move || { + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_writer(blocking_writer) + .with_ansi(false) + .without_time() + .with_filter(tracing_subscriber::filter::filter_fn(|metadata| { + metadata.target() == REVERSE_RPC_TIMING_TARGET + })), + ); + tracing::subscriber::with_default(subscriber, || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(JsonRpcClient::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + )); + }); + }); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, _request_rx) = mpsc::unbounded_channel(); + let client = JsonRpcClient::new( + tokio::io::sink(), + tokio::io::empty(), + notification_tx, + request_tx, + ); + let request = JsonRpcRequest::new(53, "hooks.invoke", None); + let now = TokioInstant::now(); + let trace = ReverseRpcTrace::new(&request, now, &RandomState::new(), timing); + trace.mark_forwarding(now); + + tokio::time::timeout( + Duration::from_secs(1), + client.write_frame( + &JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: request.id, + result: Some(serde_json::json!({})), + error: None, + }, + Some(trace), + ), + ) + .await + .expect("response acknowledgement should not wait for trace formatting") + .unwrap(); + entered_rx + .recv_timeout(Duration::from_secs(1)) + .expect("timing subscriber should be blocked after acknowledgement"); + + let (released, condvar) = &*release; + *released.lock().unwrap() = true; + condvar.notify_all(); client.force_close(); + timing_thread.join().unwrap(); } } From 0d89a78efd6e7eb26ea1687cd44b6748de61af6d Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 28 Aug 2026 14:40:20 -0600 Subject: [PATCH 4/7] Fix reverse timing opt-in controls Avoid constructing timing infrastructure when collection is disabled, preserve consumer warning-clean builds, and count saturated terminal records once. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/jsonrpc.rs | 166 ++++++++++++++++++++++++++++++++++++++------ 1 file changed, 146 insertions(+), 20 deletions(-) diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 79dacc2bee..639e6ab1ed 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -181,7 +181,7 @@ const CONTENT_LENGTH_HEADER: &str = "Content-Length: "; /// hooks dispatcher around `SessionHooks::on_hook` only. Every request phase /// includes its start offset from request receipt, and `request_complete` /// records total elapsed time. Collection is enabled only when this target -/// has an active DEBUG subscriber. +/// has an active DEBUG subscriber when the client is constructed. const REVERSE_RPC_TIMING_TARGET: &str = "github_copilot_sdk::reverse_rpc_timing"; const REVERSE_RPC_TIMING_CAPACITY: usize = 256; @@ -264,25 +264,26 @@ struct WriteCommand { struct ReverseRpcWriteTrace { trace: ReverseRpcTrace, - completed: bool, + completion_attempted: bool, } impl ReverseRpcWriteTrace { fn new(trace: ReverseRpcTrace) -> Self { Self { trace, - completed: false, + completion_attempted: false, } } fn record_complete(&mut self, completed_at: TokioInstant, succeeded: bool) { - self.completed = self.trace.record_complete(completed_at, succeeded); + self.completion_attempted = true; + let _ = self.trace.record_complete(completed_at, succeeded); } } impl Drop for ReverseRpcWriteTrace { fn drop(&mut self) { - if !self.completed { + if !self.completion_attempted { let _ = self.trace.record_complete(TokioInstant::now(), false); } } @@ -658,6 +659,13 @@ impl JsonRpcClient { /// messages to pending request channels, the notification broadcast, /// or the request-forwarding channel; and a writer actor that owns the /// underlying `AsyncWrite` and serializes frames atomically. + #[cfg_attr( + not(any(test, feature = "test-support")), + expect( + dead_code, + reason = "low-level constructor is exported only with test-support" + ) + )] pub fn new( writer: impl AsyncWrite + Unpin + Send + 'static, reader: impl AsyncRead + Unpin + Send + 'static, @@ -673,7 +681,17 @@ impl JsonRpcClient { notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, ) -> Self { - Self::new_inner(writer, reader, notification_tx, request_tx, true) + let trace_reverse_rpc = tracing::enabled!( + target: REVERSE_RPC_TIMING_TARGET, + tracing::Level::DEBUG + ); + Self::new_inner( + writer, + reader, + notification_tx, + request_tx, + trace_reverse_rpc, + ) } fn new_inner( @@ -1370,8 +1388,7 @@ impl JsonRpcClient { let (ack_tx, ack_rx) = oneshot::channel(); let enqueued_at = TokioInstant::now(); - let response_trace = reverse_rpc.clone(); - if let Some(trace) = &response_trace { + if let Some(trace) = &reverse_rpc { trace.transfer_completion_to_writer(); } if self @@ -1384,9 +1401,6 @@ impl JsonRpcClient { }) .is_err() { - if let Some(trace) = &response_trace { - trace.record_complete(TokioInstant::now(), false); - } return Err(Error::from(std::io::Error::new( std::io::ErrorKind::BrokenPipe, "writer actor has shut down", @@ -1396,15 +1410,10 @@ impl JsonRpcClient { match ack_rx.await { Ok(Ok(())) => Ok(()), Ok(Err(e)) => Err(Error::from(e)), - Err(_) => { - if let Some(trace) = &response_trace { - trace.record_complete(TokioInstant::now(), false); - } - Err(Error::from(std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "writer actor dropped ack without responding", - ))) - } + Err(_) => Err(Error::from(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "writer actor dropped ack without responding", + ))), } } @@ -1879,6 +1888,7 @@ mod tests { assert_eq!(forwarded.id, request.id); assert!(client.reverse_requests.read().is_empty()); + assert!(client.timing_task.lock().is_none()); client.force_close(); } @@ -1906,6 +1916,7 @@ mod tests { assert_eq!(forwarded.id, request.id); assert!(client.reverse_requests.read().is_empty()); + assert!(client.timing_task.lock().is_none()); client.force_close(); } @@ -1949,6 +1960,121 @@ mod tests { ); } + #[tokio::test] + async fn saturated_terminal_queue_counts_one_writer_completion_drop_once() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); + let now = TokioInstant::now(); + let filler_request = JsonRpcRequest::new(38, "filler.request", None); + let filler = + ReverseRpcTrace::new(&filler_request, now, &RandomState::new(), timing.clone()); + filler.mark_forwarding(now); + assert!(filler.record_complete(now, true)); + + let writer_request = JsonRpcRequest::new(39, "writer.request", None); + let writer = ReverseRpcTrace::new(&writer_request, now, &RandomState::new(), timing); + writer.mark_forwarding(now); + let mut writer_trace = ReverseRpcWriteTrace::new(writer); + writer_trace.record_complete(now, true); + drop(writer_trace); + + assert_eq!( + filler.inner.timing.dropped_records.load(Ordering::Relaxed), + 1 + ); + drop(filler); + tokio::spawn(JsonRpcClient::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + )); + + wait_for_trace(&trace_buffer, "phase=\"records_dropped\"").await; + let output = trace_buffer.text(); + let dropped = output + .lines() + .find(|line| line.contains("phase=\"records_dropped\"")) + .expect("terminal saturation diagnostic should be emitted"); + assert!(dropped.contains("dropped_records=1")); + assert!(output.contains("rpc_method=filler.request")); + assert!(!output.lines().any(|line| { + line.contains("rpc_method=writer.request") + && line.contains("phase=\"request_complete\"") + })); + } + + #[tokio::test] + async fn saturated_terminal_queue_counts_closed_writer_fallback_once() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); + let now = TokioInstant::now(); + let filler_request = JsonRpcRequest::new(40, "filler.request", None); + let filler = + ReverseRpcTrace::new(&filler_request, now, &RandomState::new(), timing.clone()); + filler.mark_forwarding(now); + assert!(filler.record_complete(now, true)); + + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, _request_rx) = mpsc::unbounded_channel(); + let client = JsonRpcClient::new( + tokio::io::sink(), + tokio::io::empty(), + notification_tx, + request_tx, + ); + let write_task = client + .write_task + .lock() + .take() + .expect("writer task should be running"); + write_task.abort(); + let _ = write_task.await; + + let writer_request = JsonRpcRequest::new(41, "writer.request", None); + let writer = ReverseRpcTrace::new(&writer_request, now, &RandomState::new(), timing); + writer.mark_forwarding(now); + let error = client + .write_frame( + &JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: writer_request.id, + result: Some(serde_json::json!({})), + error: None, + }, + Some(writer), + ) + .await + .unwrap_err(); + assert!(matches!(error.kind(), ErrorKind::Io)); + assert_eq!( + filler.inner.timing.dropped_records.load(Ordering::Relaxed), + 1 + ); + drop(filler); + tokio::spawn(JsonRpcClient::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + )); + + wait_for_trace(&trace_buffer, "phase=\"records_dropped\"").await; + let output = trace_buffer.text(); + let dropped = output + .lines() + .find(|line| line.contains("phase=\"records_dropped\"")) + .expect("closed-writer saturation diagnostic should be emitted"); + assert!(dropped.contains("dropped_records=1")); + assert!(output.contains("rpc_method=filler.request")); + assert!(!output.lines().any(|line| { + line.contains("rpc_method=writer.request") + && line.contains("phase=\"request_complete\"") + })); + + client.force_close(); + } + #[tokio::test(start_paused = true)] async fn reverse_request_timing_tracks_gated_scheduling_without_content() { const SENTINEL: &str = "PRIVATE_SENTINEL_DO_NOT_TRACE"; From d01d323f03829d86ca68975bd4e4bccd348ea727 Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 28 Aug 2026 16:33:22 -0600 Subject: [PATCH 5/7] Harden reverse RPC timing generations Bind timing to exact forwarded request generations, preserve forwarding order, and restrict trace labels to privacy-safe values. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/copilot_request_handler.rs | 15 +- rust/src/github_token.rs | 25 +- rust/src/jsonrpc.rs | 699 +++++++++++++++++++++------- rust/src/lib.rs | 23 +- rust/src/router.rs | 30 +- rust/src/session.rs | 20 +- 6 files changed, 599 insertions(+), 213 deletions(-) diff --git a/rust/src/copilot_request_handler.rs b/rust/src/copilot_request_handler.rs index d4bf2af1fb..df1cdcd904 100644 --- a/rust/src/copilot_request_handler.rs +++ b/rust/src/copilot_request_handler.rs @@ -1047,10 +1047,17 @@ impl CopilotRequestDispatcher { self.client.get().cloned().unwrap_or_else(Weak::new) } - pub(crate) async fn dispatch(self: &Arc, request: JsonRpcRequest) { - let _reverse_rpc = self - .client() - .and_then(|client| client.trace_reverse_request_scheduled(request.id)); + pub(crate) async fn dispatch(self: &Arc, request: crate::jsonrpc::ReverseRpcRequest) { + let (request, reverse_rpc) = request.into_dispatch(); + let dispatch = self.dispatch_inner(request); + if let Some(trace) = &reverse_rpc { + trace.scope(dispatch).await; + } else { + dispatch.await; + } + } + + async fn dispatch_inner(self: &Arc, request: JsonRpcRequest) { match request.method.as_str() { METHOD_HTTP_REQUEST_START => self.handle_start(request).await, METHOD_HTTP_REQUEST_CHUNK => self.handle_chunk(request).await, diff --git a/rust/src/github_token.rs b/rust/src/github_token.rs index 456c5a86b9..07a0ac586f 100644 --- a/rust/src/github_token.rs +++ b/rust/src/github_token.rs @@ -188,12 +188,21 @@ impl GitHubTokenRegistry { state.session_owners.clear(); } - pub(crate) async fn dispatch(&self, request: JsonRpcRequest) { + pub(crate) async fn dispatch(&self, request: crate::jsonrpc::ReverseRpcRequest) { let Some(inner) = self.client.get().and_then(Weak::upgrade) else { return; }; let client = Client::from_inner(inner); - let _reverse_rpc = client.trace_reverse_request_scheduled(request.id); + let (request, reverse_rpc) = request.into_dispatch(); + let dispatch = self.dispatch_inner(&client, request); + if let Some(trace) = &reverse_rpc { + trace.scope(dispatch).await; + } else { + dispatch.await; + } + } + + async fn dispatch_inner(&self, client: &Client, request: JsonRpcRequest) { let params = request .params .clone() @@ -202,7 +211,7 @@ impl GitHubTokenRegistry { Ok(params) => params, Err(error) => { send_error( - &client, + client, request.id, error_codes::INVALID_PARAMS, &format!("invalid params: {error}"), @@ -219,7 +228,7 @@ impl GitHubTokenRegistry { .cloned(); let Some(provider) = provider else { send_error( - &client, + client, request.id, error_codes::INTERNAL_ERROR, "unknown GitHub token provider registration", @@ -233,7 +242,7 @@ impl GitHubTokenRegistry { GitHubTokenAcquireReason::Refresh => GitHubTokenRequestReason::Refresh, GitHubTokenAcquireReason::Unknown => { send_error( - &client, + client, request.id, error_codes::INVALID_PARAMS, "unknown GitHub token acquisition reason", @@ -253,7 +262,7 @@ impl GitHubTokenRegistry { { Ok(GitHubTokenProviderResult::Token(token)) => { respond( - &client, + client, request.id, GitHubTokenAcquireResult::Token(token.into_wire()), ) @@ -261,7 +270,7 @@ impl GitHubTokenRegistry { } Ok(GitHubTokenProviderResult::Cancelled) => { respond( - &client, + client, request.id, GitHubTokenAcquireResult::Cancelled(GitHubTokenAcquireResultCancelled { kind: Default::default(), @@ -271,7 +280,7 @@ impl GitHubTokenRegistry { } Err(error) => { send_error( - &client, + client, request.id, error_codes::INTERNAL_ERROR, &format!("GitHub token provider failed: {error}"), diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index 639e6ab1ed..ddf1f2aad7 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -1,4 +1,5 @@ use std::collections::{HashMap, VecDeque, hash_map::RandomState}; +use std::future::Future; use std::hash::{BuildHasher, Hash, Hasher}; use std::sync::Arc; use std::sync::OnceLock; @@ -267,6 +268,17 @@ struct ReverseRpcWriteTrace { completion_attempted: bool, } +struct ReverseRpcResponseTrace { + request_id: u64, + trace: ReverseRpcTrace, +} + +// Response helpers retain their existing signatures while the dispatch task +// carries the exact request generation that produced the response. +tokio::task_local! { + static REVERSE_RPC_RESPONSE_TRACE: ReverseRpcResponseTrace; +} + impl ReverseRpcWriteTrace { fn new(trace: ReverseRpcTrace) -> Self { Self { @@ -305,7 +317,7 @@ enum ReverseRpcTimingEvent { }, HookCallback { trace: ReverseRpcTrace, - hook_type: String, + hook_type: &'static str, start_offset_us: u64, elapsed_us: u64, }, @@ -326,6 +338,9 @@ struct ReverseRpcTimingEmitter { impl ReverseRpcTimingEmitter { fn emit(&self, event: ReverseRpcTimingEvent) -> bool { + // Timing is measurement-only: a full bounded queue drops one logical + // record rather than delaying RPC work, including for terminal records. + // `records_dropped` is the explicit signal that the trace is incomplete. let result = if matches!(&event, ReverseRpcTimingEvent::Complete { .. }) { self.terminal_tx.try_send(event) } else { @@ -358,7 +373,7 @@ pub(crate) struct ReverseRpcTrace { struct ReverseRpcTraceInner { correlation_key: String, - method: String, + method: &'static str, received_at: TokioInstant, forwarded_at: OnceLock, timing: ReverseRpcTimingEmitter, @@ -384,6 +399,7 @@ struct ReverseRpcTimingState { impl ReverseRpcTrace { fn new( request: &JsonRpcRequest, + generation: u64, received_at: TokioInstant, correlation_hasher: &RandomState, timing: ReverseRpcTimingEmitter, @@ -400,8 +416,9 @@ impl ReverseRpcTrace { request.id, &request.method, session_id, + generation, ), - method: request.method.clone(), + method: Self::timing_method(&request.method), received_at, forwarded_at: OnceLock::new(), timing, @@ -427,6 +444,7 @@ impl ReverseRpcTrace { )); let trace = Self::new( request, + 0, received_at, &RandomState::new(), ReverseRpcTimingEmitter { @@ -444,6 +462,7 @@ impl ReverseRpcTrace { request_id: u64, method: &str, session_id: Option<&str>, + generation: u64, ) -> String { // A per-client keyed hash keeps the request-derived key stable for all // phases without making custom session IDs guessable from trace output. @@ -451,9 +470,42 @@ impl ReverseRpcTrace { session_id.unwrap_or("").hash(&mut hasher); method.hash(&mut hasher); request_id.hash(&mut hasher); + generation.hash(&mut hasher); format!("rrpc-{:016x}", hasher.finish()) } + fn timing_method(method: &str) -> &'static str { + match method { + "hooks.invoke" => "hooks.invoke", + "userInput.request" => "userInput.request", + "exitPlanMode.request" => "exitPlanMode.request", + "autoModeSwitch.request" => "autoModeSwitch.request", + "systemMessage.transform" => "systemMessage.transform", + "gitHubToken.getToken" => "gitHubToken.getToken", + "providerToken.getToken" => "providerToken.getToken", + _ if method.starts_with("sessionFs.") => "sessionFs.*", + _ if method.starts_with("canvas.") => "canvas.*", + _ if method.starts_with("llmInference.") => "llmInference.*", + _ => "unknown", + } + } + + fn timing_hook_type(hook_type: &str) -> &'static str { + match hook_type { + "preToolUse" => "preToolUse", + "preMcpToolCall" => "preMcpToolCall", + "postToolUse" => "postToolUse", + "postToolUseFailure" => "postToolUseFailure", + "userPromptSubmitted" => "userPromptSubmitted", + "userPromptTransformed" => "userPromptTransformed", + "sessionStart" => "sessionStart", + "sessionEnd" => "sessionEnd", + "errorOccurred" => "errorOccurred", + "agentStop" => "agentStop", + _ => "unknown", + } + } + fn elapsed_us(duration: std::time::Duration) -> u64 { u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) } @@ -469,18 +521,22 @@ impl ReverseRpcTrace { .expect("forwarding timestamp must be recorded exactly once"); } - fn record_forwarded(&self, succeeded: bool) { - let forwarded_at = self - .inner - .forwarded_at - .get() - .expect("forwarding timestamp must be set before forwarding"); - self.record_phase( - "request_forward", - self.inner.received_at, - forwarded_at.duration_since(self.inner.received_at), - succeeded, + fn forward(&self, send: impl FnOnce() -> Result<(), T>) -> Result<(), T> { + let forwarded_at = TokioInstant::now(); + self.mark_forwarding(forwarded_at); + let mut state = self.inner.timing_state.lock(); + let result = send(); + self.emit_phase_with_state( + &mut state, + ReverseRpcTimingEvent::Phase { + trace: self.clone(), + phase: "request_forward", + start_offset_us: 0, + elapsed_us: Self::elapsed_us(forwarded_at.duration_since(self.inner.received_at)), + succeeded: result.is_ok(), + }, ); + result } fn record_scheduled(&self, scheduled_at: TokioInstant) { @@ -505,7 +561,7 @@ impl ReverseRpcTrace { ) { self.emit_phase(ReverseRpcTimingEvent::HookCallback { trace: self.clone(), - hook_type: hook_type.to_string(), + hook_type: Self::timing_hook_type(hook_type), start_offset_us: self.start_offset_us(started_at), elapsed_us: Self::elapsed_us(elapsed), }); @@ -529,6 +585,14 @@ impl ReverseRpcTrace { fn emit_phase(&self, event: ReverseRpcTimingEvent) { let mut state = self.inner.timing_state.lock(); + self.emit_phase_with_state(&mut state, event); + } + + fn emit_phase_with_state( + &self, + state: &mut ReverseRpcTimingState, + event: ReverseRpcTimingEvent, + ) { if state.pending_completion.is_some() || state.completion_recorded { return; } @@ -587,8 +651,128 @@ impl ReverseRpcTrace { } } +#[derive(Clone)] +struct ReverseRpcRegistry(Arc); + +struct ReverseRpcRegistryInner { + traces: Mutex>, + next_generation: AtomicU64, +} + +impl ReverseRpcRegistry { + fn new() -> Self { + Self(Arc::new(ReverseRpcRegistryInner { + traces: Mutex::new(Vec::new()), + next_generation: AtomicU64::new(1), + })) + } + + fn next_generation(&self) -> u64 { + self.0.next_generation.fetch_add(1, Ordering::Relaxed) + } + + fn insert(&self, trace: ReverseRpcTrace) { + self.0.traces.lock().push(trace); + } + + fn remove(&self, trace: &ReverseRpcTrace) -> bool { + let mut traces = self.0.traces.lock(); + let Some(index) = traces + .iter() + .position(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) + else { + return false; + }; + traces.swap_remove(index); + true + } + + fn abandon_all(&self) { + let traces = std::mem::take(&mut *self.0.traces.lock()); + for trace in traces { + trace.record_abandoned(TokioInstant::now()); + } + } + + #[cfg(test)] + fn contains(&self, trace: &ReverseRpcTrace) -> bool { + self.0 + .traces + .lock() + .iter() + .any(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) + } + + #[cfg(test)] + fn is_empty(&self) -> bool { + self.0.traces.lock().is_empty() + } +} + +pub(crate) struct ReverseRpcRequest { + request: Option, + trace: Option, + registry: Option, +} + +impl std::ops::Deref for ReverseRpcRequest { + type Target = JsonRpcRequest; + + fn deref(&self) -> &Self::Target { + self.request + .as_ref() + .expect("reverse RPC request must exist until dispatch") + } +} + +impl ReverseRpcRequest { + fn new( + request: JsonRpcRequest, + trace: Option, + registry: Option, + ) -> Self { + Self { + request: Some(request), + trace, + registry, + } + } + + pub(crate) fn into_dispatch(mut self) -> (JsonRpcRequest, Option) { + let request = self + .request + .take() + .expect("reverse RPC request must exist until dispatch"); + let guard = self.trace.take().map(|trace| { + trace.record_scheduled(TokioInstant::now()); + ReverseRpcDispatchGuard { + registry: self + .registry + .take() + .expect("timed reverse RPC request must have a registry"), + request_id: request.id, + trace, + } + }); + (request, guard) + } +} + +impl Drop for ReverseRpcRequest { + fn drop(&mut self) { + if let Some(trace) = self.trace.take() + && self + .registry + .as_ref() + .is_some_and(|registry| registry.remove(&trace)) + { + trace.record_abandoned(TokioInstant::now()); + } + } +} + pub(crate) struct ReverseRpcDispatchGuard { - reverse_requests: Arc>>, + registry: ReverseRpcRegistry, request_id: u64, trace: ReverseRpcTrace, } @@ -597,31 +781,32 @@ impl ReverseRpcDispatchGuard { pub(crate) fn trace(&self) -> &ReverseRpcTrace { &self.trace } + + pub(crate) async fn scope(&self, future: F) -> F::Output { + REVERSE_RPC_RESPONSE_TRACE + .scope( + ReverseRpcResponseTrace { + request_id: self.request_id, + trace: self.trace.clone(), + }, + future, + ) + .await + } } impl Drop for ReverseRpcDispatchGuard { fn drop(&mut self) { - if remove_reverse_request_if_same(&self.reverse_requests, self.request_id, &self.trace) { + if self.registry.remove(&self.trace) { self.trace.record_abandoned(TokioInstant::now()); } } } -fn remove_reverse_request_if_same( - reverse_requests: &RwLock>, - request_id: u64, - trace: &ReverseRpcTrace, -) -> bool { - let mut reverse_requests = reverse_requests.write(); - if reverse_requests - .get(&request_id) - .is_some_and(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) - { - reverse_requests.remove(&request_id); - true - } else { - false - } +#[derive(Clone)] +enum ReverseRequestSender { + Public(mpsc::UnboundedSender), + Internal(mpsc::UnboundedSender), } /// Low-level JSON-RPC 2.0 client over Content-Length-framed streams. @@ -644,9 +829,9 @@ pub struct JsonRpcClient { /// natural request/response back-pressure of the wire. write_tx: mpsc::UnboundedSender, pending_requests: Arc>>, - reverse_requests: Arc>>, + reverse_requests: Option, notification_tx: broadcast::Sender, - request_tx: mpsc::UnboundedSender, + request_tx: ReverseRequestSender, read_task: Mutex>>, write_task: Mutex>>, timing_task: Mutex>>, @@ -672,14 +857,20 @@ impl JsonRpcClient { notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, ) -> Self { - Self::new_inner(writer, reader, notification_tx, request_tx, false) + Self::new_inner( + writer, + reader, + notification_tx, + ReverseRequestSender::Public(request_tx), + false, + ) } pub(crate) fn new_with_reverse_rpc_timing( writer: impl AsyncWrite + Unpin + Send + 'static, reader: impl AsyncRead + Unpin + Send + 'static, notification_tx: broadcast::Sender, - request_tx: mpsc::UnboundedSender, + request_tx: mpsc::UnboundedSender, ) -> Self { let trace_reverse_rpc = tracing::enabled!( target: REVERSE_RPC_TIMING_TARGET, @@ -689,7 +880,7 @@ impl JsonRpcClient { writer, reader, notification_tx, - request_tx, + ReverseRequestSender::Internal(request_tx), trace_reverse_rpc, ) } @@ -698,14 +889,14 @@ impl JsonRpcClient { writer: impl AsyncWrite + Unpin + Send + 'static, reader: impl AsyncRead + Unpin + Send + 'static, notification_tx: broadcast::Sender, - request_tx: mpsc::UnboundedSender, + request_tx: ReverseRequestSender, trace_reverse_rpc: bool, ) -> Self { let (write_tx, write_rx) = mpsc::unbounded_channel::(); let writer_span = tracing::error_span!("jsonrpc_write_loop"); let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span)); - let (timing, timing_task, correlation_hasher) = if trace_reverse_rpc { + let (timing, timing_task, correlation_hasher, reverse_requests) = if trace_reverse_rpc { let (phase_tx, phase_rx) = mpsc::channel::(REVERSE_RPC_TIMING_CAPACITY); let (terminal_tx, terminal_rx) = @@ -723,16 +914,17 @@ impl JsonRpcClient { dropped_records, ))), Some(RandomState::new()), + Some(ReverseRpcRegistry::new()), ) } else { - (None, None, None) + (None, None, None, None) }; let client = Self { request_id: AtomicU64::new(1), write_tx, pending_requests: Arc::new(RwLock::new(HashMap::new())), - reverse_requests: Arc::new(RwLock::new(HashMap::new())), + reverse_requests, notification_tx, request_tx, read_task: Mutex::new(None), @@ -774,14 +966,8 @@ impl JsonRpcClient { task.abort(); } self.pending_requests.write().clear(); - let abandoned = self - .reverse_requests - .write() - .drain() - .map(|(_, trace)| trace) - .collect::>(); - for trace in abandoned { - trace.record_abandoned(TokioInstant::now()); + if let Some(reverse_requests) = &self.reverse_requests { + reverse_requests.abandon_all(); } // Detach the timing task so it can drain the bounded queue. The // aborted read/write tasks drop the remaining senders, so it exits @@ -1034,9 +1220,9 @@ impl JsonRpcClient { async fn read_loop( reader: impl AsyncRead + Unpin + Send, pending_requests: Arc>>, - reverse_requests: Arc>>, + reverse_requests: Option, notification_tx: broadcast::Sender, - request_tx: mpsc::UnboundedSender, + request_tx: ReverseRequestSender, timing: Option, correlation_hasher: Option, ) { @@ -1103,36 +1289,53 @@ impl JsonRpcClient { let _ = notification_tx.send(notification); } JsonRpcMessage::Request(request) => { - let request_id = request.id; let trace = if tracing::enabled!( target: REVERSE_RPC_TIMING_TARGET, tracing::Level::DEBUG ) { - timing.as_ref().zip(correlation_hasher.as_ref()).map( - |(timing, correlation_hasher)| { + timing + .as_ref() + .zip(correlation_hasher.as_ref()) + .zip(reverse_requests.as_ref()) + .map(|((timing, correlation_hasher), registry)| { ReverseRpcTrace::new( &request, + registry.next_generation(), TokioInstant::now(), correlation_hasher, timing.clone(), ) - }, - ) + }) } else { None }; if let Some(trace) = &trace { - reverse_requests.write().insert(request_id, trace.clone()); - trace.mark_forwarding(TokioInstant::now()); + reverse_requests + .as_ref() + .expect("timed requests must have a registry") + .insert(trace.clone()); } - let forwarded = request_tx.send(request).is_ok(); - if let Some(trace) = &trace { - trace.record_forwarded(forwarded); - } - if !forwarded { - if let Some(trace) = reverse_requests.write().remove(&request_id) { - trace.record_abandoned(TokioInstant::now()); + let forwarded = match &request_tx { + ReverseRequestSender::Public(request_tx) => { + request_tx.send(request).is_ok() } + ReverseRequestSender::Internal(request_tx) => { + let request = ReverseRpcRequest::new( + request, + trace.clone(), + reverse_requests.clone(), + ); + let result = if let Some(trace) = &trace { + trace.forward(|| request_tx.send(request)) + } else { + request_tx.send(request) + }; + let forwarded = result.is_ok(); + drop(result); + forwarded + } + }; + if !forwarded { warn!("failed to forward JSON-RPC request, channel closed"); } } @@ -1157,13 +1360,8 @@ impl JsonRpcClient { ); pending.clear(); } - let abandoned = reverse_requests - .write() - .drain() - .map(|(_, trace)| trace) - .collect::>(); - for trace in abandoned { - trace.record_abandoned(TokioInstant::now()); + if let Some(reverse_requests) = &reverse_requests { + reverse_requests.abandon_all(); } } @@ -1353,10 +1551,15 @@ impl JsonRpcClient { } pub(crate) async fn write_response(&self, response: &JsonRpcResponse) -> Result<(), Error> { - let trace = self.reverse_requests.read().get(&response.id).cloned(); + let trace = REVERSE_RPC_RESPONSE_TRACE + .try_with(|scoped| (scoped.request_id == response.id).then(|| scoped.trace.clone())) + .ok() + .flatten(); let result = self.write_frame(response, trace.clone()).await; - if let Some(trace) = &trace { - let _ = remove_reverse_request_if_same(&self.reverse_requests, response.id, trace); + if let Some(trace) = &trace + && let Some(reverse_requests) = &self.reverse_requests + { + let _ = reverse_requests.remove(trace); } result } @@ -1416,27 +1619,6 @@ impl JsonRpcClient { ))), } } - - pub(crate) fn trace_reverse_request_scheduled( - &self, - request_id: u64, - ) -> Option { - let trace = self.reverse_requests.read().get(&request_id).cloned(); - if let Some(trace) = &trace { - trace.record_scheduled(TokioInstant::now()); - } - trace.map(|trace| ReverseRpcDispatchGuard { - reverse_requests: self.reverse_requests.clone(), - request_id, - trace, - }) - } - - pub(crate) fn abandon_reverse_request(&self, request_id: u64) { - if let Some(trace) = self.reverse_requests.write().remove(&request_id) { - trace.record_abandoned(TokioInstant::now()); - } - } } /// RAII guard that removes a pending-request entry from the map if the @@ -1786,51 +1968,79 @@ mod tests { request.id, &request.method, Some("private-session-id"), + 7, ); let repeated = ReverseRpcTrace::correlation_key( &correlation_hasher, request.id, &request.method, Some("private-session-id"), + 7, ); let different_session = ReverseRpcTrace::correlation_key( &correlation_hasher, request.id, &request.method, Some("other-session"), + 7, + ); + let different_generation = ReverseRpcTrace::correlation_key( + &correlation_hasher, + request.id, + &request.method, + Some("private-session-id"), + 8, ); assert_eq!(same, repeated); assert_ne!(same, different_session); + assert_ne!(same, different_generation); assert!(same.starts_with("rrpc-")); assert!(!same.contains("private-session-id")); } + #[test] + fn reverse_request_timing_labels_all_supported_hook_types() { + for hook_type in [ + "preToolUse", + "preMcpToolCall", + "postToolUse", + "postToolUseFailure", + "userPromptSubmitted", + "userPromptTransformed", + "sessionStart", + "sessionEnd", + "errorOccurred", + "agentStop", + ] { + assert_eq!(ReverseRpcTrace::timing_hook_type(hook_type), hook_type); + } + assert_eq!( + ReverseRpcTrace::timing_hook_type("PRIVATE_SENTINEL_DO_NOT_TRACE"), + "unknown" + ); + } + #[test] fn reverse_request_guard_only_removes_its_own_generation() { let (timing, _phase_rx, _terminal_rx, _dropped_records) = timing_channel(1); let request = JsonRpcRequest::new(17, "hooks.invoke", None); let now = TokioInstant::now(); let correlation_hasher = RandomState::new(); - let first = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing.clone()); - let second = ReverseRpcTrace::new(&request, now, &correlation_hasher, timing); - let reverse_requests = Arc::new(RwLock::new(HashMap::new())); - reverse_requests.write().insert(request.id, first.clone()); + let first = ReverseRpcTrace::new(&request, 1, now, &correlation_hasher, timing.clone()); + let second = ReverseRpcTrace::new(&request, 2, now, &correlation_hasher, timing); + let registry = ReverseRpcRegistry::new(); + registry.insert(first.clone()); + registry.insert(second.clone()); let guard = ReverseRpcDispatchGuard { - reverse_requests: reverse_requests.clone(), + registry: registry.clone(), request_id: request.id, trace: first, }; - reverse_requests.write().insert(request.id, second.clone()); drop(guard); - let retained = reverse_requests - .read() - .get(&request.id) - .cloned() - .expect("new request generation should remain tracked"); - assert!(Arc::ptr_eq(&retained.inner, &second.inner)); + assert!(registry.contains(&second)); } #[tokio::test(start_paused = true)] @@ -1850,11 +2060,10 @@ mod tests { Some(serde_json::json!({ "sessionId": "session" })), ); let received_at = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, received_at, &RandomState::new(), timing); + let trace = ReverseRpcTrace::new(&request, 1, received_at, &RandomState::new(), timing); tokio::time::advance(Duration::from_millis(5)).await; - trace.mark_forwarding(TokioInstant::now()); - trace.record_forwarded(true); + trace.forward(|| Ok::<(), ()>(())).unwrap(); tokio::time::advance(Duration::from_millis(7)).await; trace.record_scheduled(TokioInstant::now()); @@ -1875,6 +2084,85 @@ mod tests { assert!(schedule.contains("since_receive_us=12000")); } + #[tokio::test] + async fn reverse_request_forwarding_is_recorded_before_dispatch_can_start() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (timing, phase_rx, terminal_rx, dropped_records) = + timing_channel(REVERSE_RPC_TIMING_CAPACITY); + tokio::spawn(JsonRpcClient::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + )); + let request = JsonRpcRequest::new(31, "hooks.invoke", None); + let now = TokioInstant::now(); + let trace = ReverseRpcTrace::new(&request, 1, now, &RandomState::new(), timing); + let registry = ReverseRpcRegistry::new(); + registry.insert(trace.clone()); + let forwarded = + ReverseRpcRequest::new(request, Some(trace.clone()), Some(registry.clone())); + let (request_tx, request_rx) = std::sync::mpsc::channel::(); + let (received_tx, received_rx) = std::sync::mpsc::channel(); + let receiver = std::thread::spawn(move || { + let request = request_rx.recv().unwrap(); + received_tx.send(()).unwrap(); + let (_request, guard) = request.into_dispatch(); + drop(guard); + }); + + trace + .forward(|| { + request_tx.send(forwarded).map_err(|_| ())?; + received_rx.recv().map_err(|_| ())?; + Ok::<(), ()>(()) + }) + .unwrap(); + receiver.join().unwrap(); + + wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; + let output = trace_buffer.text(); + let forward = output.find("phase=\"request_forward\"").unwrap(); + let schedule = output.find("phase=\"request_schedule\"").unwrap(); + let complete = output.find("phase=\"request_complete\"").unwrap(); + assert!(forward < schedule); + assert!(schedule < complete); + assert!(registry.is_empty()); + } + + #[tokio::test] + async fn closed_forward_channel_records_failed_forward_before_abandonment() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (mut server, reader) = tokio::io::duplex(4096); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, request_rx) = mpsc::unbounded_channel(); + drop(request_rx); + let client = JsonRpcClient::new_with_reverse_rpc_timing( + tokio::io::sink(), + reader, + notification_tx, + request_tx, + ); + let request = JsonRpcRequest::new(32, "hooks.invoke", None); + + server.write_all(&frame(&request)).await.unwrap(); + + wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; + let output = trace_buffer.text(); + let forward = output + .lines() + .find(|line| line.contains("phase=\"request_forward\"")) + .expect("failed forwarding phase should be emitted"); + assert!(forward.contains("status=\"failed\"")); + assert!( + output.find("phase=\"request_forward\"").unwrap() + < output.find("phase=\"request_complete\"").unwrap() + ); + + client.force_close(); + } + #[tokio::test] async fn public_client_does_not_retain_reverse_request_timing_state() { let (mut server, reader) = tokio::io::duplex(4096); @@ -1887,7 +2175,7 @@ mod tests { let forwarded = request_rx.recv().await.unwrap(); assert_eq!(forwarded.id, request.id); - assert!(client.reverse_requests.read().is_empty()); + assert!(client.reverse_requests.is_none()); assert!(client.timing_task.lock().is_none()); client.force_close(); } @@ -1915,7 +2203,8 @@ mod tests { let forwarded = request_rx.recv().await.unwrap(); assert_eq!(forwarded.id, request.id); - assert!(client.reverse_requests.read().is_empty()); + assert!(forwarded.trace.is_none()); + assert!(client.reverse_requests.is_none()); assert!(client.timing_task.lock().is_none()); client.force_close(); } @@ -1927,7 +2216,7 @@ mod tests { let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); let request = JsonRpcRequest::new(37, "hooks.invoke", None); let now = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, now, &RandomState::new(), timing); + let trace = ReverseRpcTrace::new(&request, 1, now, &RandomState::new(), timing); trace.mark_forwarding(now); trace.record_phase("first", now, Duration::ZERO, true); @@ -1966,14 +2255,14 @@ mod tests { let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); let now = TokioInstant::now(); - let filler_request = JsonRpcRequest::new(38, "filler.request", None); + let filler_request = JsonRpcRequest::new(38, "hooks.invoke", None); let filler = - ReverseRpcTrace::new(&filler_request, now, &RandomState::new(), timing.clone()); + ReverseRpcTrace::new(&filler_request, 1, now, &RandomState::new(), timing.clone()); filler.mark_forwarding(now); assert!(filler.record_complete(now, true)); - let writer_request = JsonRpcRequest::new(39, "writer.request", None); - let writer = ReverseRpcTrace::new(&writer_request, now, &RandomState::new(), timing); + let writer_request = JsonRpcRequest::new(39, "userInput.request", None); + let writer = ReverseRpcTrace::new(&writer_request, 2, now, &RandomState::new(), timing); writer.mark_forwarding(now); let mut writer_trace = ReverseRpcWriteTrace::new(writer); writer_trace.record_complete(now, true); @@ -1997,9 +2286,9 @@ mod tests { .find(|line| line.contains("phase=\"records_dropped\"")) .expect("terminal saturation diagnostic should be emitted"); assert!(dropped.contains("dropped_records=1")); - assert!(output.contains("rpc_method=filler.request")); + assert!(output.contains("rpc_method=hooks.invoke")); assert!(!output.lines().any(|line| { - line.contains("rpc_method=writer.request") + line.contains("rpc_method=userInput.request") && line.contains("phase=\"request_complete\"") })); } @@ -2010,9 +2299,9 @@ mod tests { let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); let now = TokioInstant::now(); - let filler_request = JsonRpcRequest::new(40, "filler.request", None); + let filler_request = JsonRpcRequest::new(40, "hooks.invoke", None); let filler = - ReverseRpcTrace::new(&filler_request, now, &RandomState::new(), timing.clone()); + ReverseRpcTrace::new(&filler_request, 1, now, &RandomState::new(), timing.clone()); filler.mark_forwarding(now); assert!(filler.record_complete(now, true)); @@ -2032,8 +2321,8 @@ mod tests { write_task.abort(); let _ = write_task.await; - let writer_request = JsonRpcRequest::new(41, "writer.request", None); - let writer = ReverseRpcTrace::new(&writer_request, now, &RandomState::new(), timing); + let writer_request = JsonRpcRequest::new(41, "userInput.request", None); + let writer = ReverseRpcTrace::new(&writer_request, 2, now, &RandomState::new(), timing); writer.mark_forwarding(now); let error = client .write_frame( @@ -2066,9 +2355,9 @@ mod tests { .find(|line| line.contains("phase=\"records_dropped\"")) .expect("closed-writer saturation diagnostic should be emitted"); assert!(dropped.contains("dropped_records=1")); - assert!(output.contains("rpc_method=filler.request")); + assert!(output.contains("rpc_method=hooks.invoke")); assert!(!output.lines().any(|line| { - line.contains("rpc_method=writer.request") + line.contains("rpc_method=userInput.request") && line.contains("phase=\"request_complete\"") })); @@ -2092,7 +2381,7 @@ mod tests { ); let request = JsonRpcRequest::new( 41, - "hooks.invoke", + SENTINEL, Some(serde_json::json!({ "sessionId": SENTINEL, "hookType": "userPromptSubmitted", @@ -2108,29 +2397,28 @@ mod tests { let forwarded = request_rx.recv().await.unwrap(); tokio::time::advance(Duration::from_millis(13)).await; - let trace = client - .trace_reverse_request_scheduled(forwarded.id) - .expect("reverse request timing should be tracked"); + let (forwarded, trace) = forwarded.into_dispatch(); + let trace = trace.expect("reverse request timing should be tracked"); let callback_start = TokioInstant::now(); tokio::time::advance(Duration::from_millis(3)).await; - trace.trace().record_hook_callback( - "userPromptSubmitted", - callback_start, - callback_start.elapsed(), - ); - client - .write_response(&JsonRpcResponse { + trace + .trace() + .record_hook_callback(SENTINEL, callback_start, callback_start.elapsed()); + trace + .scope(client.write_response(&JsonRpcResponse { jsonrpc: "2.0".to_string(), id: forwarded.id, result: Some(serde_json::json!({ "output": SENTINEL })), error: None, - }) + })) .await .unwrap(); wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; let output = trace_buffer.text(); assert!(output.contains("github_copilot_sdk::reverse_rpc_timing")); + assert!(output.contains("rpc_method=unknown")); + assert!(output.contains("hook_type=\"unknown\"")); assert!(output.contains("phase=\"request_forward\"")); assert!(output.contains("phase=\"request_schedule\"")); assert!(output.contains("elapsed_us=13000")); @@ -2242,25 +2530,23 @@ mod tests { let request = JsonRpcRequest::new(53, "hooks.invoke", None); let now = TokioInstant::now(); let trace = ReverseRpcTrace::for_test(&request, now, now); - client - .reverse_requests - .write() - .insert(request.id, trace.clone()); + let registry = ReverseRpcRegistry::new(); + registry.insert(trace.clone()); let dispatch_guard = ReverseRpcDispatchGuard { - reverse_requests: client.reverse_requests.clone(), + registry: registry.clone(), request_id: request.id, trace, }; let response_task = tokio::spawn({ let client = client.clone(); async move { - client - .write_response(&JsonRpcResponse { + dispatch_guard + .scope(client.write_response(&JsonRpcResponse { jsonrpc: "2.0".to_string(), id: request.id, result: Some(serde_json::json!({})), error: None, - }) + })) .await } }); @@ -2268,7 +2554,6 @@ mod tests { assert_eq!(started_rx.recv().await, Some("write")); response_task.abort(); let _ = response_task.await; - drop(dispatch_guard); tokio::time::advance(Duration::from_millis(10)).await; assert_eq!(started_rx.recv().await, Some("flush")); @@ -2281,7 +2566,110 @@ mod tests { .collect::>(); assert_eq!(complete.len(), 1); assert!(complete[0].contains("status=\"succeeded\"")); - assert!(client.reverse_requests.read().is_empty()); + assert!(registry.is_empty()); + + client.force_close(); + } + + #[tokio::test] + async fn response_timing_uses_the_exact_dispatch_generation_when_ids_are_reused() { + let trace_buffer = TraceBuffer::default(); + let _subscriber = tracing::subscriber::set_default(trace_subscriber(trace_buffer.clone())); + let (notification_tx, _) = broadcast::channel(1); + let (request_tx, mut request_rx) = mpsc::unbounded_channel(); + let (mut server, reader) = tokio::io::duplex(4096); + let client = JsonRpcClient::new_with_reverse_rpc_timing( + tokio::io::sink(), + reader, + notification_tx, + request_tx, + ); + let params = Some(serde_json::json!({ "sessionId": "same-session" })); + let first_request = JsonRpcRequest::new(61, "hooks.invoke", params.clone()); + let second_request = JsonRpcRequest::new(61, "hooks.invoke", params); + server.write_all(&frame(&first_request)).await.unwrap(); + server.write_all(&frame(&second_request)).await.unwrap(); + let first_forwarded = request_rx.recv().await.unwrap(); + let second_forwarded = request_rx.recv().await.unwrap(); + + // Acquire the newer dispatch first to prove scheduling order cannot + // change which receipt-generation each forwarded request carries. + let (second_forwarded, second_guard) = second_forwarded.into_dispatch(); + let second_guard = second_guard.expect("second request should carry timing"); + let second_trace = second_guard.trace.clone(); + let second_correlation = second_trace.inner.correlation_key.clone(); + let (first_forwarded, first_guard) = first_forwarded.into_dispatch(); + let first_guard = first_guard.expect("first request should carry timing"); + let first_correlation = first_guard.trace.inner.correlation_key.clone(); + assert_ne!(first_correlation, second_correlation); + + first_guard + .scope(client.write_response(&JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: first_forwarded.id, + result: Some(serde_json::json!({})), + error: None, + })) + .await + .unwrap(); + wait_for_trace( + &trace_buffer, + &format!( + "correlation_key={first_correlation} rpc_method=hooks.invoke phase=\"request_complete\"" + ), + ) + .await; + assert!( + client + .reverse_requests + .as_ref() + .is_some_and(|registry| registry.contains(&second_trace)) + ); + + second_guard + .scope(client.write_response(&JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: second_forwarded.id, + result: Some(serde_json::json!({})), + error: None, + })) + .await + .unwrap(); + wait_for_trace( + &trace_buffer, + &format!( + "correlation_key={second_correlation} rpc_method=hooks.invoke phase=\"request_complete\"" + ), + ) + .await; + + let output = trace_buffer.text(); + assert_eq!( + output + .lines() + .filter(|line| { + line.contains(&format!("correlation_key={first_correlation}")) + && line.contains("phase=\"request_complete\"") + }) + .count(), + 1 + ); + assert_eq!( + output + .lines() + .filter(|line| { + line.contains(&format!("correlation_key={second_correlation}")) + && line.contains("phase=\"request_complete\"") + }) + .count(), + 1 + ); + assert!( + client + .reverse_requests + .as_ref() + .is_some_and(ReverseRpcRegistry::is_empty) + ); client.force_close(); } @@ -2306,19 +2694,19 @@ mod tests { let request = JsonRpcRequest::new(59, "hooks.invoke", None); server.write_all(&frame(&request)).await.unwrap(); let forwarded = request_rx.recv().await.unwrap(); - let dispatch_guard = client - .trace_reverse_request_scheduled(forwarded.id) - .expect("enabled timing target should track the request"); + let (forwarded, dispatch_guard) = forwarded.into_dispatch(); + let dispatch_guard = + dispatch_guard.expect("enabled timing target should track the request"); let response_task = tokio::spawn({ let client = client.clone(); async move { - client - .write_response(&JsonRpcResponse { + dispatch_guard + .scope(client.write_response(&JsonRpcResponse { jsonrpc: "2.0".to_string(), id: forwarded.id, result: Some(serde_json::json!({})), error: None, - }) + })) .await } }); @@ -2326,7 +2714,6 @@ mod tests { assert_eq!(started_rx.recv().await, Some("write")); client.force_close(); assert!(response_task.await.unwrap().is_err()); - drop(dispatch_guard); wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; let complete = trace_buffer @@ -2381,7 +2768,7 @@ mod tests { ); let request = JsonRpcRequest::new(53, "hooks.invoke", None); let now = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, now, &RandomState::new(), timing); + let trace = ReverseRpcTrace::new(&request, 1, now, &RandomState::new(), timing); trace.mark_forwarding(now); tokio::time::timeout( diff --git a/rust/src/lib.rs b/rust/src/lib.rs index f3383e4bc4..6d4e06e52f 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1023,7 +1023,7 @@ struct ClientInner { ffi_host: parking_lot::Mutex>>, rpc: JsonRpcClient, cwd: PathBuf, - request_rx: parking_lot::Mutex>>, + request_rx: parking_lot::Mutex>>, notification_tx: broadcast::Sender, router: router::SessionRouter, github_token_registry: Arc, @@ -1431,7 +1431,6 @@ impl Client { Some(dispatcher.clone()), client.inner.on_github_telemetry.clone(), client.inner.github_token_registry.clone(), - Arc::downgrade(&client.inner), ); client.rpc().llm_inference().set_provider().await?; let llm_inference_elapsed = llm_inference_start.elapsed(); @@ -1600,7 +1599,7 @@ impl Client { mode: ClientMode, ) -> Result { let setup_start = Instant::now(); - let (request_tx, request_rx) = mpsc::unbounded_channel::(); + let (request_tx, request_rx) = mpsc::unbounded_channel::(); let (notification_broadcast_tx, _) = broadcast::channel::(1024); let rpc = JsonRpcClient::new_with_reverse_rpc_timing( writer, @@ -2029,17 +2028,6 @@ impl Client { self.inner.rpc.write_response(response).await } - pub(crate) fn trace_reverse_request_scheduled( - &self, - request_id: u64, - ) -> Option { - self.inner.rpc.trace_reverse_request_scheduled(request_id) - } - - pub(crate) fn abandon_reverse_request(&self, request_id: u64) { - self.inner.rpc.abandon_reverse_request(request_id); - } - /// Reconstruct a [`Client`] handle from a shared inner pointer. pub(crate) fn from_inner(inner: Arc) -> Self { Self { inner } @@ -2049,7 +2037,9 @@ impl Client { /// /// Can only be called once — subsequent calls return `None`. #[expect(dead_code, reason = "reserved for future pub(crate) use")] - pub(crate) fn take_request_rx(&self) -> Option> { + pub(crate) fn take_request_rx( + &self, + ) -> Option> { self.inner.request_rx.lock().take() } @@ -2070,7 +2060,6 @@ impl Client { self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), self.inner.github_token_registry.clone(), - Arc::downgrade(&self.inner), ); self.inner.router.register(session_id) } @@ -2090,7 +2079,6 @@ impl Client { self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), self.inner.github_token_registry.clone(), - Arc::downgrade(&self.inner), ); let id = self.inner.github_token_registry.register(provider); github_token::GitHubTokenRegistration::new(self.inner.github_token_registry.clone(), id) @@ -2310,7 +2298,6 @@ impl Client { self.inner.llm_inference.get().cloned(), self.inner.on_github_telemetry.clone(), self.inner.github_token_registry.clone(), - Arc::downgrade(&self.inner), ); } diff --git a/rust/src/router.rs b/rust/src/router.rs index 8c9fb922c1..c02a00928d 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -1,25 +1,24 @@ use std::collections::HashMap; -use std::sync::{Arc, Weak}; +use std::sync::Arc; use parking_lot::Mutex; use tokio::sync::{broadcast, mpsc}; use tracing::warn; -use crate::jsonrpc::{JsonRpcNotification, JsonRpcRequest}; +use crate::jsonrpc::{JsonRpcNotification, ReverseRpcRequest}; use crate::types::{SessionEventNotification, SessionId}; -use crate::{Client, ClientInner}; /// Per-session channels created by the router during session registration. pub(crate) struct SessionChannels { /// Filtered `session.event` notifications for this session. pub(crate) notifications: mpsc::UnboundedReceiver, /// Filtered JSON-RPC requests (tool.call, userInput.request, etc.) for this session. - pub(crate) requests: mpsc::UnboundedReceiver, + pub(crate) requests: mpsc::UnboundedReceiver, } struct SessionSenders { notifications: mpsc::UnboundedSender, - requests: mpsc::UnboundedSender, + requests: mpsc::UnboundedSender, } /// Routes notifications and requests by sessionId to per-session channels. @@ -85,11 +84,10 @@ impl SessionRouter { pub(crate) fn ensure_started( &self, notification_tx: &broadcast::Sender, - request_rx: &Mutex>>, + request_rx: &Mutex>>, llm_inference: Option>, github_telemetry: Option, github_token_registry: Arc, - client: Weak, ) { let mut started = self.started.lock(); if *started { @@ -194,14 +192,10 @@ impl SessionRouter { if let Some(dispatcher) = &llm_inference { dispatcher.dispatch(request).await; } else { - let request_id = request.id; warn!( method = %request.method, "llmInference request with no provider registered" ); - if let Some(inner) = client.upgrade() { - Client::from_inner(inner).abandon_reverse_request(request_id); - } } continue; } @@ -218,31 +212,19 @@ impl SessionRouter { guard.get(sid).map(|s| s.requests.clone()) }; if let Some(sender) = sender { - if let Err(error) = sender.send(request) - && let Some(inner) = client.upgrade() - { - Client::from_inner(inner).abandon_reverse_request(error.0.id); - } + let _ = sender.send(request); } else { - let request_id = request.id; warn!( session_id = sid, method = %request.method, "request for unregistered session" ); - if let Some(inner) = client.upgrade() { - Client::from_inner(inner).abandon_reverse_request(request_id); - } } } else { - let request_id = request.id; warn!( method = %request.method, "request missing sessionId" ); - if let Some(inner) = client.upgrade() { - Client::from_inner(inner).abandon_reverse_request(request_id); - } } } }); diff --git a/rust/src/session.rs b/rust/src/session.rs index 398a097ddd..399c9cec73 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1610,7 +1610,7 @@ fn spawn_event_loop( .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into())); } while let Ok(request) = requests.try_recv() { - client.abandon_reverse_request(request.id); + drop(request); } } .instrument(span), @@ -2347,13 +2347,27 @@ struct RequestDispatchContext<'a> { /// Process a JSON-RPC request from the CLI. async fn handle_request( + session_id: &SessionId, + ctx: RequestDispatchContext<'_>, + request: crate::jsonrpc::ReverseRpcRequest, +) { + let (request, reverse_rpc_trace) = request.into_dispatch(); + let dispatch = handle_request_inner(session_id, ctx, request, reverse_rpc_trace.as_ref()); + if let Some(trace) = &reverse_rpc_trace { + trace.scope(dispatch).await; + } else { + dispatch.await; + } +} + +async fn handle_request_inner( session_id: &SessionId, ctx: RequestDispatchContext<'_>, request: crate::JsonRpcRequest, + reverse_rpc_trace: Option<&crate::jsonrpc::ReverseRpcDispatchGuard>, ) { let sid = session_id.clone(); let client = ctx.client; - let reverse_rpc_trace = client.trace_reverse_request_scheduled(request.id); let handlers = ctx.handlers; let hooks = ctx.hooks; let transforms = ctx.transforms; @@ -2394,7 +2408,7 @@ async fn handle_request( &sid, hook_type, input, - reverse_rpc_trace.as_ref().map(|guard| guard.trace()), + reverse_rpc_trace.map(|guard| guard.trace()), ) .await { From 5be450f67fb4119e1f547a694a273b39e0a7e77c Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 28 Aug 2026 17:25:25 -0600 Subject: [PATCH 6/7] Address reverse timing review feedback Remove disabled-path clock reads, close request receivers before shutdown draining, and isolate timing subscribers on the opt-in timing thread. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/jsonrpc.rs | 152 ++++++++++++++++++++++++++------------------ rust/src/session.rs | 1 + 2 files changed, 92 insertions(+), 61 deletions(-) diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index ddf1f2aad7..e90357a691 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -1,8 +1,7 @@ -use std::collections::{HashMap, VecDeque, hash_map::RandomState}; +use std::collections::{HashMap, VecDeque}; use std::future::Future; use std::hash::{BuildHasher, Hash, Hasher}; use std::sync::Arc; -use std::sync::OnceLock; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; @@ -185,6 +184,7 @@ const CONTENT_LENGTH_HEADER: &str = "Content-Length: "; /// has an active DEBUG subscriber when the client is constructed. const REVERSE_RPC_TIMING_TARGET: &str = "github_copilot_sdk::reverse_rpc_timing"; const REVERSE_RPC_TIMING_CAPACITY: usize = 256; +type CorrelationHasher = std::collections::hash_map::RandomState; /// Rewrites unpaired UTF-16 surrogate escapes to `\uFFFD`. /// @@ -260,7 +260,7 @@ struct WriteCommand { frame: Vec, ack: oneshot::Sender>, reverse_rpc: Option, - enqueued_at: TokioInstant, + enqueued_at: Option, } struct ReverseRpcWriteTrace { @@ -375,7 +375,7 @@ struct ReverseRpcTraceInner { correlation_key: String, method: &'static str, received_at: TokioInstant, - forwarded_at: OnceLock, + forwarded_at: std::sync::OnceLock, timing: ReverseRpcTimingEmitter, timing_state: Mutex, emitted_phase_records: AtomicU64, @@ -401,7 +401,7 @@ impl ReverseRpcTrace { request: &JsonRpcRequest, generation: u64, received_at: TokioInstant, - correlation_hasher: &RandomState, + correlation_hasher: &CorrelationHasher, timing: ReverseRpcTimingEmitter, ) -> Self { let session_id = request @@ -420,7 +420,7 @@ impl ReverseRpcTrace { ), method: Self::timing_method(&request.method), received_at, - forwarded_at: OnceLock::new(), + forwarded_at: std::sync::OnceLock::new(), timing, timing_state: Mutex::new(ReverseRpcTimingState::default()), emitted_phase_records: AtomicU64::new(0), @@ -446,7 +446,7 @@ impl ReverseRpcTrace { request, 0, received_at, - &RandomState::new(), + &CorrelationHasher::new(), ReverseRpcTimingEmitter { phase_tx, terminal_tx, @@ -458,7 +458,7 @@ impl ReverseRpcTrace { } fn correlation_key( - correlation_hasher: &RandomState, + correlation_hasher: &CorrelationHasher, request_id: u64, method: &str, session_id: Option<&str>, @@ -834,7 +834,7 @@ pub struct JsonRpcClient { request_tx: ReverseRequestSender, read_task: Mutex>>, write_task: Mutex>>, - timing_task: Mutex>>, + timing_task: Mutex>>, } impl JsonRpcClient { @@ -908,12 +908,12 @@ impl JsonRpcClient { terminal_tx, dropped_records: dropped_records.clone(), }), - Some(tokio::spawn(Self::timing_loop( + Some(Self::spawn_timing_thread( phase_rx, terminal_rx, dropped_records, - ))), - Some(RandomState::new()), + )), + Some(CorrelationHasher::new()), Some(ReverseRpcRegistry::new()), ) } else { @@ -958,6 +958,26 @@ impl JsonRpcClient { client } + fn spawn_timing_thread( + phase_rx: mpsc::Receiver, + terminal_rx: mpsc::Receiver, + dropped_records: Arc, + ) -> std::thread::JoinHandle<()> { + let dispatch = tracing::dispatcher::get_default(Clone::clone); + std::thread::Builder::new() + .name("copilot-reverse-rpc-timing".to_string()) + .spawn(move || { + tracing::dispatcher::with_default(&dispatch, || { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("reverse RPC timing runtime should start") + .block_on(Self::timing_loop(phase_rx, terminal_rx, dropped_records)); + }); + }) + .expect("reverse RPC timing thread should start") + } + pub(crate) fn force_close(&self) { if let Some(task) = self.read_task.lock().take() { task.abort(); @@ -1177,27 +1197,27 @@ impl JsonRpcClient { enqueued_at, }) = rx.recv().await { - let queue_elapsed = enqueued_at.elapsed(); - - let write_start = TokioInstant::now(); + let queue_timing = enqueued_at.map(|enqueued_at| (enqueued_at, enqueued_at.elapsed())); + let write_start = reverse_rpc.as_ref().map(|_| TokioInstant::now()); let write_result = writer.write_all(&frame).await; - let write_elapsed = write_start.elapsed(); + let write_timing = write_start.map(|write_start| (write_start, write_start.elapsed())); let write_succeeded = write_result.is_ok(); let (result, flush_timing) = match write_result { Ok(()) => { - let flush_start = TokioInstant::now(); + let flush_start = reverse_rpc.as_ref().map(|_| TokioInstant::now()); let flush_result = writer.flush().await; - let flush_elapsed = flush_start.elapsed(); let flush_succeeded = flush_result.is_ok(); ( flush_result, - Some((flush_start, flush_elapsed, flush_succeeded)), + flush_start.map(|flush_start| { + (flush_start, flush_start.elapsed(), flush_succeeded) + }), ) } Err(error) => (Err(error), None), }; - let completed_at = TokioInstant::now(); + let completed_at = reverse_rpc.as_ref().map(|_| TokioInstant::now()); let succeeded = result.is_ok(); // Caller may have dropped the ack receiver (e.g. their @@ -1207,12 +1227,19 @@ impl JsonRpcClient { if let Some(write_trace) = &mut reverse_rpc { let trace = &write_trace.trace; + let (enqueued_at, queue_elapsed) = + queue_timing.expect("timed write must include its enqueue timestamp"); + let (write_start, write_elapsed) = + write_timing.expect("timed write must include its write timestamp"); trace.record_phase("writer_queue", enqueued_at, queue_elapsed, true); trace.record_phase("write_all", write_start, write_elapsed, write_succeeded); if let Some((flush_start, flush_elapsed, flush_succeeded)) = flush_timing { trace.record_phase("flush", flush_start, flush_elapsed, flush_succeeded); } - write_trace.record_complete(completed_at, succeeded); + write_trace.record_complete( + completed_at.expect("timed write must include its completion timestamp"), + succeeded, + ); } } } @@ -1224,7 +1251,7 @@ impl JsonRpcClient { notification_tx: broadcast::Sender, request_tx: ReverseRequestSender, timing: Option, - correlation_hasher: Option, + correlation_hasher: Option, ) { let mut reader = BufReader::new(reader); @@ -1569,9 +1596,9 @@ impl JsonRpcClient { message: &T, reverse_rpc: Option, ) -> Result<(), Error> { - let encode_start = TokioInstant::now(); + let encode_start = reverse_rpc.as_ref().map(|_| TokioInstant::now()); let encoded = serde_json::to_vec(message); - if let Some(trace) = &reverse_rpc { + if let (Some(trace), Some(encode_start)) = (&reverse_rpc, encode_start) { trace.record_phase( "response_encode", encode_start, @@ -1590,7 +1617,7 @@ impl JsonRpcClient { frame.extend_from_slice(&body); let (ack_tx, ack_rx) = oneshot::channel(); - let enqueued_at = TokioInstant::now(); + let enqueued_at = reverse_rpc.as_ref().map(|_| TokioInstant::now()); if let Some(trace) = &reverse_rpc { trace.transfer_completion_to_writer(); } @@ -1962,7 +1989,7 @@ mod tests { "hooks.invoke", Some(serde_json::json!({ "sessionId": "private-session-id" })), ); - let correlation_hasher = RandomState::new(); + let correlation_hasher = CorrelationHasher::new(); let same = ReverseRpcTrace::correlation_key( &correlation_hasher, request.id, @@ -2026,7 +2053,7 @@ mod tests { let (timing, _phase_rx, _terminal_rx, _dropped_records) = timing_channel(1); let request = JsonRpcRequest::new(17, "hooks.invoke", None); let now = TokioInstant::now(); - let correlation_hasher = RandomState::new(); + let correlation_hasher = CorrelationHasher::new(); let first = ReverseRpcTrace::new(&request, 1, now, &correlation_hasher, timing.clone()); let second = ReverseRpcTrace::new(&request, 2, now, &correlation_hasher, timing); let registry = ReverseRpcRegistry::new(); @@ -2060,7 +2087,8 @@ mod tests { Some(serde_json::json!({ "sessionId": "session" })), ); let received_at = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, 1, received_at, &RandomState::new(), timing); + let trace = + ReverseRpcTrace::new(&request, 1, received_at, &CorrelationHasher::new(), timing); tokio::time::advance(Duration::from_millis(5)).await; trace.forward(|| Ok::<(), ()>(())).unwrap(); @@ -2097,7 +2125,7 @@ mod tests { )); let request = JsonRpcRequest::new(31, "hooks.invoke", None); let now = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, 1, now, &RandomState::new(), timing); + let trace = ReverseRpcTrace::new(&request, 1, now, &CorrelationHasher::new(), timing); let registry = ReverseRpcRegistry::new(); registry.insert(trace.clone()); let forwarded = @@ -2216,7 +2244,7 @@ mod tests { let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); let request = JsonRpcRequest::new(37, "hooks.invoke", None); let now = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, 1, now, &RandomState::new(), timing); + let trace = ReverseRpcTrace::new(&request, 1, now, &CorrelationHasher::new(), timing); trace.mark_forwarding(now); trace.record_phase("first", now, Duration::ZERO, true); @@ -2256,13 +2284,19 @@ mod tests { let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); let now = TokioInstant::now(); let filler_request = JsonRpcRequest::new(38, "hooks.invoke", None); - let filler = - ReverseRpcTrace::new(&filler_request, 1, now, &RandomState::new(), timing.clone()); + let filler = ReverseRpcTrace::new( + &filler_request, + 1, + now, + &CorrelationHasher::new(), + timing.clone(), + ); filler.mark_forwarding(now); assert!(filler.record_complete(now, true)); let writer_request = JsonRpcRequest::new(39, "userInput.request", None); - let writer = ReverseRpcTrace::new(&writer_request, 2, now, &RandomState::new(), timing); + let writer = + ReverseRpcTrace::new(&writer_request, 2, now, &CorrelationHasher::new(), timing); writer.mark_forwarding(now); let mut writer_trace = ReverseRpcWriteTrace::new(writer); writer_trace.record_complete(now, true); @@ -2300,8 +2334,13 @@ mod tests { let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(1); let now = TokioInstant::now(); let filler_request = JsonRpcRequest::new(40, "hooks.invoke", None); - let filler = - ReverseRpcTrace::new(&filler_request, 1, now, &RandomState::new(), timing.clone()); + let filler = ReverseRpcTrace::new( + &filler_request, + 1, + now, + &CorrelationHasher::new(), + timing.clone(), + ); filler.mark_forwarding(now); assert!(filler.record_complete(now, true)); @@ -2322,7 +2361,8 @@ mod tests { let _ = write_task.await; let writer_request = JsonRpcRequest::new(41, "userInput.request", None); - let writer = ReverseRpcTrace::new(&writer_request, 2, now, &RandomState::new(), timing); + let writer = + ReverseRpcTrace::new(&writer_request, 2, now, &CorrelationHasher::new(), timing); writer.mark_forwarding(now); let error = client .write_frame( @@ -2454,7 +2494,7 @@ mod tests { frame: frame(&serde_json::json!({})), ack: first_ack_tx, reverse_rpc: None, - enqueued_at: TokioInstant::now(), + enqueued_at: None, }) .unwrap(); assert_eq!(started_rx.recv().await, Some("write")); @@ -2479,7 +2519,7 @@ mod tests { }), ack: second_ack_tx, reverse_rpc: Some(ReverseRpcWriteTrace::new(trace)), - enqueued_at: TokioInstant::now(), + enqueued_at: Some(TokioInstant::now()), }) .unwrap(); @@ -2726,7 +2766,7 @@ mod tests { assert!(complete[0].contains("status=\"failed\"")); } - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[tokio::test] async fn slow_timing_subscriber_does_not_delay_response_ack() { let (entered_tx, entered_rx) = std::sync::mpsc::channel(); let release = Arc::new((std::sync::Mutex::new(false), std::sync::Condvar::new())); @@ -2736,27 +2776,17 @@ mod tests { }; let (timing, phase_rx, terminal_rx, dropped_records) = timing_channel(REVERSE_RPC_TIMING_CAPACITY); - let timing_thread = std::thread::spawn(move || { - let subscriber = tracing_subscriber::registry().with( - tracing_subscriber::fmt::layer() - .with_writer(blocking_writer) - .with_ansi(false) - .without_time() - .with_filter(tracing_subscriber::filter::filter_fn(|metadata| { - metadata.target() == REVERSE_RPC_TIMING_TARGET - })), - ); - tracing::subscriber::with_default(subscriber, || { - tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap() - .block_on(JsonRpcClient::timing_loop( - phase_rx, - terminal_rx, - dropped_records, - )); - }); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_writer(blocking_writer) + .with_ansi(false) + .without_time() + .with_filter(tracing_subscriber::filter::filter_fn(|metadata| { + metadata.target() == REVERSE_RPC_TIMING_TARGET + })), + ); + let timing_thread = tracing::subscriber::with_default(subscriber, || { + JsonRpcClient::spawn_timing_thread(phase_rx, terminal_rx, dropped_records) }); let (notification_tx, _) = broadcast::channel(1); let (request_tx, _request_rx) = mpsc::unbounded_channel(); @@ -2768,7 +2798,7 @@ mod tests { ); let request = JsonRpcRequest::new(53, "hooks.invoke", None); let now = TokioInstant::now(); - let trace = ReverseRpcTrace::new(&request, 1, now, &RandomState::new(), timing); + let trace = ReverseRpcTrace::new(&request, 1, now, &CorrelationHasher::new(), timing); trace.mark_forwarding(now); tokio::time::timeout( diff --git a/rust/src/session.rs b/rust/src/session.rs index 399c9cec73..58835bfa6d 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1609,6 +1609,7 @@ fn spawn_event_loop( .tx .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into())); } + requests.close(); while let Ok(request) = requests.try_recv() { drop(request); } From aef0dea2022b5e5f0cef388201d7a5070f1149d4 Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 28 Aug 2026 18:35:41 -0600 Subject: [PATCH 7/7] Harden reverse timing shutdown Ensure timing startup fails closed, synchronize terminal ownership through force-close, and let the dedicated timing thread drain and exit even when handlers retain trace senders. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/jsonrpc.rs | 415 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 319 insertions(+), 96 deletions(-) diff --git a/rust/src/jsonrpc.rs b/rust/src/jsonrpc.rs index e90357a691..1dc4e9368f 100644 --- a/rust/src/jsonrpc.rs +++ b/rust/src/jsonrpc.rs @@ -12,6 +12,7 @@ use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWrite use tokio::sync::{broadcast, mpsc, oneshot}; use tokio::task::JoinHandle; use tokio::time::Instant as TokioInstant; +use tokio_util::sync::CancellationToken; use tracing::{Instrument, debug, error, warn}; use crate::{Error, ErrorKind, ProtocolErrorKind}; @@ -265,12 +266,14 @@ struct WriteCommand { struct ReverseRpcWriteTrace { trace: ReverseRpcTrace, + registry: Option, completion_attempted: bool, } struct ReverseRpcResponseTrace { request_id: u64, trace: ReverseRpcTrace, + registry: ReverseRpcRegistry, } // Response helpers retain their existing signatures while the dispatch task @@ -280,23 +283,32 @@ tokio::task_local! { } impl ReverseRpcWriteTrace { - fn new(trace: ReverseRpcTrace) -> Self { + fn new(trace: ReverseRpcTrace, registry: Option) -> Self { Self { trace, + registry, completion_attempted: false, } } fn record_complete(&mut self, completed_at: TokioInstant, succeeded: bool) { self.completion_attempted = true; - let _ = self.trace.record_complete(completed_at, succeeded); + if let Some(registry) = &self.registry { + registry.complete(&self.trace, completed_at, succeeded); + } else { + let _ = self.trace.record_complete(completed_at, succeeded); + } } } impl Drop for ReverseRpcWriteTrace { fn drop(&mut self) { if !self.completion_attempted { - let _ = self.trace.record_complete(TokioInstant::now(), false); + if let Some(registry) = &self.registry { + registry.complete(&self.trace, TokioInstant::now(), false); + } else { + let _ = self.trace.record_complete(TokioInstant::now(), false); + } } } } @@ -441,6 +453,7 @@ impl ReverseRpcTrace { phase_rx, terminal_rx, dropped_records.clone(), + CancellationToken::new(), )); let trace = Self::new( request, @@ -606,11 +619,17 @@ impl ReverseRpcTrace { self.record_complete_with_state(&mut state, completed_at, succeeded) } - fn record_abandoned(&self, completed_at: TokioInstant) { + fn record_abandoned(&self, completed_at: TokioInstant) -> bool { let mut state = self.inner.timing_state.lock(); - if !state.writer_owns_completion { - let _ = self.record_complete_with_state(&mut state, completed_at, false); + if state.writer_owns_completion { + return false; } + let _ = self.record_complete_with_state(&mut state, completed_at, false); + true + } + + fn record_force_abandoned(&self, completed_at: TokioInstant) { + let _ = self.record_complete(completed_at, false); } fn record_complete_with_state( @@ -655,57 +674,119 @@ impl ReverseRpcTrace { struct ReverseRpcRegistry(Arc); struct ReverseRpcRegistryInner { - traces: Mutex>, + state: Mutex, next_generation: AtomicU64, } +struct ReverseRpcRegistryState { + traces: Vec, + force_closed: bool, +} + impl ReverseRpcRegistry { fn new() -> Self { Self(Arc::new(ReverseRpcRegistryInner { - traces: Mutex::new(Vec::new()), + state: Mutex::new(ReverseRpcRegistryState { + traces: Vec::new(), + force_closed: false, + }), next_generation: AtomicU64::new(1), })) } - fn next_generation(&self) -> u64 { - self.0.next_generation.fetch_add(1, Ordering::Relaxed) + fn register( + &self, + request: &JsonRpcRequest, + received_at: TokioInstant, + correlation_hasher: &CorrelationHasher, + timing: ReverseRpcTimingEmitter, + ) -> Option { + let mut state = self.0.state.lock(); + if state.force_closed { + return None; + } + let trace = ReverseRpcTrace::new( + request, + self.0.next_generation.fetch_add(1, Ordering::Relaxed), + received_at, + correlation_hasher, + timing, + ); + state.traces.push(trace.clone()); + Some(trace) } + #[cfg(test)] fn insert(&self, trace: ReverseRpcTrace) { - self.0.traces.lock().push(trace); + let mut state = self.0.state.lock(); + assert!(!state.force_closed); + state.traces.push(trace); } - fn remove(&self, trace: &ReverseRpcTrace) -> bool { - let mut traces = self.0.traces.lock(); - let Some(index) = traces + fn abandon(&self, trace: &ReverseRpcTrace) { + let mut state = self.0.state.lock(); + let Some(index) = state + .traces .iter() .position(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) else { - return false; + return; }; - traces.swap_remove(index); - true + if trace.record_abandoned(TokioInstant::now()) { + state.traces.swap_remove(index); + } + } + + fn complete(&self, trace: &ReverseRpcTrace, completed_at: TokioInstant, succeeded: bool) { + let mut state = self.0.state.lock(); + let Some(index) = state + .traces + .iter() + .position(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) + else { + drop(state); + let _ = trace.record_complete(completed_at, succeeded); + return; + }; + let _ = trace.record_complete(completed_at, succeeded); + state.traces.swap_remove(index); } fn abandon_all(&self) { - let traces = std::mem::take(&mut *self.0.traces.lock()); - for trace in traces { - trace.record_abandoned(TokioInstant::now()); + let mut state = self.0.state.lock(); + let mut index = 0; + while index < state.traces.len() { + let trace = state.traces[index].clone(); + if trace.record_abandoned(TokioInstant::now()) { + state.traces.swap_remove(index); + } else { + index += 1; + } + } + } + + fn force_abandon_all(&self) { + let mut state = self.0.state.lock(); + state.force_closed = true; + for trace in &state.traces { + trace.record_force_abandoned(TokioInstant::now()); } + state.traces.clear(); } #[cfg(test)] fn contains(&self, trace: &ReverseRpcTrace) -> bool { self.0 - .traces + .state .lock() + .traces .iter() .any(|current| Arc::ptr_eq(¤t.inner, &trace.inner)) } #[cfg(test)] fn is_empty(&self) -> bool { - self.0.traces.lock().is_empty() + self.0.state.lock().traces.is_empty() } } @@ -761,12 +842,9 @@ impl ReverseRpcRequest { impl Drop for ReverseRpcRequest { fn drop(&mut self) { if let Some(trace) = self.trace.take() - && self - .registry - .as_ref() - .is_some_and(|registry| registry.remove(&trace)) + && let Some(registry) = &self.registry { - trace.record_abandoned(TokioInstant::now()); + registry.abandon(&trace); } } } @@ -788,6 +866,7 @@ impl ReverseRpcDispatchGuard { ReverseRpcResponseTrace { request_id: self.request_id, trace: self.trace.clone(), + registry: self.registry.clone(), }, future, ) @@ -797,9 +876,7 @@ impl ReverseRpcDispatchGuard { impl Drop for ReverseRpcDispatchGuard { fn drop(&mut self) { - if self.registry.remove(&self.trace) { - self.trace.record_abandoned(TokioInstant::now()); - } + self.registry.abandon(&self.trace); } } @@ -835,6 +912,7 @@ pub struct JsonRpcClient { read_task: Mutex>>, write_task: Mutex>>, timing_task: Mutex>>, + timing_shutdown: Option, } impl JsonRpcClient { @@ -896,29 +974,27 @@ impl JsonRpcClient { let writer_span = tracing::error_span!("jsonrpc_write_loop"); let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span)); - let (timing, timing_task, correlation_hasher, reverse_requests) = if trace_reverse_rpc { - let (phase_tx, phase_rx) = - mpsc::channel::(REVERSE_RPC_TIMING_CAPACITY); - let (terminal_tx, terminal_rx) = - mpsc::channel::(REVERSE_RPC_TIMING_CAPACITY); - let dropped_records = Arc::new(AtomicU64::new(0)); - ( - Some(ReverseRpcTimingEmitter { - phase_tx, - terminal_tx, - dropped_records: dropped_records.clone(), - }), - Some(Self::spawn_timing_thread( - phase_rx, - terminal_rx, - dropped_records, - )), - Some(CorrelationHasher::new()), - Some(ReverseRpcRegistry::new()), - ) - } else { - (None, None, None, None) - }; + let (timing, timing_task, timing_shutdown, correlation_hasher, reverse_requests) = + if trace_reverse_rpc { + match Self::start_reverse_rpc_timing_with(Self::spawn_timing_thread) { + Some(( + timing, + timing_task, + timing_shutdown, + correlation_hasher, + reverse_requests, + )) => ( + Some(timing), + Some(timing_task), + Some(timing_shutdown), + Some(correlation_hasher), + Some(reverse_requests), + ), + None => (None, None, None, None, None), + } + } else { + (None, None, None, None, None) + }; let client = Self { request_id: AtomicU64::new(1), @@ -930,6 +1006,7 @@ impl JsonRpcClient { read_task: Mutex::new(None), write_task: Mutex::new(Some(write_task)), timing_task: Mutex::new(timing_task), + timing_shutdown, }; let pending_requests = client.pending_requests.clone(); @@ -958,24 +1035,87 @@ impl JsonRpcClient { client } + fn start_reverse_rpc_timing_with( + spawn: impl FnOnce( + mpsc::Receiver, + mpsc::Receiver, + Arc, + CancellationToken, + ) -> std::io::Result>, + ) -> Option<( + ReverseRpcTimingEmitter, + std::thread::JoinHandle<()>, + CancellationToken, + CorrelationHasher, + ReverseRpcRegistry, + )> { + let (phase_tx, phase_rx) = + mpsc::channel::(REVERSE_RPC_TIMING_CAPACITY); + let (terminal_tx, terminal_rx) = + mpsc::channel::(REVERSE_RPC_TIMING_CAPACITY); + let dropped_records = Arc::new(AtomicU64::new(0)); + let timing_shutdown = CancellationToken::new(); + match spawn( + phase_rx, + terminal_rx, + dropped_records.clone(), + timing_shutdown.clone(), + ) { + Ok(timing_task) => Some(( + ReverseRpcTimingEmitter { + phase_tx, + terminal_tx, + dropped_records, + }, + timing_task, + timing_shutdown, + CorrelationHasher::new(), + ReverseRpcRegistry::new(), + )), + Err(error) => { + warn!( + error = %error, + "failed to start reverse RPC timing thread; timing disabled" + ); + None + } + } + } + fn spawn_timing_thread( phase_rx: mpsc::Receiver, terminal_rx: mpsc::Receiver, dropped_records: Arc, - ) -> std::thread::JoinHandle<()> { + shutdown: CancellationToken, + ) -> std::io::Result> { let dispatch = tracing::dispatcher::get_default(Clone::clone); - std::thread::Builder::new() + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let runtime_slot = Arc::new(Mutex::new(Some(runtime))); + let thread_runtime_slot = runtime_slot.clone(); + let result = std::thread::Builder::new() .name("copilot-reverse-rpc-timing".to_string()) .spawn(move || { tracing::dispatcher::with_default(&dispatch, || { - tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .expect("reverse RPC timing runtime should start") - .block_on(Self::timing_loop(phase_rx, terminal_rx, dropped_records)); + let runtime = thread_runtime_slot + .lock() + .take() + .expect("reverse RPC timing runtime must be transferred once"); + runtime.block_on(Self::timing_loop( + phase_rx, + terminal_rx, + dropped_records, + shutdown, + )); }); - }) - .expect("reverse RPC timing thread should start") + }); + if result.is_err() + && let Some(runtime) = runtime_slot.lock().take() + { + runtime.shutdown_background(); + } + result } pub(crate) fn force_close(&self) { @@ -987,11 +1127,13 @@ impl JsonRpcClient { } self.pending_requests.write().clear(); if let Some(reverse_requests) = &self.reverse_requests { - reverse_requests.abandon_all(); + reverse_requests.force_abandon_all(); } - // Detach the timing task so it can drain the bounded queue. The - // aborted read/write tasks drop the remaining senders, so it exits - // once those final diagnostics are emitted. + if let Some(timing_shutdown) = &self.timing_shutdown { + timing_shutdown.cancel(); + } + // The timing thread closes its receivers, drains accepted records, + // and exits without waiting for traces held by in-flight handlers. let _ = self.timing_task.lock().take(); } @@ -999,14 +1141,21 @@ impl JsonRpcClient { mut phase_rx: mpsc::Receiver, mut terminal_rx: mpsc::Receiver, dropped_records: Arc, + shutdown: CancellationToken, ) { let mut phase_closed = false; let mut terminal_closed = false; + let mut shutting_down = false; let mut pending_terminals = VecDeque::new(); while !phase_closed || !terminal_closed { Self::record_dropped_timing_records(&dropped_records); tokio::select! { biased; + _ = shutdown.cancelled(), if !shutting_down => { + shutting_down = true; + phase_rx.close(); + terminal_rx.close(); + } event = terminal_rx.recv(), if !terminal_closed => { if let Some(event) = event { Self::record_or_defer_timing_event(event, &mut pending_terminals); @@ -1220,11 +1369,6 @@ impl JsonRpcClient { let completed_at = reverse_rpc.as_ref().map(|_| TokioInstant::now()); let succeeded = result.is_ok(); - // Caller may have dropped the ack receiver (e.g. their - // `await` was cancelled); that's fine — we still completed - // the write, which was the whole point. - let _ = ack.send(result); - if let Some(write_trace) = &mut reverse_rpc { let trace = &write_trace.trace; let (enqueued_at, queue_elapsed) = @@ -1241,6 +1385,11 @@ impl JsonRpcClient { succeeded, ); } + + // Caller may have dropped the ack receiver (e.g. their + // `await` was cancelled); that's fine — we still completed + // the write, which was the whole point. + let _ = ack.send(result); } } @@ -1324,10 +1473,9 @@ impl JsonRpcClient { .as_ref() .zip(correlation_hasher.as_ref()) .zip(reverse_requests.as_ref()) - .map(|((timing, correlation_hasher), registry)| { - ReverseRpcTrace::new( + .and_then(|((timing, correlation_hasher), registry)| { + registry.register( &request, - registry.next_generation(), TokioInstant::now(), correlation_hasher, timing.clone(), @@ -1336,12 +1484,6 @@ impl JsonRpcClient { } else { None }; - if let Some(trace) = &trace { - reverse_requests - .as_ref() - .expect("timed requests must have a registry") - .insert(trace.clone()); - } let forwarded = match &request_tx { ReverseRequestSender::Public(request_tx) => { request_tx.send(request).is_ok() @@ -1574,27 +1716,28 @@ impl JsonRpcClient { /// drops the ack receiver; the actor still completes the frame and /// flushes. A partial frame can never appear on the wire. pub async fn write(&self, message: &T) -> Result<(), Error> { - self.write_frame(message, None).await + self.write_frame(message, None, None).await } pub(crate) async fn write_response(&self, response: &JsonRpcResponse) -> Result<(), Error> { - let trace = REVERSE_RPC_RESPONSE_TRACE - .try_with(|scoped| (scoped.request_id == response.id).then(|| scoped.trace.clone())) + let reverse_rpc = REVERSE_RPC_RESPONSE_TRACE + .try_with(|scoped| { + (scoped.request_id == response.id) + .then(|| (scoped.trace.clone(), scoped.registry.clone())) + }) .ok() .flatten(); - let result = self.write_frame(response, trace.clone()).await; - if let Some(trace) = &trace - && let Some(reverse_requests) = &self.reverse_requests - { - let _ = reverse_requests.remove(trace); - } - result + let (trace, registry) = reverse_rpc + .map(|(trace, registry)| (Some(trace), Some(registry))) + .unwrap_or((None, None)); + self.write_frame(response, trace, registry).await } async fn write_frame( &self, message: &T, reverse_rpc: Option, + reverse_rpc_registry: Option, ) -> Result<(), Error> { let encode_start = reverse_rpc.as_ref().map(|_| TokioInstant::now()); let encoded = serde_json::to_vec(message); @@ -1626,7 +1769,8 @@ impl JsonRpcClient { .send(WriteCommand { frame, ack: ack_tx, - reverse_rpc: reverse_rpc.map(ReverseRpcWriteTrace::new), + reverse_rpc: reverse_rpc + .map(|trace| ReverseRpcWriteTrace::new(trace, reverse_rpc_registry)), enqueued_at, }) .is_err() @@ -2080,6 +2224,7 @@ mod tests { phase_rx, terminal_rx, dropped_records, + CancellationToken::new(), )); let request = JsonRpcRequest::new( 29, @@ -2122,6 +2267,7 @@ mod tests { phase_rx, terminal_rx, dropped_records, + CancellationToken::new(), )); let request = JsonRpcRequest::new(31, "hooks.invoke", None); let now = TokioInstant::now(); @@ -2205,6 +2351,7 @@ mod tests { assert_eq!(forwarded.id, request.id); assert!(client.reverse_requests.is_none()); assert!(client.timing_task.lock().is_none()); + assert!(client.timing_shutdown.is_none()); client.force_close(); } @@ -2234,9 +2381,37 @@ mod tests { assert!(forwarded.trace.is_none()); assert!(client.reverse_requests.is_none()); assert!(client.timing_task.lock().is_none()); + assert!(client.timing_shutdown.is_none()); client.force_close(); } + #[test] + fn timing_thread_spawn_failure_disables_timing_without_panicking() { + let timing = JsonRpcClient::start_reverse_rpc_timing_with(|_, _, _, _| { + Err(std::io::Error::other("injected timing thread failure")) + }); + + assert!(timing.is_none()); + } + + #[test] + fn force_closed_registry_rejects_late_reverse_requests() { + let (timing, _phase_rx, _terminal_rx, _dropped_records) = + timing_channel(REVERSE_RPC_TIMING_CAPACITY); + let registry = ReverseRpcRegistry::new(); + registry.force_abandon_all(); + + let trace = registry.register( + &JsonRpcRequest::new(35, "hooks.invoke", None), + TokioInstant::now(), + &CorrelationHasher::new(), + timing, + ); + + assert!(trace.is_none()); + assert!(registry.is_empty()); + } + #[tokio::test] async fn saturated_timing_queue_drops_records_and_reports_the_count() { let trace_buffer = TraceBuffer::default(); @@ -2259,6 +2434,7 @@ mod tests { phase_rx, terminal_rx, dropped_records, + CancellationToken::new(), )); wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; @@ -2298,7 +2474,7 @@ mod tests { let writer = ReverseRpcTrace::new(&writer_request, 2, now, &CorrelationHasher::new(), timing); writer.mark_forwarding(now); - let mut writer_trace = ReverseRpcWriteTrace::new(writer); + let mut writer_trace = ReverseRpcWriteTrace::new(writer, None); writer_trace.record_complete(now, true); drop(writer_trace); @@ -2311,6 +2487,7 @@ mod tests { phase_rx, terminal_rx, dropped_records, + CancellationToken::new(), )); wait_for_trace(&trace_buffer, "phase=\"records_dropped\"").await; @@ -2373,6 +2550,7 @@ mod tests { error: None, }, Some(writer), + None, ) .await .unwrap_err(); @@ -2386,6 +2564,7 @@ mod tests { phase_rx, terminal_rx, dropped_records, + CancellationToken::new(), )); wait_for_trace(&trace_buffer, "phase=\"records_dropped\"").await; @@ -2518,7 +2697,7 @@ mod tests { error: None, }), ack: second_ack_tx, - reverse_rpc: Some(ReverseRpcWriteTrace::new(trace)), + reverse_rpc: Some(ReverseRpcWriteTrace::new(trace, None)), enqueued_at: Some(TokioInstant::now()), }) .unwrap(); @@ -2752,10 +2931,22 @@ mod tests { }); assert_eq!(started_rx.recv().await, Some("write")); + response_task.abort(); + assert!(response_task.await.unwrap_err().is_cancelled()); + assert!( + client + .reverse_requests + .as_ref() + .is_some_and(|registry| !registry.is_empty()) + ); + let timing_thread = client + .timing_task + .lock() + .take() + .expect("enabled timing target should start the timing thread"); client.force_close(); - assert!(response_task.await.unwrap().is_err()); + timing_thread.join().unwrap(); - wait_for_trace(&trace_buffer, "phase=\"request_complete\"").await; let complete = trace_buffer .text() .lines() @@ -2786,7 +2977,13 @@ mod tests { })), ); let timing_thread = tracing::subscriber::with_default(subscriber, || { - JsonRpcClient::spawn_timing_thread(phase_rx, terminal_rx, dropped_records) + JsonRpcClient::spawn_timing_thread( + phase_rx, + terminal_rx, + dropped_records, + CancellationToken::new(), + ) + .unwrap() }); let (notification_tx, _) = broadcast::channel(1); let (request_tx, _request_rx) = mpsc::unbounded_channel(); @@ -2811,6 +3008,7 @@ mod tests { error: None, }, Some(trace), + None, ), ) .await @@ -2826,4 +3024,29 @@ mod tests { client.force_close(); timing_thread.join().unwrap(); } + + #[test] + fn timing_thread_shutdown_does_not_wait_for_trace_senders() { + let (timing, phase_rx, terminal_rx, dropped_records) = + timing_channel(REVERSE_RPC_TIMING_CAPACITY); + let shutdown = CancellationToken::new(); + let timing_thread = JsonRpcClient::spawn_timing_thread( + phase_rx, + terminal_rx, + dropped_records, + shutdown.clone(), + ) + .unwrap(); + let (joined_tx, joined_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + timing_thread.join().unwrap(); + joined_tx.send(()).unwrap(); + }); + + shutdown.cancel(); + joined_rx + .recv_timeout(Duration::from_secs(1)) + .expect("timing thread should exit while timing senders remain alive"); + drop(timing); + } }