1use std::{future::Future, path::PathBuf, sync::Arc};
4
5use async_trait::async_trait;
6use thiserror::Error;
7use tokio_util::sync::CancellationToken;
8
9use crate::{
10 ApprovalGate, ApprovalRequest, ContextPolicy, CoreEvent, EventSink, HistoryLimits, Message,
11 PROGRESS_INTERVAL, ProgressSink, Provider, ProviderError, ProviderErrorKind, ProviderRequest,
12 TextDeltaSink, ToolApprovals, ToolCall, ToolContext, ToolError, ToolFailure, ToolOutput,
13 ToolRegistry, TurnInput, Usage,
14};
15
16const DENIED: &str = "tool call denied by policy or user";
18
19#[derive(Debug, Clone)]
21pub struct AgentConfig {
22 pub system_prompt: String,
23 pub max_steps: usize,
25 pub history_limits: HistoryLimits,
26}
27
28#[derive(Debug, Clone)]
30pub struct TurnOutcome {
31 pub steps: usize,
32 pub usage: Usage,
33}
34
35#[derive(Debug, Error)]
37pub enum AgentError {
38 #[error("turn cancelled")]
39 Cancelled,
40 #[error("{0}")]
41 Provider(String),
42 #[error("{0}")]
43 ContextLimit(String),
44 #[error("agent reached its maximum step count")]
45 StepLimit,
46 #[error("{0}")]
47 HistoryLimit(String),
48 #[error("{0}")]
49 ResponseLimit(String),
50 #[error("{0}")]
51 ToolLimit(String),
52 #[error("{0}")]
53 Internal(String),
54}
55
56pub struct AgentRuntime {
71 provider: Arc<dyn Provider>,
72 tools: Arc<ToolRegistry>,
73 context: Arc<dyn ContextPolicy>,
74 config: AgentConfig,
75 workspace: PathBuf,
76}
77
78impl AgentRuntime {
79 pub fn new(
80 provider: Arc<dyn Provider>,
81 tools: Arc<ToolRegistry>,
82 context: Arc<dyn ContextPolicy>,
83 config: AgentConfig,
84 workspace: PathBuf,
85 ) -> Self {
86 Self {
87 provider,
88 tools,
89 context,
90 config,
91 workspace,
92 }
93 }
94
95 pub fn model(&self) -> &str {
96 self.provider.model()
97 }
98
99 pub async fn run_turn(
100 &self,
101 history: &mut Vec<Message>,
102 prompt: impl Into<TurnInput>,
103 sink: Arc<dyn EventSink>,
104 approvals: Arc<dyn ApprovalGate>,
105 cancellation: CancellationToken,
106 ) -> Result<TurnOutcome, AgentError> {
107 let checkpoint = history.clone();
108 let result = self
109 .run_turn_inner(history, prompt.into(), sink, approvals, cancellation)
110 .await;
111 if result.is_err() {
112 *history = checkpoint;
113 }
114 result
115 }
116
117 async fn run_turn_inner(
118 &self,
119 history: &mut Vec<Message>,
120 prompt: TurnInput,
121 sink: Arc<dyn EventSink>,
122 approvals: Arc<dyn ApprovalGate>,
123 cancellation: CancellationToken,
124 ) -> Result<TurnOutcome, AgentError> {
125 if cancellation.is_cancelled() {
126 return Err(AgentError::Cancelled);
127 }
128 history.push(Message::User {
129 content: prompt.text,
130 images: prompt.images,
131 });
132 self.enforce_history_limits(history, sink.as_ref()).await?;
133 let specs = self.tools.specs();
134 let mut usage = Usage::default();
135
136 for step in 1..=self.config.max_steps {
137 if cancellation.is_cancelled() {
138 return Err(AgentError::Cancelled);
139 }
140 let selection = self
141 .context
142 .select(history, &self.config.system_prompt, &specs)
143 .map_err(|error| AgentError::ContextLimit(error.to_string()))?;
144 if selection.removed_messages > 0 {
145 sink.emit(CoreEvent::ContextCompacted {
146 before_tokens: selection.before_tokens,
147 after_tokens: selection.after_tokens,
148 removed_messages: selection.removed_messages,
149 })
150 .await?;
151 }
152 let delta_sink: Arc<dyn TextDeltaSink> = Arc::new(ForwardDeltas {
153 sink: Arc::clone(&sink),
154 });
155 let response = self
156 .provider
157 .complete(
158 ProviderRequest {
159 system_prompt: self.config.system_prompt.clone(),
160 messages: selection.messages,
161 tools: specs.clone(),
162 },
163 delta_sink,
164 cancellation.child_token(),
165 )
166 .await
167 .map_err(map_provider_error)?;
168 usage.add(&response.usage);
169 sink.emit(CoreEvent::AssistantCompleted {
170 content: response.content.clone(),
171 })
172 .await?;
173 let calls = response.tool_calls.clone();
174 history.push(Message::Assistant {
175 content: response.content,
176 tool_calls: response.tool_calls,
177 });
178 self.enforce_history_limits(history, sink.as_ref()).await?;
179 if calls.is_empty() {
180 return Ok(TurnOutcome { steps: step, usage });
181 }
182
183 for call in calls {
184 if cancellation.is_cancelled() {
185 return Err(AgentError::Cancelled);
186 }
187 sink.emit(CoreEvent::ToolProposed {
188 call_id: call.id.clone(),
189 name: call.name.clone(),
190 arguments: call.arguments.clone(),
191 })
192 .await?;
193 let Some(tool) = self.tools.get(&call.name) else {
194 let error = ToolError::new(
195 ToolFailure::UnknownTool,
196 format!("unknown tool: {}", call.name),
197 );
198 self.record_tool_error(history, sink.as_ref(), &call, error)
199 .await?;
200 self.enforce_history_limits(history, sink.as_ref()).await?;
201 continue;
202 };
203 let risk = match tool.risk(&call.arguments) {
204 Ok(risk) => risk,
205 Err(error) => {
206 self.record_tool_error(history, sink.as_ref(), &call, error)
207 .await?;
208 self.enforce_history_limits(history, sink.as_ref()).await?;
209 continue;
210 }
211 };
212 let summary = match tool.approval_summary(&call.arguments) {
213 Ok(summary) => summary,
214 Err(error) => {
215 self.record_tool_error(history, sink.as_ref(), &call, error)
216 .await?;
217 self.enforce_history_limits(history, sink.as_ref()).await?;
218 continue;
219 }
220 };
221 let approved = approvals
222 .approve(
223 ApprovalRequest {
224 call_id: call.id.clone(),
225 name: call.name.clone(),
226 risk,
227 cwd: self.workspace.clone(),
228 summary,
229 },
230 cancellation.child_token(),
231 )
232 .await?;
233 let output = if approved {
234 sink.emit(CoreEvent::ToolStarted {
235 call_id: call.id.clone(),
236 name: call.name.clone(),
237 })
238 .await?;
239 let progress = ProgressSink::buffered();
240 let execution = tool.execute(
241 call.arguments.clone(),
242 ToolContext {
243 call_id: call.id.clone(),
244 workspace: self.workspace.clone(),
245 cancellation: cancellation.child_token(),
246 background_job: false,
247 job: None,
248 progress: progress.clone(),
249 approvals: ToolApprovals::new(Arc::clone(&approvals), call.id.clone()),
250 },
251 );
252 forward_progress(execution, &progress, sink.as_ref(), &call.id)
253 .await
254 .unwrap_or_else(ToolOutput::from)
255 } else {
256 ToolOutput::failed(ToolFailure::Denied, DENIED)
257 };
258 sink.emit(CoreEvent::ToolCompleted {
259 call_id: call.id.clone(),
260 name: call.name.clone(),
261 output: output.clone(),
262 })
263 .await?;
264 let is_error = output.is_error();
265 history.push(Message::Tool {
266 call_id: call.id,
267 name: call.name,
268 content: output.content,
269 is_error,
270 });
271 self.enforce_history_limits(history, sink.as_ref()).await?;
272 }
273 }
274 Err(AgentError::StepLimit)
275 }
276
277 async fn record_tool_error(
279 &self,
280 history: &mut Vec<Message>,
281 sink: &dyn EventSink,
282 call: &ToolCall,
283 error: ToolError,
284 ) -> Result<(), AgentError> {
285 let output = ToolOutput::from(error);
286 sink.emit(CoreEvent::ToolCompleted {
287 call_id: call.id.clone(),
288 name: call.name.clone(),
289 output: output.clone(),
290 })
291 .await?;
292 history.push(Message::Tool {
293 call_id: call.id.clone(),
294 name: call.name.clone(),
295 content: output.content,
296 is_error: true,
297 });
298 Ok(())
299 }
300
301 pub(crate) async fn enforce_history_limits(
302 &self,
303 history: &mut Vec<Message>,
304 sink: &dyn EventSink,
305 ) -> Result<(), AgentError> {
306 crate::history::enforce_limits(history, &self.config.history_limits, sink).await
307 }
308}
309
310async fn forward_progress<T>(
314 execution: impl Future<Output = T>,
315 progress: &ProgressSink,
316 sink: &dyn EventSink,
317 call_id: &str,
318) -> T {
319 let mut execution = std::pin::pin!(execution);
320 let mut ticker = tokio::time::interval_at(
321 tokio::time::Instant::now() + PROGRESS_INTERVAL,
322 PROGRESS_INTERVAL,
323 );
324 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
325 let mut last_sent: Option<tokio::time::Instant> = None;
326 let mut forwarding = true;
327 let result = loop {
328 tokio::select! {
329 biased;
330 result = &mut execution => break result,
331 _ = ticker.tick(), if forwarding => {
332 if let Some(text) = progress.take() {
333 let event = CoreEvent::ToolProgress { call_id: call_id.to_owned(), text };
334 forwarding = sink.emit(event).await.is_ok();
335 last_sent = Some(tokio::time::Instant::now());
336 }
337 }
338 }
339 };
340 if forwarding
342 && last_sent.is_none_or(|sent| sent.elapsed() >= PROGRESS_INTERVAL)
343 && let Some(text) = progress.take()
344 {
345 let _ = sink
346 .emit(CoreEvent::ToolProgress {
347 call_id: call_id.to_owned(),
348 text,
349 })
350 .await;
351 }
352 result
353}
354
355fn map_provider_error(error: ProviderError) -> AgentError {
356 match error.kind {
357 ProviderErrorKind::Provider => AgentError::Provider(error.message),
358 ProviderErrorKind::ResponseLimit => AgentError::ResponseLimit(error.message),
359 ProviderErrorKind::ToolLimit => AgentError::ToolLimit(error.message),
360 ProviderErrorKind::Cancelled => AgentError::Cancelled,
361 }
362}
363
364struct ForwardDeltas {
365 sink: Arc<dyn EventSink>,
366}
367
368#[async_trait]
369impl TextDeltaSink for ForwardDeltas {
370 async fn push(&self, delta: &str) -> Result<(), ProviderError> {
371 self.sink
372 .emit(CoreEvent::AssistantDelta {
373 content: delta.to_owned(),
374 })
375 .await
376 .map_err(|error| match error {
377 AgentError::Cancelled => {
378 ProviderError::new(ProviderErrorKind::Cancelled, "turn cancelled")
379 }
380 AgentError::ResponseLimit(message) => {
381 ProviderError::new(ProviderErrorKind::ResponseLimit, message)
382 }
383 AgentError::ToolLimit(message) => {
384 ProviderError::new(ProviderErrorKind::ToolLimit, message)
385 }
386 error => ProviderError::new(ProviderErrorKind::Provider, error.to_string()),
387 })
388 }
389}
390
391#[cfg(test)]
392mod tests;