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