forge_sdk/client/
control_dispatch.rs

1//! Inbound `control_request` dispatch: [`ControlDispatchHandle`]
2//! routes permission checks, MCP JSON-RPC, and hook callbacks to the
3//! appropriate handler and writes the matching `control_response` via
4//! a clonable [`SharedWriter`].
5//!
6//! Internal: built once during [`Client::spawn`](crate::Client::spawn),
7//! cloned and moved into a `tokio::spawn`'d task per inbound
8//! `control_request` by the reader task in [`crate::client::runtime`].
9//! The clonable writer + the `tokio::spawn` mean a slow callback cannot
10//! block the read loop, and cancellation of the actor's `select!` cannot
11//! drop a `control_response` write mid-flight.
12//!
13//! Used the same way during the synchronous init handshake - the
14//! init loop calls [`ControlDispatchHandle::dispatch`] directly for
15//! interleaved `control_request`s, since no concurrent reader exists
16//! yet.
17
18use std::collections::HashMap;
19use std::sync::Arc;
20
21use crate::Error;
22use crate::control::{
23    AllowBehavior, ControlRequest, ControlRequestKind, ControlResponse, ControlResponseKind,
24    ControlResponseType,
25};
26use crate::hooks::callback::ErasedHookCallback;
27use crate::hooks::{HookContext, HookDecision, HookKind};
28use crate::mcp::orchestration::{Dispatched, McpHosts};
29use crate::mcp::protocol::JsonRpcRequest;
30use crate::permissions::CanUseToolCallback;
31use crate::transport::process::SharedWriter;
32use forge_primitives::hooks::outputs::encode_updated_input_wrapper;
33use forge_primitives::{PermissionDecision, ToolPermissionContext};
34
35// =============================================================================
36// Detached dispatch - inbound `control_request`s go through
37// `dispatch`, which writes the matching `control_response` via a
38// clonable [`SharedWriter`]. The reader task in
39// [`crate::client::runtime`] `tokio::spawn`s a fresh task per
40// inbound request so a slow callback can't block the read loop AND
41// cancellation of the actor's `select!` over a command channel +
42// `next_event` cannot drop the response write mid-flight.
43// =============================================================================
44
45/// Clonable bundle of state + writer that dispatches a single
46/// `control_request`. Internal - built once during
47/// [`Client::spawn`](crate::Client::spawn) and cloned by the reader
48/// task per inbound request.
49///
50/// Each field is `Arc`-backed (or `Clone`); cloning the handle is
51/// cheap. Designed to be moved into a `tokio::spawn`'d task per
52/// inbound `control_request`.
53#[derive(Clone)]
54pub(crate) struct ControlDispatchHandle {
55    writer: Arc<SharedWriter>,
56    can_use_tool: Option<Arc<dyn CanUseToolCallback>>,
57    auto_approve_tool: Option<crate::options::AutoApproveToolPredicate>,
58    mcp_hosts: McpHosts,
59    hook_callbacks: HashMap<String, Arc<dyn ErasedHookCallback>>,
60    session_id: crate::client::runtime::SharedSessionId,
61}
62
63impl std::fmt::Debug for ControlDispatchHandle {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct("ControlDispatchHandle")
66            .field("writer", &self.writer)
67            .field("can_use_tool", &self.can_use_tool.as_ref().map(|_| "<callback>"))
68            .field("auto_approve_tool", &self.auto_approve_tool.as_ref().map(|_| "<predicate>"))
69            .field("mcp_hosts", &self.mcp_hosts)
70            .field("hook_callbacks", &format!("<{} hooks>", self.hook_callbacks.len()))
71            .field("session_id", &"<shared>")
72            .finish()
73    }
74}
75
76/// The subtype a request crossed as, for a failure record that has to say
77/// which kind of control request went unanswered.
78///
79/// **An unknown subtype keeps its own string**: that is the one case the
80/// record is most useful for, and folding it to a literal "unknown" would
81/// discard the only name the CLI gave it.
82fn control_subtype(kind: &ControlRequestKind) -> std::borrow::Cow<'_, str> {
83    match kind {
84        ControlRequestKind::CanUseTool { .. } => std::borrow::Cow::Borrowed("can_use_tool"),
85        ControlRequestKind::McpMessage { .. } => std::borrow::Cow::Borrowed("mcp_message"),
86        ControlRequestKind::HookCallback { .. } => std::borrow::Cow::Borrowed("hook_callback"),
87        ControlRequestKind::Unknown { subtype, .. } => std::borrow::Cow::Borrowed(subtype.as_str()),
88    }
89}
90
91impl ControlDispatchHandle {
92    pub(crate) fn new(
93        writer: Arc<SharedWriter>,
94        can_use_tool: Option<Arc<dyn CanUseToolCallback>>,
95        auto_approve_tool: Option<crate::options::AutoApproveToolPredicate>,
96        mcp_hosts: McpHosts,
97        hook_callbacks: HashMap<String, Arc<dyn ErasedHookCallback>>,
98        session_id: crate::client::runtime::SharedSessionId,
99    ) -> Self {
100        Self { writer, can_use_tool, auto_approve_tool, mcp_hosts, hook_callbacks, session_id }
101    }
102
103    /// Capture a session id from an incoming Message. No-op once a
104    /// non-empty id has been bound. Called by the reader task on every
105    /// message.
106    pub(crate) fn capture_session_id_from(&self, msg: &forge_primitives::Message) {
107        if let Some(id) = msg.session_id()
108            && !id.is_empty()
109        {
110            let mut current = self.session_id.write();
111            if current.is_empty() {
112                *current = id.to_string();
113                tracing::debug!(session_id = %*current, "client session_id bound");
114            }
115        }
116    }
117
118    /// Dispatch one inbound `control_request`. Routes the request
119    /// to the right handler (MCP, hook callback, `can_use_tool`) and
120    /// writes the matching `control_response` via the cloned
121    /// [`SharedWriter`]. Safe to call from a `tokio::spawn`'d task -
122    /// runs to completion regardless of caller cancellation.
123    ///
124    /// On handler failure the CLI is best-effort answered with an
125    /// error `control_response` before the error surfaces - an
126    /// unanswered request hangs the subprocess's in-flight tool call.
127    ///
128    /// # Errors
129    ///
130    /// Same shape as `Client::handle_control`: encode/serialise
131    /// failures from [`Error::message_parse`] / [`Error::encode`],
132    /// and write failures from the underlying writer.
133    pub(crate) async fn dispatch(&self, req: ControlRequest) -> Result<(), Error> {
134        let request_id = req.request_id.clone();
135        // Owned before the request moves into the handler: an unknown
136        // subtype's own string is the record's only name for it.
137        let subtype = control_subtype(&req.request).into_owned();
138        let result = self.dispatch_inner(req).await;
139        if let Err(err) = &result {
140            // **The record this failure gets, and the only one that can name
141            // the session.** A stalled stdin surfaces here as a hung turn,
142            // and the error response written next is the CLI's only answer -
143            // unless that write fails too, which its own record carries.
144            tracing::warn!(
145                target: crate::logging::targets::BRIDGE_LIFECYCLE,
146                event_name = "control_dispatch_failed",
147                session_id = %self.session_id.read(),
148                %request_id,
149                subtype = %subtype,
150                error = %err,
151                "a control request failed; the error response is the CLI's only answer",
152            );
153            self.write_error_response(&request_id, &err.to_string()).await;
154        }
155        result
156    }
157
158    /// Best-effort error `control_response` for a request we could not
159    /// handle. The CLI blocks on an unanswered request, so any frame
160    /// carrying a `request_id` must be answered even on failure.
161    pub(crate) async fn write_error_response(&self, request_id: &str, error: &str) {
162        let resp = ControlResponse {
163            ty: ControlResponseType::ControlResponse,
164            response: ControlResponseKind::Error {
165                request_id: request_id.to_string(),
166                error: error.to_string(),
167            },
168        };
169        match serde_json::to_string(&resp) {
170            Ok(mut line) => {
171                line.push('\n');
172                if let Err(write_err) = self.writer.write_line(&line).await {
173                    // **The CLI stays blocked on this request**, so the
174                    // record names the session and the request it is blocked
175                    // on: without them a hung turn has no owner.
176                    tracing::warn!(
177                        target: crate::logging::targets::BRIDGE_LIFECYCLE,
178                        event_name = "control_error_response_failed",
179                        session_id = %self.session_id.read(),
180                        %request_id,
181                        error = %write_err,
182                        "the error control_response could not be written; the CLI is waiting on it",
183                    );
184                }
185            }
186            Err(e) => {
187                tracing::debug!(error = %e, "error control_response encode failed");
188            }
189        }
190    }
191
192    async fn dispatch_inner(&self, req: ControlRequest) -> Result<(), Error> {
193        let original_input = req.original_tool_input().cloned();
194
195        if let ControlRequestKind::Unknown { subtype, raw } = &req.request {
196            tracing::warn!(
197                %subtype,
198                raw = %raw,
199                request_id = %req.request_id,
200                "unknown control_request subtype - responding with error, session continues"
201            );
202            return self.write_unsupported(&req).await;
203        }
204
205        if let ControlRequestKind::McpMessage { server_name, message } = &req.request {
206            return self.handle_mcp(&req, server_name, message).await;
207        }
208
209        if let ControlRequestKind::HookCallback { callback_id, input, tool_use_id } = &req.request {
210            return self.handle_hook(&req, callback_id, input, tool_use_id.as_deref()).await;
211        }
212
213        let decision = match (&self.can_use_tool, &req.request) {
214            (
215                Some(cb),
216                ControlRequestKind::CanUseTool {
217                    tool_name,
218                    input,
219                    tool_use_id,
220                    agent_id,
221                    permission_suggestions,
222                    blocked_path,
223                    decision_reason,
224                    title,
225                    display_name,
226                    description,
227                },
228            ) => {
229                // Auto-approve fast-path. When the caller configured
230                // an `auto_approve_tool` predicate on Options
231                // (typically: "any tool whose name starts with
232                // mcp__<server>__ for one of my in-process MCP
233                // servers"), tools that match skip the permission
234                // prompt. The predicate's call sites stay in the
235                // configurator (forge-agent / forge-workspace) so
236                // forge-sdk doesn't hardcode server names.
237                if self.auto_approve_tool.as_ref().is_some_and(|p| p(tool_name)) {
238                    tracing::trace!(
239                        target: "forge_sdk::control_dispatch",
240                        tool_name = %tool_name,
241                        "auto-approving via Options::auto_approve_tool predicate"
242                    );
243                    PermissionDecision::allow()
244                } else {
245                    let suggestions: Vec<forge_primitives::PermissionUpdate> =
246                        permission_suggestions
247                            .iter()
248                            .filter_map(|v| match serde_json::from_value(v.clone()) {
249                                Ok(s) => Some(s),
250                                Err(e) => {
251                                    tracing::warn!(
252                                        error = %e,
253                                        raw = %v,
254                                        "permission_suggestion failed to decode; dropping (CLI schema drift?)"
255                                    );
256                                    None
257                                }
258                            })
259                            .collect();
260                    let ctx = ToolPermissionContext::new(
261                        tool_name.clone(),
262                        input.clone(),
263                        tool_use_id.clone(),
264                        agent_id.clone(),
265                    )
266                    .with_suggestions(suggestions)
267                    .with_display(
268                        blocked_path.clone(),
269                        decision_reason.clone(),
270                        title.clone(),
271                        display_name.clone(),
272                        description.clone(),
273                    );
274                    cb.call(ctx).await
275                }
276            }
277            (None, ControlRequestKind::CanUseTool { .. }) => {
278                PermissionDecision::deny("no permission callback registered")
279            }
280            _ => {
281                return self.write_unsupported(&req).await;
282            }
283        };
284
285        let behavior = if decision.is_allow() {
286            let perms = decision.updated_permissions();
287            let updated_permissions = if perms.is_empty() {
288                None
289            } else {
290                Some(
291                    serde_json::to_value(perms)
292                        .map_err(|e| Error::encode("updated_permissions", e))?,
293                )
294            };
295            AllowBehavior::Allow {
296                updated_input: decision
297                    .updated_input()
298                    .cloned()
299                    .or(original_input)
300                    .unwrap_or(serde_json::Value::Null),
301                updated_permissions,
302            }
303        } else {
304            AllowBehavior::Deny {
305                message: decision.reason().unwrap_or("denied").to_string(),
306                interrupt: false,
307            }
308        };
309
310        let resp = req
311            .build_response(behavior)
312            .map_err(|e| Error::message_parse(format!("could not build control response: {e}")))?;
313        let mut line = serde_json::to_string(&resp).map_err(|e| {
314            Error::message_parse(format!("could not serialise control response: {e}"))
315        })?;
316        line.push('\n');
317        self.writer.write_line(&line).await
318    }
319
320    async fn write_unsupported(&self, req: &ControlRequest) -> Result<(), Error> {
321        let resp = ControlResponse {
322            ty: ControlResponseType::ControlResponse,
323            response: ControlResponseKind::Error {
324                request_id: req.request_id.clone(),
325                error: "unsupported control-request subtype".into(),
326            },
327        };
328        let mut line = serde_json::to_string(&resp)
329            .map_err(|e| Error::message_parse(format!("error response serialise: {e}")))?;
330        line.push('\n');
331        self.writer.write_line(&line).await
332    }
333
334    async fn handle_mcp(
335        &self,
336        req: &ControlRequest,
337        server_name: &str,
338        message: &serde_json::Value,
339    ) -> Result<(), Error> {
340        let jsonrpc: JsonRpcRequest = serde_json::from_value(message.clone())
341            .map_err(|e| Error::message_parse(format!("bad JSON-RPC envelope: {e}")))?;
342
343        let jsonrpc_response = match self.mcp_hosts.dispatch(server_name, &jsonrpc).await {
344            Dispatched::Response(r) => serde_json::to_value(&r)
345                .map_err(|e| Error::message_parse(format!("mcp response serialise: {e}")))?,
346            Dispatched::Notification => serde_json::json!({"jsonrpc": "2.0", "result": {}}),
347            Dispatched::UnknownServer => {
348                let resp = ControlResponse {
349                    ty: ControlResponseType::ControlResponse,
350                    response: ControlResponseKind::Error {
351                        request_id: req.request_id.clone(),
352                        error: format!("unknown MCP server: {server_name}"),
353                    },
354                };
355                let mut line = serde_json::to_string(&resp)
356                    .map_err(|e| Error::message_parse(format!("error response serialise: {e}")))?;
357                line.push('\n');
358                return self.writer.write_line(&line).await;
359            }
360        };
361
362        let wrapper = serde_json::json!({"mcp_response": jsonrpc_response});
363        let resp = ControlResponse {
364            ty: ControlResponseType::ControlResponse,
365            response: ControlResponseKind::Success {
366                request_id: req.request_id.clone(),
367                response: wrapper,
368            },
369        };
370        let mut line = serde_json::to_string(&resp)
371            .map_err(|e| Error::message_parse(format!("mcp control response serialise: {e}")))?;
372        line.push('\n');
373        self.writer.write_line(&line).await
374    }
375
376    async fn handle_hook(
377        &self,
378        req: &ControlRequest,
379        callback_id: &str,
380        input: &serde_json::Value,
381        tool_use_id: Option<&str>,
382    ) -> Result<(), Error> {
383        let received_at = std::time::Instant::now();
384        let event_name = input
385            .get("hook_event_name")
386            .and_then(serde_json::Value::as_str)
387            .unwrap_or("Unknown")
388            .to_string();
389        let kind = HookKind::from_wire(&event_name);
390        let session_id = self.session_id.read().clone();
391
392        // The three traces below bracket the hook dispatch so a
393        // CLI-side hook timeout (control_cancel_request) can always be
394        // attributed: either forge answered late or the CLI never read
395        // the response.
396        tracing::debug!(
397            target: crate::logging::targets::BRIDGE_LIFECYCLE,
398            session_id = %session_id,
399            hook = %event_name,
400            %callback_id,
401            request_id = %req.request_id,
402            "hook dispatch-received",
403        );
404
405        let ctx = HookContext {
406            kind,
407            tool_name: input
408                .get("tool_name")
409                .and_then(serde_json::Value::as_str)
410                .map(str::to_string),
411            session_id,
412            tool_use_id: tool_use_id.map(str::to_string),
413        };
414
415        let decision = if let Some(cb) = self.hook_callbacks.get(callback_id) {
416            let decision = cb.call_erased(input.clone(), ctx).await;
417            tracing::debug!(
418                target: crate::logging::targets::BRIDGE_LIFECYCLE,
419                session_id = %self.session_id.read(),
420                hook = %event_name,
421                %callback_id,
422                elapsed_ms = u64::try_from(received_at.elapsed().as_millis()).unwrap_or(u64::MAX),
423                "hook callback-done",
424            );
425            decision
426        } else {
427            tracing::warn!(%callback_id, "hook_callback for unknown id; passthrough");
428            HookDecision::passthrough()
429        };
430
431        let ctrl = if decision.is_deferred() {
432            let mut defer_body = serde_json::Map::new();
433            defer_body.insert("async".into(), serde_json::Value::Bool(true));
434            if let Some(timeout) = decision.defer_timeout_ms() {
435                defer_body.insert("asyncTimeout".into(), serde_json::json!(timeout));
436            }
437            ControlResponse {
438                ty: ControlResponseType::ControlResponse,
439                response: ControlResponseKind::Success {
440                    request_id: req.request_id.clone(),
441                    response: serde_json::Value::Object(defer_body),
442                },
443            }
444        } else {
445            let mut response_body = match (decision.is_allow(), decision.reason()) {
446                (true, _) => serde_json::json!({}),
447                (false, Some(reason)) => serde_json::json!({"decision": "block", "reason": reason}),
448                (false, None) => serde_json::json!({"decision": "block"}),
449            };
450
451            if let Some(updated) = decision.updated_input()
452                && let Some(wrapper) = encode_updated_input_wrapper(kind, updated)
453                && let Some(map) = response_body.as_object_mut()
454            {
455                map.insert("hookSpecificOutput".into(), wrapper);
456            }
457
458            ControlResponse {
459                ty: ControlResponseType::ControlResponse,
460                response: ControlResponseKind::Success {
461                    request_id: req.request_id.clone(),
462                    response: response_body,
463                },
464            }
465        };
466
467        let encode_label =
468            if decision.is_deferred() { "deferred hook payload" } else { "hook response" };
469        let mut line = serde_json::to_string(&ctrl).map_err(|e| Error::encode(encode_label, e))?;
470        line.push('\n');
471        self.writer.write_line(&line).await?;
472        tracing::debug!(
473            target: crate::logging::targets::BRIDGE_LIFECYCLE,
474            session_id = %self.session_id.read(),
475            hook = %event_name,
476            %callback_id,
477            elapsed_ms = u64::try_from(received_at.elapsed().as_millis()).unwrap_or(u64::MAX),
478            "hook response-written",
479        );
480        Ok(())
481    }
482}
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487    use crate::control::ControlRequestType;
488    use crate::mcp::McpServerBuilder;
489
490    /// A handler failure (here: unparseable JSON-RPC envelope for a
491    /// registered MCP server) must still answer the CLI with an error
492    /// control_response - an unanswered request_id hangs the
493    /// subprocess's in-flight tool call forever.
494    #[tokio::test]
495    async fn dispatch_answers_with_error_response_on_handler_failure() {
496        let (writer, mut lines) = SharedWriter::test_stub();
497        let hosts = McpHosts::new(
498            vec![("forge".to_owned(), McpServerBuilder::new("forge", "1.0").build())],
499            HashMap::new(),
500        );
501        let handle = ControlDispatchHandle::new(
502            Arc::new(writer),
503            None,
504            None,
505            hosts,
506            HashMap::new(),
507            Arc::new(parking_lot::RwLock::new(String::new())),
508        );
509
510        let req = ControlRequest {
511            ty: ControlRequestType::ControlRequest,
512            request_id: "req-err-1".to_owned(),
513            request: ControlRequestKind::McpMessage {
514                server_name: "forge".to_owned(),
515                message: serde_json::Value::String("not a jsonrpc envelope".to_owned()),
516            },
517        };
518
519        let result = handle.dispatch(req).await;
520        assert!(result.is_err(), "parse failure still surfaces to the caller");
521
522        let line = lines.recv().await.expect("an error control_response was written");
523        let v: serde_json::Value = serde_json::from_str(&line).expect("valid json line");
524        assert_eq!(v["type"], "control_response");
525        assert_eq!(v["response"]["subtype"], "error");
526        assert_eq!(v["response"]["request_id"], "req-err-1");
527    }
528
529    /// An unroutable message must not take the notification path, which
530    /// answers an id-less empty result - that would report success for a
531    /// request no server ever saw.
532    #[tokio::test]
533    async fn unknown_server_answers_error_not_empty_success() {
534        let (writer, mut lines) = SharedWriter::test_stub();
535        let handle = ControlDispatchHandle::new(
536            Arc::new(writer),
537            None,
538            None,
539            McpHosts::new(Vec::new(), HashMap::new()),
540            HashMap::new(),
541            Arc::new(parking_lot::RwLock::new(String::new())),
542        );
543
544        let req = ControlRequest {
545            ty: ControlRequestType::ControlRequest,
546            request_id: "req-unknown-1".to_owned(),
547            request: ControlRequestKind::McpMessage {
548                server_name: "nope".to_owned(),
549                message: serde_json::json!({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}),
550            },
551        };
552        handle.dispatch(req).await.expect("unknown server is answered, not an Err");
553
554        let line = lines.recv().await.expect("a control_response was written");
555        let v: serde_json::Value = serde_json::from_str(&line).expect("valid json line");
556        assert_eq!(v["response"]["subtype"], "error");
557        assert_eq!(v["response"]["request_id"], "req-unknown-1");
558        assert!(v["response"]["response"]["mcp_response"].is_null(), "must not be a success body");
559    }
560
561    /// A control_request whose subtype forge does not know must answer
562    /// the CLI with an error response (so the CLI's awaiting tool call
563    /// unblocks) and must not error the dispatch itself - the session
564    /// continues.
565    #[tokio::test]
566    async fn unknown_subtype_answers_error_and_continues() {
567        let (writer, mut lines) = SharedWriter::test_stub();
568        let handle = ControlDispatchHandle::new(
569            Arc::new(writer),
570            None,
571            None,
572            McpHosts::new(Vec::new(), HashMap::new()),
573            HashMap::new(),
574            Arc::new(parking_lot::RwLock::new(String::new())),
575        );
576
577        let req = ControlRequest {
578            ty: ControlRequestType::ControlRequest,
579            request_id: "req-unknown-subtype-1".to_owned(),
580            request: ControlRequestKind::Unknown {
581                subtype: "hologram_project".to_owned(),
582                raw: serde_json::json!({"subtype": "hologram_project"}),
583            },
584        };
585        handle.dispatch(req).await.expect("unknown subtype is answered, dispatch continues");
586
587        let line = lines.recv().await.expect("an error control_response was written");
588        let v: serde_json::Value = serde_json::from_str(&line).expect("valid json line");
589        assert_eq!(v["type"], "control_response");
590        assert_eq!(v["response"]["subtype"], "error");
591        assert_eq!(v["response"]["request_id"], "req-unknown-subtype-1");
592        assert!(
593            v["response"]["error"].as_str().is_some_and(|e| e.contains("unsupported")),
594            "the error names the unsupported subtype: {v}"
595        );
596    }
597}
forge_sdk/client/control_dispatch.rs

forge_sdk/client/
control_dispatch.rs
