Rust SDK: PR #1367 review follow-ups (#1382)

Author: SteveSandersonMS

rust/README.md

@@ -375,7 +375,7 @@ Tools are named types (not closures) — visible in stack traces and navigable v
 
 Tools without an attached handler (`Tool::with_handler` never called) are declaration-only: the SDK advertises them on the wire but doesn't dispatch invocations to anything. Useful when another connected client services the tool.
 
-For trivial tools that don't need a named type, [`define_tool`](crate::tool::define_tool) collapses the definition to a single expression and returns a fully-formed `Tool` with handler attached:
+For trivial tools that don't need a named type, the `define_tool` helper function (available with the `derive` feature) collapses the definition to a single expression and returns a fully-formed `Tool` with handler attached:
 
 ```rust,ignore
 use github_copilot_sdk::tool::{define_tool, JsonSchema};
@@ -672,7 +672,7 @@ ergonomics the dynamically-typed SDKs don't.
   `Session` value to thread in, and the SDK already prefers traits over
   boxed closures for handler-shaped APIs (`PermissionHandler`, `ToolHandler`,
   `SessionHooks`,
-  `ToolHandler`).
+  `SystemMessageTransform`).
 
 ```rust,ignore
 use std::sync::Arc;

rust/src/lib.rs

@@ -937,26 +937,20 @@ impl Client {
         // to the server. For Tcp, the SDK auto-generates one when the
         // caller leaves it unset so the loopback listener is safe by
         // default.
-        let (mut options, effective_connection_token) = {
-            let mut options = options;
-            let effective = match &mut options.transport {
-                Transport::Stdio => None,
-                Transport::Tcp {
-                    connection_token, ..
-                } => {
-                    if connection_token.is_none() {
-                        *connection_token = Some(generate_connection_token());
-                    }
-                    connection_token.clone()
-                }
-                Transport::External {
-                    connection_token, ..
-                } => connection_token.clone(),
-            };
-            (options, effective)
+        let mut options = options;
+        let effective_connection_token: Option<String> = match &mut options.transport {
+            Transport::Stdio => None,
+            Transport::Tcp {
+                connection_token, ..
+            } => Some(
+                connection_token
+                    .get_or_insert_with(generate_connection_token)
+                    .clone(),
+            ),
+            Transport::External {
+                connection_token, ..
+            } => connection_token.clone(),
         };
-        let _ = &mut options;
-        let effective_connection_token: Option<String> = effective_connection_token;
         let session_fs_config = options.session_fs.clone();
         let session_fs_sqlite_declared = session_fs_config
             .as_ref()

rust/src/session.rs

@@ -122,7 +122,7 @@ impl Drop for PendingSessionRegistration {
 /// Owns an internal event loop that dispatches events to the per-callback
 /// handlers installed on the session config.
 ///
-/// Protocol methods (`send`, `get_messages`, `abort`, etc.) automatically
+/// Protocol methods (`send`, `get_events`, `abort`, etc.) automatically
 /// inject the session ID into RPC params.
 ///
 /// Call [`destroy`](Self::destroy) for graceful cleanup (RPC + local). If dropped
@@ -788,45 +788,27 @@ impl Client {
         if let Some(transforms) = config.system_message_transform.clone() {
             inject_transform_sections(&mut config, transforms.as_ref());
         }
-        let wire = config.to_wire(session_id.clone());
+        let (wire, mut runtime) = config.into_wire(session_id.clone())?;
 
         let permission_handler = crate::permission::resolve_handler(
-            config.permission_handler.take(),
-            config.permission_policy.take(),
+            runtime.permission_handler.take(),
+            runtime.permission_policy.take(),
         );
-        let elicitation_handler = config.elicitation_handler.take();
-        let user_input_handler = config.user_input_handler.take();
-        let exit_plan_mode_handler = config.exit_plan_mode_handler.take();
-        let auto_mode_switch_handler = config.auto_mode_switch_handler.take();
-        let mut tool_map: HashMap<String, Arc<dyn crate::tool::ToolHandler>> = HashMap::new();
-        if let Some(tools) = config.tools.as_mut() {
-            for tool in tools.iter_mut() {
-                if let Some(handler) = tool.handler.take() {
-                    if tool_map.contains_key(&tool.name) {
-                        return Err(Error::InvalidConfig(format!(
-                            "duplicate tool handler registered for name {:?}",
-                            tool.name
-                        )));
-                    }
-                    tool_map.insert(tool.name.clone(), handler);
-                }
-            }
-        }
         let handlers = SessionHandlers {
             permission: permission_handler,
-            elicitation: elicitation_handler,
-            user_input: user_input_handler,
-            exit_plan_mode: exit_plan_mode_handler,
-            auto_mode_switch: auto_mode_switch_handler,
-            tools: Arc::new(tool_map),
+            elicitation: runtime.elicitation_handler.take(),
+            user_input: runtime.user_input_handler.take(),
+            exit_plan_mode: runtime.exit_plan_mode_handler.take(),
+            auto_mode_switch: runtime.auto_mode_switch_handler.take(),
+            tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
         };
-        let hooks = config.hooks_handler.take();
-        let transforms = config.system_message_transform.take();
-        let tools_count = config.tools.as_ref().map_or(0, Vec::len);
-        let commands_count = config.commands.as_ref().map_or(0, Vec::len);
+        let hooks = runtime.hooks_handler.take();
+        let transforms = runtime.system_message_transform.take();
+        let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
+        let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
         let has_hooks = hooks.is_some();
-        let command_handlers = build_command_handler_map(config.commands.as_deref());
-        let session_fs_provider = config.session_fs_provider.take();
+        let command_handlers = build_command_handler_map(runtime.commands.as_deref());
+        let session_fs_provider = runtime.session_fs_provider.take();
         if self.inner.session_fs_configured && session_fs_provider.is_none() {
             return Err(Error::Session(SessionError::SessionFsProviderRequired));
         }
@@ -943,45 +925,27 @@ impl Client {
         if let Some(transforms) = config.system_message_transform.clone() {
             inject_transform_sections_resume(&mut config, transforms.as_ref());
         }
-        let wire = config.to_wire();
+        let (wire, mut runtime) = config.into_wire()?;
 
         let permission_handler = crate::permission::resolve_handler(
-            config.permission_handler.take(),
-            config.permission_policy.take(),
+            runtime.permission_handler.take(),
+            runtime.permission_policy.take(),
         );
-        let elicitation_handler = config.elicitation_handler.take();
-        let user_input_handler = config.user_input_handler.take();
-        let exit_plan_mode_handler = config.exit_plan_mode_handler.take();
-        let auto_mode_switch_handler = config.auto_mode_switch_handler.take();
-        let mut tool_map: HashMap<String, Arc<dyn crate::tool::ToolHandler>> = HashMap::new();
-        if let Some(tools) = config.tools.as_mut() {
-            for tool in tools.iter_mut() {
-                if let Some(handler) = tool.handler.take() {
-                    if tool_map.contains_key(&tool.name) {
-                        return Err(Error::InvalidConfig(format!(
-                            "duplicate tool handler registered for name {:?}",
-                            tool.name
-                        )));
-                    }
-                    tool_map.insert(tool.name.clone(), handler);
-                }
-            }
-        }
         let handlers = SessionHandlers {
             permission: permission_handler,
-            elicitation: elicitation_handler,
-            user_input: user_input_handler,
-            exit_plan_mode: exit_plan_mode_handler,
-            auto_mode_switch: auto_mode_switch_handler,
-            tools: Arc::new(tool_map),
+            elicitation: runtime.elicitation_handler.take(),
+            user_input: runtime.user_input_handler.take(),
+            exit_plan_mode: runtime.exit_plan_mode_handler.take(),
+            auto_mode_switch: runtime.auto_mode_switch_handler.take(),
+            tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
         };
-        let hooks = config.hooks_handler.take();
-        let transforms = config.system_message_transform.take();
-        let tools_count = config.tools.as_ref().map_or(0, Vec::len);
-        let commands_count = config.commands.as_ref().map_or(0, Vec::len);
+        let hooks = runtime.hooks_handler.take();
+        let transforms = runtime.system_message_transform.take();
+        let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
+        let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
         let has_hooks = hooks.is_some();
-        let command_handlers = build_command_handler_map(config.commands.as_deref());
-        let session_fs_provider = config.session_fs_provider.take();
+        let command_handlers = build_command_handler_map(runtime.commands.as_deref());
+        let session_fs_provider = runtime.session_fs_provider.take();
         if self.inner.session_fs_configured && session_fs_provider.is_none() {
             return Err(Error::Session(SessionError::SessionFsProviderRequired));
         }
@@ -1464,12 +1428,12 @@ async fn handle_notification(
             );
             tokio::spawn(
                 async move {
-                    if data.tool_call_id.is_empty() || data.tool_name.is_empty() {
-                        let error_msg = if data.tool_call_id.is_empty() {
-                            "Missing toolCallId"
-                        } else {
-                            "Missing toolName"
-                        };
+                    // `tool_name.is_empty()` would have produced a `None`
+                    // lookup in `handlers.tools` and short-circuited at the
+                    // outer guard above, so only the tool_call_id check is
+                    // reachable here.
+                    if data.tool_call_id.is_empty() {
+                        let error_msg = "Missing toolCallId";
                         let rpc_start = Instant::now();
                         let _ = client
                             .call(

rust/src/tool.rs

@@ -57,11 +57,13 @@ pub fn schema_for<T: schemars::JsonSchema>() -> serde_json::Value {
 }
 
 /// Convert a JSON Schema [`Value`](serde_json::Value) into the
-/// [`Tool::parameters`] map shape expected by the protocol.
+/// [`Tool::parameters`](crate::types::Tool::parameters) map shape
+/// expected by the protocol.
 ///
 /// Panics if the input is not a JSON object — tool parameter schemas
 /// are always top-level objects (`{"type": "object", ...}`). Pair with
-/// [`schema_for`] or a `serde_json::json!(...)` literal.
+/// `schema_for` (available with the `derive` feature) or a
+/// `serde_json::json!(...)` literal.
 ///
 /// Use [`try_tool_parameters`] when the schema comes from dynamic input and
 /// should return a recoverable error instead of panicking.
@@ -179,13 +181,15 @@ pub fn convert_mcp_call_tool_result(value: &serde_json::Value) -> Option<ToolRes
 ///
 /// Implement this trait when you want to bind a Rust function to a tool
 /// name and have the SDK dispatch matching `external_tool.requested`
-/// broadcasts to it. Attach the impl to a [`Tool`] via [`Tool::with_handler`].
+/// broadcasts to it. Attach the impl to a [`Tool`](crate::types::Tool)
+/// via [`Tool::with_handler`](crate::types::Tool::with_handler).
 ///
 /// Named handler types (e.g. `struct MyTool;`) are visible in stack
 /// traces and navigable via "go to definition", which is preferable to
 /// closure-based alternatives for non-trivial tools. For trivial tools,
-/// [`define_tool`] wraps a free `async fn` or closure into a [`Tool`]
-/// with the handler already attached.
+/// the `define_tool` helper function (available with the `derive`
+/// feature) wraps a free `async fn` or closure into a [`Tool`](crate::types::Tool) with
+/// the handler already attached.
 ///
 /// # Example
 ///