cookie_agent_engine/
delegation_api.rs

1use std::{future::Future, pin::Pin};
2
3use cookie_agent_protocol::{
4    AgentId, InvocationId, PersistedToolResult as ToolResult, RunId, SessionId, ToolCallId,
5};
6use serde::{Deserialize, Serialize};
7
8use crate::runtime::{Engine, EngineError};
9
10pub(crate) fn delegate_result_matches_child(
11    result: &ToolResult,
12    child_session_id: SessionId,
13) -> bool {
14    let expected = serde_json::json!(child_session_id);
15    result
16        .metadata
17        .get("session_id")
18        .or_else(|| result.metadata.get("child_session_id"))
19        == Some(&expected)
20}
21
22/// Immutable arguments for one delegate-tool invocation.
23#[derive(Clone, Debug, Deserialize, Serialize)]
24pub struct DelegateInvocation {
25    pub parent_session_id: SessionId,
26    pub parent_run_id: RunId,
27    pub parent_tool_call_id: ToolCallId,
28    pub agent_type: AgentId,
29    pub description: String,
30    pub prompt: String,
31    pub background: bool,
32    /// Handle or full UUID of an existing subagent to resume. Resolved against
33    /// the parent's tree at admission.
34    pub resume_session_id: Option<String>,
35    pub inherit_context: bool,
36}
37
38/// Stable child identity returned to the delegate tool provider.
39#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
40pub struct DelegateHandle {
41    pub invocation_id: InvocationId,
42    pub child_session_id: SessionId,
43    pub child_run_id: Option<RunId>,
44}
45
46/// A delegate wait that cancels its child if its consumer abandons the wait.
47pub struct DelegateAwait {
48    pub(crate) future: Pin<Box<dyn Future<Output = Result<ToolResult, EngineError>> + Send>>,
49    pub(crate) engine: Engine,
50    pub(crate) runtime: Option<tokio::runtime::Handle>,
51    pub(crate) handle: DelegateHandle,
52    pub(crate) completed: bool,
53}
54
55impl Future for DelegateAwait {
56    type Output = Result<ToolResult, EngineError>;
57
58    fn poll(
59        mut self: Pin<&mut Self>,
60        context: &mut std::task::Context<'_>,
61    ) -> std::task::Poll<Self::Output> {
62        let result = self.future.as_mut().poll(context);
63        if result.is_ready() {
64            self.completed = true;
65        }
66        result
67    }
68}
69
70impl Drop for DelegateAwait {
71    fn drop(&mut self) {
72        if self.completed {
73            return;
74        }
75        // Delegate waits are created and polled from the Tokio tool task. If that
76        // task is dropped, retain the cancellation in a detached runtime task.
77        // This closes the abandoned-tool-call child-run leak.
78        if let Some(runtime) = self
79            .runtime
80            .clone()
81            .or_else(|| tokio::runtime::Handle::try_current().ok())
82        {
83            let engine = self.engine.clone();
84            let cancel_engine = engine.clone();
85            let handle = self.handle;
86            let _ = engine.spawn_admission_task(&runtime, async move {
87                let _ = cancel_engine.cancel_delegate(handle).await;
88            });
89        }
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn delegate_result_requires_the_recorded_child_identity() {
99        let child = SessionId::new_v7();
100        let other = SessionId::new_v7();
101        for key in ["session_id", "child_session_id"] {
102            for (value, expected) in [
103                (serde_json::json!(child), true),
104                (serde_json::json!(other), false),
105                (serde_json::Value::Null, false),
106            ] {
107                let result = ToolResult {
108                    display: None,
109                    retained_output: None,
110                    title: cookie_agent_protocol::SafeDisplayText::new("Delegate cancelled")
111                        .unwrap(),
112                    output: String::new(),
113                    metadata: serde_json::json!({key: value}),
114                    truncation: None,
115                    attachments: Vec::new(),
116                    additional_messages: Vec::new(),
117                };
118                assert_eq!(delegate_result_matches_child(&result, child), expected);
119            }
120        }
121    }
122}
cookie_agent_engine/delegation_api.rs

cookie_agent_engine/
delegation_api.rs
