phi_agent/bridge/
server.rs

1//! Protocol server — adapts [`AgentRuntime`] to the bridge protocol.
2//!
3//! Tool calls use a single-slot pattern: the serve loop pushes a receiver
4//! before each tool call, and ProxyTool pops it.  This handles sequential
5//! tool calls cleanly; parallel calls can be added later via a FIFO queue.
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use agent_base::{
11    AgentBuilder, AgentResult, AgentRuntime, RunOutcome, RuntimeEvent, SessionId, Tool, ToolContext, ToolControlFlow,
12    ToolMetadata, ToolOutput,
13};
14use async_trait::async_trait;
15use serde_json::Value;
16use tokio::sync::{Mutex, mpsc};
17
18#[derive(Clone)]
19pub struct ProtocolServer {
20    runtime: AgentRuntime,
21    /// Single-slot: the next tool call's response receiver.
22    /// serve loop pushes, ProxyTool pops.
23    slot: Arc<Mutex<Option<mpsc::UnboundedReceiver<AgentResult<ToolOutput>>>>>,
24    /// Map external_id → SessionId so that runs with the same
25    /// external_id reuse the same session.
26    sessions: Arc<Mutex<HashMap<String, SessionId>>>,
27}
28
29impl ProtocolServer {
30    pub fn new(runtime: AgentRuntime) -> Self {
31        Self {
32            runtime,
33            slot: Arc::new(Mutex::new(None)),
34            sessions: Arc::new(Mutex::new(HashMap::new())),
35        }
36    }
37
38    pub fn from_builder(builder: AgentBuilder) -> Result<Self, agent_base::AgentError> {
39        let runtime = builder.build()?;
40        Ok(Self::new(runtime))
41    }
42
43    /// Register a Python-side tool.
44    pub async fn register_tool(&self, name: String, _description: String, _parameters: Value) {
45        let proxy = ProxyTool { name, slot: self.slot.clone() };
46        let tools_arc = self.runtime.tools_mut();
47        let mut tools = tools_arc.write().await;
48        tools.register(proxy);
49    }
50
51    /// Set up the response channel for the NEXT tool call.
52    /// Returns the sender — keep it; send the result when the SDK replies.
53    pub async fn prepare_tool_call(&self) -> mpsc::UnboundedSender<AgentResult<ToolOutput>> {
54        let (tx, rx) = mpsc::unbounded_channel();
55        *self.slot.lock().await = Some(rx);
56        tx
57    }
58
59    pub async fn create_session(&self, external_id: Option<String>) -> (SessionId, Option<String>) {
60        let sid = self.runtime.create_session().await;
61        // NOTE: We intentionally do NOT set sid.external_id because
62        // agent_base's run_turn() hangs when external_id is Some.
63        // Session reuse is handled by get_or_create_session() which
64        // maintains its own external_id → SessionId map.
65        let ext = external_id.clone();
66        (sid, ext)
67    }
68
69    /// Get or create a session by external_id.
70    ///
71    /// If ``external_id`` is ``Some`` and a session with that id already
72    /// exists, it is reused (preserving conversation history).  Otherwise
73    /// a new session is created and registered.
74    pub async fn get_or_create_session(
75        &self,
76        external_id: Option<String>,
77    ) -> SessionId {
78        if let Some(ref ext) = external_id {
79            let mut sessions = self.sessions.lock().await;
80            if let Some(sid) = sessions.get(ext) {
81                return sid.clone();
82            }
83            // Create new and register
84            let (sid, _) = self.create_session(Some(ext.clone())).await;
85            sessions.insert(ext.clone(), sid.clone());
86            return sid;
87        }
88        // No external_id — always create new
89        self.create_session(None).await.0
90    }
91
92    pub fn subscribe_events(&self) -> tokio::sync::broadcast::Receiver<RuntimeEvent> {
93        self.runtime.subscribe_runtime_events()
94    }
95
96    pub async fn run_turn<F>(&self, sid: &SessionId, input: &str, f: F) -> AgentResult<RunOutcome>
97    where
98        F: FnMut(RuntimeEvent) -> AgentResult<()> + Send,
99    {
100        self.runtime.run_turn(sid.clone(), input, f).await
101    }
102
103    pub fn cancel(&self) {
104        self.runtime.cancel();
105    }
106
107    /// List all registered tools with their metadata, sorted by name.
108    pub async fn list_tools(&self) -> Vec<ToolMetadata> {
109        let tools = self.runtime.tools_mut();
110        let registry = tools.read().await;
111        registry.metadatas()
112    }
113}
114
115// ── ProxyTool ─────────────────────────────────────────────────────────
116
117struct ProxyTool {
118    name: String,
119    slot: Arc<Mutex<Option<mpsc::UnboundedReceiver<AgentResult<ToolOutput>>>>>,
120}
121
122#[async_trait]
123impl Tool for ProxyTool {
124    fn name(&self) -> &'static str {
125        Box::leak(self.name.clone().into_boxed_str())
126    }
127
128    fn definition(&self) -> Value {
129        serde_json::json!({
130            "type": "function",
131            "function": {
132                "name": self.name,
133                "description": "Proxy tool",
134                "parameters": { "type": "object", "properties": {} }
135            }
136        })
137    }
138
139    async fn call(&self, _args: &Value, _ctx: &ToolContext) -> AgentResult<ToolOutput> {
140        let mut rx = self
141            .slot
142            .lock()
143            .await
144            .take()
145            .ok_or_else(|| agent_base::AgentError::internal("no tool call slot prepared"))?;
146
147        match rx.recv().await {
148            Some(result) => result,
149            None => Ok(ToolOutput {
150                summary: "Tool call cancelled".to_string(),
151                raw: None,
152                control_flow: ToolControlFlow::Break,
153                truncation: None,
154            }),
155        }
156    }
157}
phi_agent/bridge/server.rs

phi_agent/bridge/
server.rs
