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

phi_agent/bridge/
server.rs
