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 command: tool.approval_command(&call.arguments),
230 },
231 cancellation.child_token(),
232 )
233 .await?;
234 let output = if approved {
235 sink.emit(CoreEvent::ToolStarted {
236 call_id: call.id.clone(),
237 name: call.name.clone(),
238 })
239 .await?;
240 let progress = ProgressSink::buffered();
241 let execution = tool.execute(
242 call.arguments.clone(),
243 ToolContext {
244 call_id: call.id.clone(),
245 workspace: self.workspace.clone(),
246 cancellation: cancellation.child_token(),
247 background_job: false,
248 job: None,
249 progress: progress.clone(),
250 approvals: ToolApprovals::new(Arc::clone(&approvals), call.id.clone()),
251 },
252 );
253 forward_progress(execution, &progress, sink.as_ref(), &call.id)
254 .await
255 .unwrap_or_else(ToolOutput::from)
256 } else {
257 ToolOutput::failed(ToolFailure::Denied, DENIED)
258 };
259 sink.emit(CoreEvent::ToolCompleted {
260 call_id: call.id.clone(),
261 name: call.name.clone(),
262 output: output.clone(),
263 })
264 .await?;
265 let is_error = output.is_error();
266 history.push(Message::Tool {
267 call_id: call.id,
268 name: call.name,
269 content: output.content,
270 is_error,
271 });
272 self.enforce_history_limits(history, sink.as_ref()).await?;
273 }
274 }
275 Err(AgentError::StepLimit)
276 }
277
278 async fn record_tool_error(
280 &self,
281 history: &mut Vec<Message>,
282 sink: &dyn EventSink,
283 call: &ToolCall,
284 error: ToolError,
285 ) -> Result<(), AgentError> {
286 let output = ToolOutput::from(error);
287 sink.emit(CoreEvent::ToolCompleted {
288 call_id: call.id.clone(),
289 name: call.name.clone(),
290 output: output.clone(),
291 })
292 .await?;
293 history.push(Message::Tool {
294 call_id: call.id.clone(),
295 name: call.name.clone(),
296 content: output.content,
297 is_error: true,
298 });
299 Ok(())
300 }
301
302 pub(crate) async fn enforce_history_limits(
303 &self,
304 history: &mut Vec<Message>,
305 sink: &dyn EventSink,
306 ) -> Result<(), AgentError> {
307 crate::history::enforce_limits(history, &self.config.history_limits, sink).await
308 }
309}
310
311async fn forward_progress<T>(
315 execution: impl Future<Output = T>,
316 progress: &ProgressSink,
317 sink: &dyn EventSink,
318 call_id: &str,
319) -> T {
320 let mut execution = std::pin::pin!(execution);
321 let mut ticker = tokio::time::interval_at(
322 tokio::time::Instant::now() + PROGRESS_INTERVAL,
323 PROGRESS_INTERVAL,
324 );
325 ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
326 let mut last_sent: Option<tokio::time::Instant> = None;
327 let mut forwarding = true;
328 let result = loop {
329 tokio::select! {
330 biased;
331 result = &mut execution => break result,
332 _ = ticker.tick(), if forwarding => {
333 if let Some(text) = progress.take() {
334 let event = CoreEvent::ToolProgress { call_id: call_id.to_owned(), text };
335 forwarding = sink.emit(event).await.is_ok();
336 last_sent = Some(tokio::time::Instant::now());
337 }
338 }
339 }
340 };
341 if forwarding
343 && last_sent.is_none_or(|sent| sent.elapsed() >= PROGRESS_INTERVAL)
344 && let Some(text) = progress.take()
345 {
346 let _ = sink
347 .emit(CoreEvent::ToolProgress {
348 call_id: call_id.to_owned(),
349 text,
350 })
351 .await;
352 }
353 result
354}
355
356fn map_provider_error(error: ProviderError) -> AgentError {
357 match error.kind {
358 ProviderErrorKind::Provider => AgentError::Provider(error.message),
359 ProviderErrorKind::ResponseLimit => AgentError::ResponseLimit(error.message),
360 ProviderErrorKind::ToolLimit => AgentError::ToolLimit(error.message),
361 ProviderErrorKind::Cancelled => AgentError::Cancelled,
362 }
363}
364
365struct ForwardDeltas {
366 sink: Arc<dyn EventSink>,
367}
368
369#[async_trait]
370impl TextDeltaSink for ForwardDeltas {
371 async fn push(&self, delta: &str) -> Result<(), ProviderError> {
372 self.sink
373 .emit(CoreEvent::AssistantDelta {
374 content: delta.to_owned(),
375 })
376 .await
377 .map_err(|error| match error {
378 AgentError::Cancelled => {
379 ProviderError::new(ProviderErrorKind::Cancelled, "turn cancelled")
380 }
381 AgentError::ResponseLimit(message) => {
382 ProviderError::new(ProviderErrorKind::ResponseLimit, message)
383 }
384 AgentError::ToolLimit(message) => {
385 ProviderError::new(ProviderErrorKind::ToolLimit, message)
386 }
387 error => ProviderError::new(ProviderErrorKind::Provider, error.to_string()),
388 })
389 }
390}
391
392#[cfg(test)]
393mod tests;