antigravity_sdk_rust/
connection.rs

1//! Connection abstractions for communication with the agent execution harness.
2//!
3//! This module defines the [`Connection`] trait, which provides a transport-agnostic interface
4//! for sending prompts, receiving step updates, and executing tool and question handshakes.
5
6use crate::types::{QuestionHookResult, Step, ToolResult};
7#[cfg(test)]
8use futures_util::StreamExt;
9use futures_util::stream::BoxStream;
10
11/// Trait representing an active session connection with the agentic harness.
12///
13/// A connection manages process or network lifecycle, handshaking, and event streaming.
14/// Common implementations include [`LocalConnection`](crate::local::LocalConnection) which spawns
15/// a local helper subprocess and upgrades communication to a WebSocket protocol.
16pub trait Connection: Send + Sync {
17    /// Returns the conversation ID for the connection.
18    fn conversation_id(&self) -> &str;
19
20    /// Returns whether the connection is currently idle.
21    fn is_idle(&self) -> bool;
22
23    /// Subscribes to the stream of step updates from the connection.
24    fn receive_steps(&self) -> BoxStream<'static, Result<Step, anyhow::Error>>;
25
26    /// Sends a standard user text prompt to the connection.
27    fn send(
28        &self,
29        content: &str,
30    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
31
32    /// Sends a trigger notification message to the connection.
33    fn send_trigger_notification(
34        &self,
35        content: &str,
36    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
37
38    /// Sends a halt/cancellation request to the connection.
39    fn send_halt_request(
40        &self,
41    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
42
43    /// Sends approval/rejection for a tool execution confirmation request.
44    fn send_tool_confirmation(
45        &self,
46        trajectory_id: &str,
47        step_index: u32,
48        accepted: bool,
49    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
50
51    /// Sends the result of a tool execution back to the connection.
52    fn send_tool_response(
53        &self,
54        id: &str,
55        result: ToolResult,
56    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
57
58    /// Sends the answers to user questions back to the connection.
59    fn send_question_response(
60        &self,
61        trajectory_id: &str,
62        step_index: u32,
63        answers: QuestionHookResult,
64    ) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
65
66    /// Gracefully closes the connection.
67    fn disconnect(&self) -> impl std::future::Future<Output = Result<(), anyhow::Error>> + Send;
68}
69
70/// A target-agnostic wrapper enum for active connection types.
71#[derive(Clone)]
72pub enum AnyConnection {
73    #[cfg(not(target_arch = "wasm32"))]
74    Local(std::sync::Arc<crate::local::LocalConnection>),
75    #[cfg(target_arch = "wasm32")]
76    Wasm(std::sync::Arc<crate::wasm::WasmConnection>),
77    #[cfg(test)]
78    Mock(std::sync::Arc<MockConnection>),
79}
80
81impl std::fmt::Debug for AnyConnection {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        match self {
84            #[cfg(not(target_arch = "wasm32"))]
85            Self::Local(_) => f.write_str("AnyConnection::Local"),
86            #[cfg(target_arch = "wasm32")]
87            Self::Wasm(_) => f.write_str("AnyConnection::Wasm"),
88            #[cfg(test)]
89            Self::Mock(_) => f.write_str("AnyConnection::Mock"),
90        }
91    }
92}
93
94impl Connection for AnyConnection {
95    fn conversation_id(&self) -> &str {
96        match self {
97            #[cfg(not(target_arch = "wasm32"))]
98            Self::Local(c) => c.conversation_id(),
99            #[cfg(target_arch = "wasm32")]
100            Self::Wasm(c) => c.conversation_id(),
101            #[cfg(test)]
102            Self::Mock(c) => c.conversation_id(),
103        }
104    }
105
106    fn is_idle(&self) -> bool {
107        match self {
108            #[cfg(not(target_arch = "wasm32"))]
109            Self::Local(c) => c.is_idle(),
110            #[cfg(target_arch = "wasm32")]
111            Self::Wasm(c) => c.is_idle(),
112            #[cfg(test)]
113            Self::Mock(c) => c.is_idle(),
114        }
115    }
116
117    fn receive_steps(&self) -> BoxStream<'static, Result<Step, anyhow::Error>> {
118        match self {
119            #[cfg(not(target_arch = "wasm32"))]
120            Self::Local(c) => c.receive_steps(),
121            #[cfg(target_arch = "wasm32")]
122            Self::Wasm(c) => c.receive_steps(),
123            #[cfg(test)]
124            Self::Mock(c) => c.receive_steps(),
125        }
126    }
127
128    async fn send(&self, content: &str) -> Result<(), anyhow::Error> {
129        match self {
130            #[cfg(not(target_arch = "wasm32"))]
131            Self::Local(c) => c.send(content).await,
132            #[cfg(target_arch = "wasm32")]
133            Self::Wasm(c) => c.send(content).await,
134            #[cfg(test)]
135            Self::Mock(c) => c.send(content).await,
136        }
137    }
138
139    async fn send_trigger_notification(&self, content: &str) -> Result<(), anyhow::Error> {
140        match self {
141            #[cfg(not(target_arch = "wasm32"))]
142            Self::Local(c) => c.send_trigger_notification(content).await,
143            #[cfg(target_arch = "wasm32")]
144            Self::Wasm(c) => c.send_trigger_notification(content).await,
145            #[cfg(test)]
146            Self::Mock(c) => c.send_trigger_notification(content).await,
147        }
148    }
149
150    async fn send_halt_request(&self) -> Result<(), anyhow::Error> {
151        match self {
152            #[cfg(not(target_arch = "wasm32"))]
153            Self::Local(c) => c.send_halt_request().await,
154            #[cfg(target_arch = "wasm32")]
155            Self::Wasm(c) => c.send_halt_request().await,
156            #[cfg(test)]
157            Self::Mock(c) => c.send_halt_request().await,
158        }
159    }
160
161    async fn send_tool_confirmation(
162        &self,
163        trajectory_id: &str,
164        step_index: u32,
165        accepted: bool,
166    ) -> Result<(), anyhow::Error> {
167        match self {
168            #[cfg(not(target_arch = "wasm32"))]
169            Self::Local(c) => {
170                c.send_tool_confirmation(trajectory_id, step_index, accepted)
171                    .await
172            }
173            #[cfg(target_arch = "wasm32")]
174            Self::Wasm(c) => {
175                c.send_tool_confirmation(trajectory_id, step_index, accepted)
176                    .await
177            }
178            #[cfg(test)]
179            Self::Mock(c) => {
180                c.send_tool_confirmation(trajectory_id, step_index, accepted)
181                    .await
182            }
183        }
184    }
185
186    async fn send_tool_response(&self, id: &str, result: ToolResult) -> Result<(), anyhow::Error> {
187        match self {
188            #[cfg(not(target_arch = "wasm32"))]
189            Self::Local(c) => c.send_tool_response(id, result).await,
190            #[cfg(target_arch = "wasm32")]
191            Self::Wasm(c) => c.send_tool_response(id, result).await,
192            #[cfg(test)]
193            Self::Mock(c) => c.send_tool_response(id, result).await,
194        }
195    }
196
197    async fn send_question_response(
198        &self,
199        trajectory_id: &str,
200        step_index: u32,
201        answers: QuestionHookResult,
202    ) -> Result<(), anyhow::Error> {
203        match self {
204            #[cfg(not(target_arch = "wasm32"))]
205            Self::Local(c) => {
206                c.send_question_response(trajectory_id, step_index, answers)
207                    .await
208            }
209            #[cfg(target_arch = "wasm32")]
210            Self::Wasm(c) => {
211                c.send_question_response(trajectory_id, step_index, answers)
212                    .await
213            }
214            #[cfg(test)]
215            Self::Mock(c) => {
216                c.send_question_response(trajectory_id, step_index, answers)
217                    .await
218            }
219        }
220    }
221
222    async fn disconnect(&self) -> Result<(), anyhow::Error> {
223        match self {
224            #[cfg(not(target_arch = "wasm32"))]
225            Self::Local(c) => c.disconnect().await,
226            #[cfg(target_arch = "wasm32")]
227            Self::Wasm(c) => c.disconnect().await,
228            #[cfg(test)]
229            Self::Mock(c) => c.disconnect().await,
230        }
231    }
232}
233
234#[cfg(test)]
235pub struct MockConnection {
236    pub id: String,
237    pub is_idle: std::sync::atomic::AtomicBool,
238    pub steps_to_yield: std::sync::Mutex<Vec<Step>>,
239    pub sent_prompts: std::sync::Mutex<Vec<String>>,
240}
241
242#[cfg(test)]
243impl std::fmt::Debug for MockConnection {
244    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245        f.debug_struct("MockConnection")
246            .field("id", &self.id)
247            .field(
248                "is_idle",
249                &self.is_idle.load(std::sync::atomic::Ordering::Relaxed),
250            )
251            .finish_non_exhaustive()
252    }
253}
254
255#[cfg(test)]
256impl MockConnection {
257    pub fn new(id: &str) -> Self {
258        Self {
259            id: id.to_string(),
260            is_idle: std::sync::atomic::AtomicBool::new(true),
261            steps_to_yield: std::sync::Mutex::new(Vec::new()),
262            sent_prompts: std::sync::Mutex::new(Vec::new()),
263        }
264    }
265}
266
267#[cfg(test)]
268impl Connection for MockConnection {
269    fn conversation_id(&self) -> &str {
270        &self.id
271    }
272
273    fn is_idle(&self) -> bool {
274        self.is_idle.load(std::sync::atomic::Ordering::SeqCst)
275    }
276
277    fn receive_steps(&self) -> BoxStream<'static, Result<Step, anyhow::Error>> {
278        let steps = self
279            .steps_to_yield
280            .lock()
281            .unwrap_or_else(std::sync::PoisonError::into_inner)
282            .clone();
283        futures_util::stream::iter(steps).map(Ok).boxed()
284    }
285
286    async fn send(&self, content: &str) -> Result<(), anyhow::Error> {
287        self.sent_prompts
288            .lock()
289            .unwrap_or_else(std::sync::PoisonError::into_inner)
290            .push(content.to_string());
291        Ok(())
292    }
293
294    async fn send_trigger_notification(&self, _content: &str) -> Result<(), anyhow::Error> {
295        Ok(())
296    }
297
298    async fn send_halt_request(&self) -> Result<(), anyhow::Error> {
299        Ok(())
300    }
301
302    async fn send_tool_confirmation(
303        &self,
304        _trajectory_id: &str,
305        _step_index: u32,
306        _accepted: bool,
307    ) -> Result<(), anyhow::Error> {
308        Ok(())
309    }
310
311    async fn send_tool_response(
312        &self,
313        _id: &str,
314        _result: ToolResult,
315    ) -> Result<(), anyhow::Error> {
316        Ok(())
317    }
318
319    async fn send_question_response(
320        &self,
321        _trajectory_id: &str,
322        _step_index: u32,
323        _answers: QuestionHookResult,
324    ) -> Result<(), anyhow::Error> {
325        Ok(())
326    }
327
328    async fn disconnect(&self) -> Result<(), anyhow::Error> {
329        Ok(())
330    }
331}
antigravity_sdk_rust/connection.rs

antigravity_sdk_rust/
connection.rs
