Add reverse RPC and hook callback timing traces

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`)

rust/src/copilot_request_handler.rs

@@ -1047,7 +1047,17 @@ impl CopilotRequestDispatcher {
         self.client.get().cloned().unwrap_or_else(Weak::new)
     }
 
-    pub(crate) async fn dispatch(self: &Arc<Self>, request: JsonRpcRequest) {
+    pub(crate) async fn dispatch(self: &Arc<Self>, 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<Self>, 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,

rust/src/github_token.rs

@@ -188,11 +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 (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()
@@ -201,7 +211,7 @@ impl GitHubTokenRegistry {
             Ok(params) => params,
             Err(error) => {
                 send_error(
-                    &client,
+                    client,
                     request.id,
                     error_codes::INVALID_PARAMS,
                     &format!("invalid params: {error}"),
@@ -218,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",
@@ -232,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",
@@ -252,15 +262,15 @@ impl GitHubTokenRegistry {
         {
             Ok(GitHubTokenProviderResult::Token(token)) => {
                 respond(
-                    &client,
+                    client,
                     request.id,
                     GitHubTokenAcquireResult::Token(token.into_wire()),
                 )
                 .await;
             }
             Ok(GitHubTokenProviderResult::Cancelled) => {
                 respond(
-                    &client,
+                    client,
                     request.id,
                     GitHubTokenAcquireResult::Cancelled(GitHubTokenAcquireResultCancelled {
                         kind: Default::default(),
@@ -270,7 +280,7 @@ impl GitHubTokenRegistry {
             }
             Err(error) => {
                 send_error(
-                    &client,
+                    client,
                     request.id,
                     error_codes::INTERNAL_ERROR,
                     &format!("GitHub token provider failed: {error}"),