1use 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 pub fn with_default_extensions(mut self, extensions: Extensions) -> Self {
56 self.default_extensions = Arc::new(extensions);
57 self
58 }
59
60 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 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}