Skip to main content

orchestral_runtime/
agent_sdk.rs

1//! Provider-neutral SDK surface over the Agent Protocol control plane.
2//!
3//! This layer owns no planning or model loop. CLI, HTTP adapters, and embedded
4//! applications can therefore share exactly the same `AgentController`
5//! sequencing, commands, inspection, and durable event semantics.
6
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9
10use orchestral_core::agent_connector::AgentSessionActionInvocation;
11use orchestral_core::agent_protocol::{
12    reference::AgentRunStatus,
13    wire::{
14        AgentCommand, AgentCommandEnvelope, AgentExecutionRef, AgentJournalRecord,
15        AgentRunEnvelope, AgentRunView, AgentSessionId, CommandAck, CommandId, Content,
16        ContentBody, Extensions, RequestId, RequestResolution, ResourceBinding, RunId,
17    },
18    AGENT_PROTOCOL_V1,
19};
20use tokio::sync::broadcast;
21
22use crate::agent_control::{AgentControlError, AgentControlEvent, AgentController};
23
24#[derive(Clone)]
25pub struct AgentClient {
26    controller: Arc<AgentController>,
27    session_id: AgentSessionId,
28    resources: Arc<Vec<ResourceBinding>>,
29    default_extensions: Arc<Extensions>,
30    next_run: Arc<AtomicU64>,
31}
32
33impl AgentClient {
34    pub fn new(controller: Arc<AgentController>, session_id: AgentSessionId) -> Self {
35        Self {
36            controller,
37            session_id,
38            resources: Arc::new(Vec::new()),
39            default_extensions: Arc::new(Extensions::new()),
40            next_run: Arc::new(AtomicU64::new(1)),
41        }
42    }
43
44    pub fn with_resources(mut self, resources: Vec<ResourceBinding>) -> Self {
45        self.resources = Arc::new(resources);
46        self
47    }
48
49    pub fn session_id(&self) -> &AgentSessionId {
50        &self.session_id
51    }
52
53    /// Binds Host metadata to every new Run. Per-Run values cannot replace a
54    /// conflicting Host value. Extensions are descriptive, not Tool grants.
55    pub fn with_default_extensions(mut self, extensions: Extensions) -> Self {
56        self.default_extensions = Arc::new(extensions);
57        self
58    }
59
60    /// Reattaches a durable Run in this Session through the Controller's
61    /// existing recovery contract. Never starts a replacement Run.
62    pub async fn resume_run(&self, run_id: &RunId) -> Result<AgentRunHandle, AgentSdkError> {
63        let mut view = self.controller.inspect(run_id).await?;
64        if view.execution.session_id != self.session_id {
65            return Err(AgentSdkError::InvalidInput(
66                "Run belongs to another Session".to_owned(),
67            ));
68        }
69        if matches!(
70            view.state,
71            orchestral_core::agent_protocol::wire::AgentRunState::Unknown { .. }
72        ) {
73            view = self.controller.recover(run_id).await?;
74        }
75        Ok(AgentRunHandle {
76            controller: self.controller.clone(),
77            run_id: run_id.clone(),
78            execution: view.execution,
79        })
80    }
81
82    fn merged_extensions(&self, extensions: Extensions) -> Result<Extensions, AgentSdkError> {
83        let mut merged = self.default_extensions.as_ref().clone();
84        for (key, value) in extensions {
85            if merged.get(&key).is_some_and(|existing| existing != &value) {
86                return Err(AgentSdkError::InvalidInput(format!(
87                    "Run extension conflicts with Host metadata: {key}"
88                )));
89            }
90            merged.insert(key, value);
91        }
92        Ok(merged)
93    }
94
95    pub fn controller(&self) -> &Arc<AgentController> {
96        &self.controller
97    }
98
99    pub async fn start_text(
100        &self,
101        input: impl Into<String>,
102    ) -> Result<AgentRunHandle, AgentSdkError> {
103        let input = input.into();
104        if input.trim().is_empty() {
105            return Err(AgentSdkError::InvalidInput(
106                "Agent input must not be empty".to_owned(),
107            ));
108        }
109        let sequence = self.next_run.fetch_add(1, Ordering::SeqCst);
110        self.start_with_run_id(
111            RunId::new(format!(
112                "sdk-{}-{sequence}-{}",
113                self.session_id.as_str(),
114                uuid::Uuid::new_v4()
115            )),
116            vec![Content::text(input)],
117        )
118        .await
119    }
120
121    pub async fn start_with_run_id(
122        &self,
123        run_id: RunId,
124        input: Vec<Content>,
125    ) -> Result<AgentRunHandle, AgentSdkError> {
126        self.start_with_run_id_and_extensions(run_id, input, Extensions::new())
127            .await
128    }
129
130    pub async fn start_with_run_id_and_extensions(
131        &self,
132        run_id: RunId,
133        input: Vec<Content>,
134        extensions: Extensions,
135    ) -> Result<AgentRunHandle, AgentSdkError> {
136        let mut run = AgentRunEnvelope::new(
137            AGENT_PROTOCOL_V1,
138            self.session_id.clone(),
139            run_id.clone(),
140            input,
141        )?;
142        run.spec.resources = self.resources.as_ref().clone();
143        run.spec.extensions = self.merged_extensions(extensions)?;
144        let run = AgentRunEnvelope::seal(run.spec)?;
145        let execution = self.controller.start(run).await?;
146        Ok(AgentRunHandle {
147            controller: self.controller.clone(),
148            run_id,
149            execution,
150        })
151    }
152
153    pub async fn start_session_action_with_run_id(
154        &self,
155        run_id: RunId,
156        title: impl Into<String>,
157        action: AgentSessionActionInvocation,
158    ) -> Result<AgentRunHandle, AgentSdkError> {
159        let title = title.into();
160        if title.trim().is_empty() {
161            return Err(AgentSdkError::InvalidInput(
162                "Agent session action title must not be empty".to_owned(),
163            ));
164        }
165        let mut run = AgentRunEnvelope::new(
166            AGENT_PROTOCOL_V1,
167            self.session_id.clone(),
168            run_id.clone(),
169            vec![Content::text(title)],
170        )?;
171        run.spec.resources = self.resources.as_ref().clone();
172        action
173            .insert_into(&mut run.spec)
174            .map_err(|error| AgentSdkError::InvalidInput(error.to_string()))?;
175        run.spec.extensions = self.merged_extensions(run.spec.extensions)?;
176        let run = AgentRunEnvelope::seal(run.spec)?;
177        let execution = self.controller.start(run).await?;
178        Ok(AgentRunHandle {
179            controller: self.controller.clone(),
180            run_id,
181            execution,
182        })
183    }
184
185    /// Convenience path for non-interactive calls. It returns on a terminal
186    /// state or when Host input/approval is required, never by bypassing the
187    /// Agent request protocol.
188    pub async fn run_text(&self, input: impl Into<String>) -> Result<AgentTurn, AgentSdkError> {
189        self.start_text(input).await?.wait_until_blocked().await
190    }
191}
192
193#[derive(Clone)]
194pub struct AgentRunHandle {
195    controller: Arc<AgentController>,
196    run_id: RunId,
197    execution: AgentExecutionRef,
198}
199
200impl AgentRunHandle {
201    pub fn run_id(&self) -> &RunId {
202        &self.run_id
203    }
204
205    pub fn execution(&self) -> &AgentExecutionRef {
206        &self.execution
207    }
208
209    pub async fn inspect(&self) -> Result<AgentRunView, AgentSdkError> {
210        Ok(self.controller.inspect(&self.run_id).await?)
211    }
212
213    pub async fn events(
214        &self,
215        after_run_seq: u64,
216    ) -> Result<Vec<AgentJournalRecord>, AgentSdkError> {
217        Ok(self.controller.events(&self.run_id, after_run_seq).await?)
218    }
219
220    pub async fn subscribe(&self) -> Result<broadcast::Receiver<AgentControlEvent>, AgentSdkError> {
221        Ok(self.controller.subscribe(&self.run_id).await?)
222    }
223
224    pub async fn command(
225        &self,
226        command: AgentCommandEnvelope,
227    ) -> Result<CommandAck, AgentSdkError> {
228        if command.run_id != self.run_id {
229            return Err(AgentSdkError::InvalidInput(
230                "Agent command belongs to another Run".to_owned(),
231            ));
232        }
233        Ok(self.controller.command(command).await?)
234    }
235
236    pub async fn cancel(&self, reason: impl Into<String>) -> Result<CommandAck, AgentSdkError> {
237        Ok(self.controller.cancel(&self.run_id, reason).await?)
238    }
239
240    pub async fn steer_text(&self, input: impl Into<String>) -> Result<CommandAck, AgentSdkError> {
241        let input = input.into();
242        if input.trim().is_empty() {
243            return Err(AgentSdkError::InvalidInput(
244                "Steer input must not be empty".to_owned(),
245            ));
246        }
247        let command = AgentCommandEnvelope::new(
248            CommandId::new(format!("sdk-steer-{}", uuid::Uuid::new_v4())),
249            self.run_id.clone(),
250            None,
251            AgentCommand::Steer {
252                content: vec![Content::text(input)],
253            },
254        )?;
255        self.command(command).await
256    }
257
258    pub async fn resolve_input_text(
259        &self,
260        request_id: RequestId,
261        input: impl Into<String>,
262    ) -> Result<CommandAck, AgentSdkError> {
263        let input = input.into();
264        if input.trim().is_empty() {
265            return Err(AgentSdkError::InvalidInput(
266                "Input resolution must not be empty".to_owned(),
267            ));
268        }
269        let command = AgentCommandEnvelope::new(
270            CommandId::new(format!("sdk-input-{}", uuid::Uuid::new_v4())),
271            self.run_id.clone(),
272            Some(request_id),
273            AgentCommand::ResolveRequest {
274                response: RequestResolution::Input {
275                    content: vec![Content::text(input)],
276                },
277            },
278        )?;
279        self.command(command).await
280    }
281
282    pub async fn wait_until_blocked(&self) -> Result<AgentTurn, AgentSdkError> {
283        let mut events = self.subscribe().await?;
284        loop {
285            let view = self.inspect().await?;
286            if is_stable_turn_boundary(&view) {
287                return Ok(AgentTurn {
288                    run_id: self.run_id.clone(),
289                    view,
290                });
291            }
292            match events.recv().await {
293                Ok(_) | Err(broadcast::error::RecvError::Lagged(_)) => {}
294                Err(broadcast::error::RecvError::Closed) => {
295                    let view = self.inspect().await?;
296                    if is_stable_turn_boundary(&view) {
297                        return Ok(AgentTurn {
298                            run_id: self.run_id.clone(),
299                            view,
300                        });
301                    }
302                    return Err(AgentSdkError::ControlStreamClosed(self.run_id.clone()));
303                }
304            }
305        }
306    }
307}
308
309fn is_stable_turn_boundary(view: &AgentRunView) -> bool {
310    view.state.is_terminal()
311        || view.state.status() == AgentRunStatus::Unknown
312        || !view.pending_requests.is_empty()
313}
314
315#[derive(Debug, Clone)]
316pub struct AgentTurn {
317    pub run_id: RunId,
318    pub view: AgentRunView,
319}
320
321impl AgentTurn {
322    pub fn status(&self) -> AgentRunStatus {
323        self.view.state.status()
324    }
325
326    pub fn is_waiting(&self) -> bool {
327        !self.view.state.is_terminal() && !self.view.pending_requests.is_empty()
328    }
329
330    pub fn final_text(&self) -> Option<&str> {
331        let delivery = self.view.delivery.as_ref()?;
332        match &delivery.final_response.body {
333            ContentBody::Inline(serde_json::Value::String(text)) => Some(text),
334            _ => None,
335        }
336    }
337}
338
339#[derive(Debug, thiserror::Error)]
340#[non_exhaustive]
341pub enum AgentSdkError {
342    #[error("invalid Agent SDK input: {0}")]
343    InvalidInput(String),
344    #[error(transparent)]
345    Protocol(#[from] orchestral_core::agent_protocol::wire::AgentProtocolError),
346    #[error(transparent)]
347    Control(#[from] AgentControlError),
348    #[error("Agent control stream closed before Run {0} reached a stable boundary")]
349    ControlStreamClosed(RunId),
350}