1use std::collections::BTreeSet;
2use std::future::poll_fn;
3use std::sync::Arc;
4use std::task::Poll;
5
6use serde::{Deserialize, Serialize};
7use thiserror::Error;
8
9use crate::{
10 BoxFuture, CachePolicy, CancellationSignal, ChatMessage, ContextProjector,
11 DEFAULT_HISTORY_CONTEXT_TOKENS, DeterministicContextProjector, Gateway, GatewayError,
12 GatewayEvent, GatewayEventSink, GatewayRequest, LARGE_TOOL_RESULT_BYTES, NeverCancelled,
13 PermissionDecision, PermissionEngine, PermissionRequest, Role, TOOL_RESULT_PREVIEW_BYTES,
14 ToolArgumentIntegrity, ToolError, ToolExecutionProvenance, ToolOutput, ToolPreparation,
15 ToolRegistry, ToolReview, Usage,
16};
17
18#[derive(Clone, Debug)]
19pub struct AgentOptions {
20 pub model: String,
21 pub max_steps: usize,
22 pub max_output_tokens: Option<u32>,
23 pub history_context_tokens: usize,
24}
25
26impl AgentOptions {
27 pub fn new(model: impl Into<String>) -> Self {
28 Self {
29 model: model.into(),
30 max_steps: 50,
31 max_output_tokens: None,
32 history_context_tokens: DEFAULT_HISTORY_CONTEXT_TOKENS,
33 }
34 }
35}
36
37#[derive(Clone, Debug, Default)]
38pub struct AgentRequest {
39 pub history: Vec<ChatMessage>,
40 pub prompt: String,
41}
42
43#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
44#[serde(rename_all = "snake_case")]
45pub enum AgentStopReason {
46 Complete,
47 StepLimit,
48 Cancelled,
49}
50
51#[derive(Clone, Debug)]
52pub struct AgentResult {
53 pub messages: Vec<ChatMessage>,
54 pub output: String,
55 pub usage: Usage,
56 pub steps: usize,
57 pub stop_reason: AgentStopReason,
58 pub delivery_ambiguous: bool,
59}
60
61#[derive(Clone, Copy, Debug, Eq, PartialEq)]
62pub enum ApprovalDecision {
63 AllowOnce,
64 AllowForSession,
65 Deny,
66}
67
68#[derive(Clone, Copy, Debug, Eq, PartialEq)]
69pub enum ApprovalKind {
70 User,
71 Automatic,
72}
73
74#[derive(Clone, Debug)]
75pub struct ApprovalRequest {
76 pub kind: ApprovalKind,
77 pub tool_call_id: String,
78 pub tool_name: String,
79 pub arguments_json: String,
80 pub permission_requests: Vec<PermissionRequest>,
81 pub irreversible: bool,
82 pub review: Option<ToolReview>,
83}
84
85#[derive(Debug, Error)]
86pub enum ApprovalError {
87 #[error("approval is unavailable: {0}")]
88 Unavailable(String),
89}
90
91pub trait ApprovalHandler: Send {
92 fn review<'a>(
93 &'a mut self,
94 request: ApprovalRequest,
95 ) -> BoxFuture<'a, Result<ApprovalDecision, ApprovalError>>;
96}
97
98#[derive(Clone, Copy, Debug)]
99pub struct StaticApprovalHandler {
100 decision: ApprovalDecision,
101}
102
103impl StaticApprovalHandler {
104 pub fn allow_once() -> Self {
105 Self {
106 decision: ApprovalDecision::AllowOnce,
107 }
108 }
109
110 pub fn deny() -> Self {
111 Self {
112 decision: ApprovalDecision::Deny,
113 }
114 }
115}
116
117impl ApprovalHandler for StaticApprovalHandler {
118 fn review<'a>(
119 &'a mut self,
120 _request: ApprovalRequest,
121 ) -> BoxFuture<'a, Result<ApprovalDecision, ApprovalError>> {
122 Box::pin(async move { Ok(self.decision) })
123 }
124}
125
126#[derive(Clone, Debug, PartialEq)]
127pub enum AgentEvent {
128 Gateway(GatewayEvent),
129 ToolStarted {
130 id: String,
131 name: String,
132 arguments_json: String,
133 },
134 ToolFinished {
135 id: String,
136 name: String,
137 is_error: bool,
138 output: ToolOutput,
139 },
140}
141
142pub trait AgentEventSink: Send {
143 fn emit(&mut self, event: AgentEvent);
144}
145
146#[derive(Debug, Error)]
147pub enum AgentError {
148 #[error(transparent)]
149 Gateway(#[from] GatewayError),
150 #[error(transparent)]
151 Approval(#[from] ApprovalError),
152 #[error(transparent)]
153 ProjectContext(#[from] crate::ScopedProjectContextError),
154}
155
156pub struct Agent {
162 gateway: Arc<dyn Gateway>,
163 tools: Arc<ToolRegistry>,
164 options: AgentOptions,
165 context_projector: Arc<dyn ContextProjector>,
166}
167
168impl Agent {
169 pub fn new(gateway: Arc<dyn Gateway>, tools: Arc<ToolRegistry>, options: AgentOptions) -> Self {
170 Self {
171 gateway,
172 tools,
173 options,
174 context_projector: Arc::new(DeterministicContextProjector),
175 }
176 }
177
178 pub fn with_context_projector(mut self, projector: Arc<dyn ContextProjector>) -> Self {
179 self.context_projector = projector;
180 self
181 }
182
183 pub fn run<'a>(
184 &'a self,
185 request: AgentRequest,
186 tool_context: &'a crate::ToolContext,
187 permissions: &'a mut PermissionEngine,
188 approvals: &'a mut dyn ApprovalHandler,
189 events: &'a mut dyn AgentEventSink,
190 ) -> BoxFuture<'a, Result<AgentResult, AgentError>> {
191 self.run_controlled(
192 request,
193 tool_context,
194 permissions,
195 approvals,
196 events,
197 Arc::new(NeverCancelled),
198 )
199 }
200
201 pub fn run_controlled<'a>(
202 &'a self,
203 request: AgentRequest,
204 tool_context: &'a crate::ToolContext,
205 permissions: &'a mut PermissionEngine,
206 approvals: &'a mut dyn ApprovalHandler,
207 events: &'a mut dyn AgentEventSink,
208 cancellation: Arc<dyn CancellationSignal>,
209 ) -> BoxFuture<'a, Result<AgentResult, AgentError>> {
210 Box::pin(async move {
211 let mut messages = request.history;
212 messages.push(ChatMessage::text(Role::User, request.prompt));
213 let mut usage = Usage::default();
214 let mut output = String::new();
215 let mut delivery_ambiguous = false;
216 let mut invocation_context = tool_context.clone();
217 invocation_context.cancellation = cancellation.clone();
218
219 let mut step = 0usize;
220 while self.options.max_steps == 0 || step < self.options.max_steps {
221 if cancellation.is_cancelled() {
222 return Ok(cancelled_result(
223 messages,
224 output,
225 usage,
226 step,
227 delivery_ambiguous,
228 ));
229 }
230 step += 1;
231 let projection = self
232 .context_projector
233 .project(&messages, self.options.history_context_tokens);
234 let gateway_request = GatewayRequest {
235 model: self.options.model.clone(),
236 messages: projection.messages,
237 tools: self.tools.advertisements(),
238 tool_choice: crate::ToolChoice::Auto,
239 max_output_tokens: self.options.max_output_tokens,
240 };
241 let response = {
242 let mut bridge = GatewayEventBridge(events);
243 match self.gateway.complete(gateway_request, &mut bridge).await {
244 Err(GatewayError::Cancelled) if cancellation.is_cancelled() => {
245 return Ok(cancelled_result(
246 messages,
247 output,
248 usage,
249 step,
250 delivery_ambiguous,
251 ));
252 }
253 result => result?,
254 }
255 };
256 accumulate_usage(&mut usage, response.usage);
257 delivery_ambiguous |= response.delivery_ambiguous;
258 if let Some(content) = &response.content {
259 output.push_str(content);
260 }
261
262 messages.push(ChatMessage {
263 role: Role::Assistant,
264 content: response.content.clone(),
265 tool_call_id: None,
266 tool_name: None,
267 tool_calls: response.tool_calls.clone(),
268 permission_feedback: false,
269 cache_policy: CachePolicy::Default,
270 });
271
272 if cancellation.is_cancelled() {
273 return Ok(cancelled_result(
274 messages,
275 output,
276 usage,
277 step,
278 delivery_ambiguous,
279 ));
280 }
281
282 if response.tool_calls.is_empty() {
283 return Ok(AgentResult {
284 messages,
285 output,
286 usage,
287 steps: step,
288 stop_reason: AgentStopReason::Complete,
289 delivery_ambiguous,
290 });
291 }
292
293 let calls = response.tool_calls;
294 let (context_delta, context_deferred) =
295 select_scoped_project_context(&self.tools, &invocation_context, &calls)?;
296 if let Some(delta) = context_delta {
297 insert_stable_system_context(&mut messages, delta);
298 }
299 let outcomes = execute_tool_batch(
300 &self.tools,
301 &invocation_context,
302 permissions,
303 approvals,
304 events,
305 &calls,
306 &context_deferred,
307 )
308 .await?;
309 for (call, (tool_output, permission_feedback)) in calls.into_iter().zip(outcomes) {
310 messages.push(ChatMessage {
311 role: Role::Tool,
312 content: Some(tool_output.content),
313 tool_call_id: Some(call.id),
314 tool_name: Some(call.name),
315 tool_calls: Vec::new(),
316 permission_feedback,
317 cache_policy: CachePolicy::Default,
318 });
319 }
320 }
321
322 Ok(AgentResult {
323 messages,
324 output,
325 usage,
326 steps: step,
327 stop_reason: AgentStopReason::StepLimit,
328 delivery_ambiguous,
329 })
330 })
331 }
332}
333
334fn cancelled_result(
335 messages: Vec<ChatMessage>,
336 output: String,
337 usage: Usage,
338 steps: usize,
339 delivery_ambiguous: bool,
340) -> AgentResult {
341 AgentResult {
342 messages,
343 output,
344 usage,
345 steps,
346 stop_reason: AgentStopReason::Cancelled,
347 delivery_ambiguous,
348 }
349}
350
351async fn execute_tool_batch(
352 tools: &ToolRegistry,
353 context: &crate::ToolContext,
354 permissions: &mut PermissionEngine,
355 approvals: &mut dyn ApprovalHandler,
356 events: &mut dyn AgentEventSink,
357 calls: &[crate::ToolCall],
358 context_deferred: &[bool],
359) -> Result<Vec<(ToolOutput, bool)>, AgentError> {
360 let mut outcomes = Vec::with_capacity(calls.len());
361 let mut index = 0usize;
362 while index < calls.len() {
363 if context.cancellation.is_cancelled() {
364 outcomes.extend(
365 calls[index..]
366 .iter()
367 .map(|_| (error_output(ToolError::Cancelled.to_string()), false)),
368 );
369 break;
370 }
371 if context_deferred.get(index).copied().unwrap_or(false) {
372 let call = &calls[index];
373 emit_tool_started(events, call);
374 let output = error_output(
375 "tool execution deferred because new scoped project instructions were loaded; review the new rules and retry the action",
376 );
377 emit_tool_finished(events, call, &output);
378 outcomes.push((output, false));
379 index += 1;
380 continue;
381 }
382 let parallel_len = parallel_read_prefix_len(tools, &calls[index..]);
383 if parallel_len >= 2 {
384 let batch = &calls[index..index + parallel_len];
385 for call in batch {
386 emit_tool_started(events, call);
387 }
388 let mut executions = Vec::with_capacity(batch.len());
389 for call in batch {
390 let admission =
391 admit_local_tool(tools, context, permissions, approvals, call).await?;
392 executions.push(admission.into_future(context.clone()));
393 }
394 let mut completed = join_ordered(executions).await;
395 for (call, outcome) in batch.iter().zip(&mut completed) {
396 finalize_tool_output(context, call, &mut outcome.0);
397 emit_tool_finished(events, call, &outcome.0);
398 }
399 outcomes.extend(completed);
400 index += parallel_len;
401 continue;
402 }
403
404 let call = &calls[index];
405 emit_tool_started(events, call);
406 let mut outcome = execute_one_tool(tools, context, permissions, approvals, call).await?;
407 finalize_tool_output(context, call, &mut outcome.0);
408 emit_tool_finished(events, call, &outcome.0);
409 outcomes.push(outcome);
410 index += 1;
411 }
412 Ok(outcomes)
413}
414
415fn select_scoped_project_context(
416 tools: &ToolRegistry,
417 context: &crate::ToolContext,
418 calls: &[crate::ToolCall],
419) -> Result<(Option<String>, Vec<bool>), AgentError> {
420 let mut deferred = vec![false; calls.len()];
421 let Some(provider) = &context.project_context else {
422 return Ok((None, deferred));
423 };
424 let mut targets = BTreeSet::new();
425 let mut sensitive_with_targets = Vec::new();
426 for (index, call) in calls.iter().enumerate() {
427 if call.provenance != ToolExecutionProvenance::FxLocal
428 || call.argument_integrity == ToolArgumentIntegrity::MalformedJson
429 {
430 continue;
431 }
432 let Ok((tool, arguments)) = tools.validate_call(call) else {
433 continue;
434 };
435 let Ok(call_targets) = tool.project_context_targets(context, &arguments) else {
436 continue;
437 };
438 if call_targets.is_empty() {
439 continue;
440 }
441 let sensitive = !matches!(tool.effect(&arguments), Ok(crate::ToolEffect::Read));
442 targets.extend(call_targets);
443 if sensitive {
444 sensitive_with_targets.push(index);
445 }
446 }
447 if targets.is_empty() {
448 return Ok((None, deferred));
449 }
450 let targets = targets.into_iter().collect::<Vec<_>>();
451 let delta = provider.select(&targets)?;
452 if delta.is_some() {
453 for index in sensitive_with_targets {
454 deferred[index] = true;
455 }
456 }
457 Ok((delta, deferred))
458}
459
460fn insert_stable_system_context(messages: &mut Vec<ChatMessage>, content: String) {
461 let prefix_end = messages
462 .iter()
463 .take_while(|message| message.role == Role::System)
464 .count();
465 messages.insert(prefix_end, ChatMessage::text(Role::System, content));
466}
467
468fn parallel_read_prefix_len(tools: &ToolRegistry, calls: &[crate::ToolCall]) -> usize {
469 calls
470 .iter()
471 .take_while(|call| {
472 if call.provenance != ToolExecutionProvenance::FxLocal
473 || call.argument_integrity == ToolArgumentIntegrity::MalformedJson
474 {
475 return false;
476 }
477 let Ok((tool, arguments)) = tools.validate_call(call) else {
478 return false;
479 };
480 matches!(tool.effect(&arguments), Ok(crate::ToolEffect::Read))
481 })
482 .count()
483}
484
485fn emit_tool_started(events: &mut dyn AgentEventSink, call: &crate::ToolCall) {
486 events.emit(AgentEvent::ToolStarted {
487 id: call.id.clone(),
488 name: call.name.clone(),
489 arguments_json: call.arguments_json.clone(),
490 });
491}
492
493fn emit_tool_finished(
494 events: &mut dyn AgentEventSink,
495 call: &crate::ToolCall,
496 output: &ToolOutput,
497) {
498 events.emit(AgentEvent::ToolFinished {
499 id: call.id.clone(),
500 name: call.name.clone(),
501 is_error: output.is_error,
502 output: output.clone(),
503 });
504}
505
506async fn execute_one_tool(
507 tools: &ToolRegistry,
508 context: &crate::ToolContext,
509 permissions: &mut PermissionEngine,
510 approvals: &mut dyn ApprovalHandler,
511 call: &crate::ToolCall,
512) -> Result<(ToolOutput, bool), AgentError> {
513 if call.provenance == ToolExecutionProvenance::Provider {
514 let content = call
515 .provider_result
516 .clone()
517 .unwrap_or_else(|| "provider-executed tool returned no result".into());
518 return Ok((
519 ToolOutput {
520 original_bytes: content.len(),
521 content,
522 is_error: false,
523 structured: None,
524 truncated: false,
525 durable_content: None,
526 },
527 false,
528 ));
529 }
530 if call.argument_integrity == ToolArgumentIntegrity::MalformedJson {
531 return Ok((error_output("tool arguments were malformed JSON"), false));
532 }
533 let admission = admit_local_tool(tools, context, permissions, approvals, call).await?;
534 Ok(admission.into_future(context.clone()).await)
535}
536
537enum LocalToolAdmission {
538 Immediate(ToolOutput, bool),
539 Direct {
540 tool: Arc<dyn crate::Tool>,
541 arguments: serde_json::Value,
542 },
543 Prepared(crate::PreparedToolCall),
544}
545
546impl LocalToolAdmission {
547 fn into_future(self, context: crate::ToolContext) -> BoxFuture<'static, (ToolOutput, bool)> {
548 Box::pin(async move {
549 match self {
550 Self::Immediate(output, permission_feedback) => (output, permission_feedback),
551 Self::Direct { tool, arguments } => {
552 if context.cancellation.is_cancelled() {
553 return (error_output(ToolError::Cancelled.to_string()), false);
554 }
555 match tool.execute(&context, arguments).await {
556 Ok(output) => (output, false),
557 Err(error) => (error_output(error.to_string()), false),
558 }
559 }
560 Self::Prepared(prepared) => {
561 if context.cancellation.is_cancelled() {
562 return (error_output(ToolError::Cancelled.to_string()), false);
563 }
564 match prepared.commit(&context).await {
565 Ok(output) => (output, false),
566 Err(error) => (error_output(error.to_string()), false),
567 }
568 }
569 }
570 })
571 }
572}
573
574async fn admit_local_tool(
575 tools: &ToolRegistry,
576 context: &crate::ToolContext,
577 permissions: &mut PermissionEngine,
578 approvals: &mut dyn ApprovalHandler,
579 call: &crate::ToolCall,
580) -> Result<LocalToolAdmission, AgentError> {
581 let (tool, arguments) = match tools.validate_call(call) {
582 Ok(validated) => validated,
583 Err(error) => {
584 return Ok(LocalToolAdmission::Immediate(
585 error_output(error.to_string()),
586 false,
587 ));
588 }
589 };
590 let preparation = match tool.prepare(context, &arguments) {
591 Ok(preparation) => preparation,
592 Err(error) => {
593 return Ok(LocalToolAdmission::Immediate(
594 error_output(error.to_string()),
595 false,
596 ));
597 }
598 };
599 let (requests, irreversible, review) = match &preparation {
600 ToolPreparation::Direct {
601 permission_requests,
602 irreversible,
603 } => (permission_requests, *irreversible, None),
604 ToolPreparation::Prepared(prepared) => (
605 &prepared.permission_requests,
606 prepared.irreversible,
607 prepared.review.clone(),
608 ),
609 };
610
611 let decisions: Vec<_> = requests
612 .iter()
613 .map(|request| permissions.decide(request))
614 .collect();
615 if decisions.contains(&PermissionDecision::Deny) {
616 return Ok(LocalToolAdmission::Immediate(
617 error_output("tool permission denied"),
618 true,
619 ));
620 }
621 if decisions.contains(&PermissionDecision::Ask)
622 || decisions.contains(&PermissionDecision::AutoReview)
623 {
624 let kind = if decisions.contains(&PermissionDecision::Ask) {
625 ApprovalKind::User
626 } else {
627 ApprovalKind::Automatic
628 };
629 let review_request = ApprovalRequest {
630 kind,
631 tool_call_id: call.id.clone(),
632 tool_name: tool.name().to_owned(),
633 arguments_json: call.arguments_json.clone(),
634 permission_requests: requests.clone(),
635 irreversible,
636 review,
637 };
638 match approvals.review(review_request).await? {
639 ApprovalDecision::AllowOnce => {}
640 ApprovalDecision::AllowForSession => {
641 for (request, decision) in requests.iter().zip(decisions) {
642 if matches!(
643 decision,
644 PermissionDecision::Ask | PermissionDecision::AutoReview
645 ) {
646 permissions.grant_request_for_session(request);
647 }
648 }
649 }
650 ApprovalDecision::Deny => {
651 let reason = match kind {
652 ApprovalKind::User => "tool permission denied by user",
653 ApprovalKind::Automatic => "tool permission denied by automatic safety review",
654 };
655 return Ok(LocalToolAdmission::Immediate(error_output(reason), true));
656 }
657 }
658 }
659
660 if context.cancellation.is_cancelled() {
661 return Ok(LocalToolAdmission::Immediate(
662 error_output(ToolError::Cancelled.to_string()),
663 false,
664 ));
665 }
666
667 Ok(match preparation {
668 ToolPreparation::Direct { .. } => LocalToolAdmission::Direct { tool, arguments },
669 ToolPreparation::Prepared(prepared) => LocalToolAdmission::Prepared(prepared),
670 })
671}
672
673async fn join_ordered<T>(futures: Vec<BoxFuture<'static, T>>) -> Vec<T> {
674 let mut futures = futures.into_iter().map(Some).collect::<Vec<_>>();
675 let mut outputs = (0..futures.len()).map(|_| None).collect::<Vec<_>>();
676 poll_fn(move |task_context| {
677 let mut remaining = 0usize;
678 for (index, future) in futures.iter_mut().enumerate() {
679 let Some(pending) = future.as_mut() else {
680 continue;
681 };
682 match pending.as_mut().poll(task_context) {
683 Poll::Ready(output) => {
684 outputs[index] = Some(output);
685 *future = None;
686 }
687 Poll::Pending => remaining += 1,
688 }
689 }
690 if remaining == 0 {
691 Poll::Ready(
692 outputs
693 .iter_mut()
694 .map(|output| output.take().expect("completed future has output"))
695 .collect(),
696 )
697 } else {
698 Poll::Pending
699 }
700 })
701 .await
702}
703
704fn error_output(content: impl Into<String>) -> ToolOutput {
705 let content = content.into();
706 let original_bytes = content.len();
707 ToolOutput {
708 content,
709 is_error: true,
710 structured: None,
711 original_bytes,
712 truncated: false,
713 durable_content: None,
714 }
715}
716
717fn finalize_tool_output(
718 context: &crate::ToolContext,
719 call: &crate::ToolCall,
720 output: &mut ToolOutput,
721) {
722 let Some(durable) = output.durable_content.take() else {
723 if let std::borrow::Cow::Owned(redacted) = crate::redact_secrets(&output.content) {
724 output.content = redacted;
725 output.structured = None;
726 }
727 return;
728 };
729 output.original_bytes = output.original_bytes.max(durable.len());
730 let durable = crate::redact_secrets(&durable).into_owned();
731 if durable.len() > LARGE_TOOL_RESULT_BYTES
732 && let Some(store) = &context.tool_results
733 {
734 match store.store(&call.id, &call.name, &durable) {
735 Ok(stored) => {
736 let preview = utf8_prefix(&durable, TOOL_RESULT_PREVIEW_BYTES);
737 output.content = format!(
738 "<tool_result_preview handle=\"{}\" stored_bytes=\"{}\">\n{}\n</tool_result_preview>\n<tool_result_handle>{}</tool_result_handle>\nFull result is stored outside session JSON. Use read_tool_result with this handle to inspect a byte range or literal query.",
739 stored.handle, stored.stored_bytes, preview, stored.handle
740 );
741 output.structured = None;
742 output.truncated = true;
743 return;
744 }
745 Err(error) => {
746 *output = error_output(format!("tool result storage failed: {error}"));
747 return;
748 }
749 }
750 }
751
752 let (content, truncated) = bounded_inline_output(
753 &durable,
754 context.limits.max_result_bytes,
755 call.name.as_str(),
756 );
757 output.content = content;
758 output.truncated |= truncated;
759 if truncated {
760 output.structured = None;
761 }
762}
763
764fn bounded_inline_output(content: &str, max_bytes: usize, tool_name: &str) -> (String, bool) {
765 if content.len() <= max_bytes {
766 return (content.to_owned(), false);
767 }
768 let marker = format!(
769 "\n... [tool result truncated for {tool_name}: original {} bytes; cap is {max_bytes} bytes]\n",
770 content.len()
771 );
772 if marker.len() >= max_bytes {
773 return (utf8_prefix(&marker, max_bytes).to_owned(), true);
774 }
775 let prefix = utf8_prefix(content, max_bytes - marker.len());
776 (format!("{prefix}{marker}"), true)
777}
778
779fn utf8_prefix(content: &str, max_bytes: usize) -> &str {
780 let mut end = content.len().min(max_bytes);
781 while end > 0 && !content.is_char_boundary(end) {
782 end -= 1;
783 }
784 &content[..end]
785}
786
787fn accumulate_usage(total: &mut Usage, next: Usage) {
788 total.input_tokens = sum_optional(total.input_tokens, next.input_tokens);
789 total.output_tokens = sum_optional(total.output_tokens, next.output_tokens);
790}
791
792fn sum_optional(left: Option<u64>, right: Option<u64>) -> Option<u64> {
793 match (left, right) {
794 (None, None) => None,
795 (left, right) => Some(left.unwrap_or(0).saturating_add(right.unwrap_or(0))),
796 }
797}
798
799struct GatewayEventBridge<'a>(&'a mut dyn AgentEventSink);
800
801impl GatewayEventSink for GatewayEventBridge<'_> {
802 fn emit(&mut self, event: GatewayEvent) {
803 self.0.emit(AgentEvent::Gateway(event));
804 }
805}
806
807#[cfg(test)]
808mod tests {
809 use std::collections::VecDeque;
810 use std::sync::Mutex;
811 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
812 use std::task::{Context, Poll, Waker};
813
814 use serde_json::{Value, json};
815
816 use super::*;
817 use crate::{
818 FinishReason, GatewayResponse, PermissionMode, ScopedProjectContextError,
819 ScopedProjectContextProvider, StoredToolResult, Tool, ToolCall, ToolContext, ToolEffect,
820 ToolError, ToolResultMatch, ToolResultPage, ToolResultStore, ToolResultStoreError,
821 };
822
823 struct ScriptedGateway {
824 responses: Mutex<VecDeque<GatewayResponse>>,
825 requests: Mutex<Vec<GatewayRequest>>,
826 }
827
828 struct CancellingGateway {
829 cancellation: Arc<AtomicBool>,
830 executions: Arc<AtomicUsize>,
831 }
832
833 struct CancelledGateway {
834 cancellation: Arc<AtomicBool>,
835 }
836
837 #[derive(Default)]
838 struct MemoryToolResultStore {
839 stored: Mutex<Vec<String>>,
840 }
841
842 struct OneScopedDelta(AtomicUsize);
843
844 impl ScopedProjectContextProvider for OneScopedDelta {
845 fn select(
846 &self,
847 targets: &[std::path::PathBuf],
848 ) -> Result<Option<String>, ScopedProjectContextError> {
849 assert_eq!(
850 targets,
851 [std::path::PathBuf::from("/workspace/nested/file.rs")]
852 );
853 Ok((self.0.fetch_add(1, Ordering::SeqCst) == 0).then(|| "NESTED PROJECT RULE".into()))
854 }
855
856 fn fork_session(&self) -> Arc<dyn ScopedProjectContextProvider> {
857 Arc::new(Self(AtomicUsize::new(0)))
858 }
859 }
860
861 struct ScopedWriteTool(Arc<AtomicUsize>);
862
863 impl Tool for ScopedWriteTool {
864 fn name(&self) -> &str {
865 "scoped_write"
866 }
867
868 fn description(&self) -> &str {
869 "Test scoped write."
870 }
871
872 fn input_schema(&self) -> Value {
873 json!({"type": "object"})
874 }
875
876 fn effect(&self, _: &Value) -> Result<ToolEffect, ToolError> {
877 Ok(ToolEffect::Write)
878 }
879
880 fn project_context_targets(
881 &self,
882 _context: &ToolContext,
883 _arguments: &Value,
884 ) -> Result<Vec<std::path::PathBuf>, ToolError> {
885 Ok(vec!["/workspace/nested/file.rs".into()])
886 }
887
888 fn execute<'a>(
889 &'a self,
890 _context: &'a ToolContext,
891 _arguments: Value,
892 ) -> BoxFuture<'a, Result<ToolOutput, ToolError>> {
893 Box::pin(async move {
894 self.0.fetch_add(1, Ordering::SeqCst);
895 Ok(ToolOutput {
896 content: "written".into(),
897 is_error: false,
898 structured: None,
899 original_bytes: 7,
900 truncated: false,
901 durable_content: None,
902 })
903 })
904 }
905 }
906
907 impl ToolResultStore for MemoryToolResultStore {
908 fn store(
909 &self,
910 _tool_call_id: &str,
911 _tool_name: &str,
912 content: &str,
913 ) -> Result<StoredToolResult, ToolResultStoreError> {
914 self.stored.lock().unwrap().push(content.to_owned());
915 Ok(StoredToolResult {
916 handle: "result-test.txt".into(),
917 stored_bytes: content.len(),
918 })
919 }
920
921 fn read_range(
922 &self,
923 _handle: &str,
924 _start_byte: usize,
925 _byte_count: usize,
926 ) -> Result<ToolResultPage, ToolResultStoreError> {
927 Err(ToolResultStoreError::NotFound)
928 }
929
930 fn search(
931 &self,
932 _handle: &str,
933 _query: &str,
934 _max_matches: usize,
935 ) -> Result<Vec<ToolResultMatch>, ToolResultStoreError> {
936 Err(ToolResultStoreError::NotFound)
937 }
938 }
939
940 impl Gateway for CancelledGateway {
941 fn complete<'a>(
942 &'a self,
943 _request: GatewayRequest,
944 _events: &'a mut dyn GatewayEventSink,
945 ) -> BoxFuture<'a, Result<GatewayResponse, GatewayError>> {
946 Box::pin(async move {
947 self.cancellation.store(true, Ordering::Release);
948 Err(GatewayError::Cancelled)
949 })
950 }
951 }
952
953 impl Gateway for CancellingGateway {
954 fn complete<'a>(
955 &'a self,
956 _request: GatewayRequest,
957 _events: &'a mut dyn GatewayEventSink,
958 ) -> BoxFuture<'a, Result<GatewayResponse, GatewayError>> {
959 Box::pin(async move {
960 self.executions.fetch_add(1, Ordering::SeqCst);
961 self.cancellation.store(true, Ordering::Release);
962 Ok(response(
963 Some("partial"),
964 vec![tool_call(ToolExecutionProvenance::FxLocal)],
965 ))
966 })
967 }
968 }
969
970 struct AtomicCancellation(Arc<AtomicBool>);
971
972 impl CancellationSignal for AtomicCancellation {
973 fn is_cancelled(&self) -> bool {
974 self.0.load(Ordering::Acquire)
975 }
976 }
977
978 impl Gateway for ScriptedGateway {
979 fn complete<'a>(
980 &'a self,
981 request: GatewayRequest,
982 events: &'a mut dyn GatewayEventSink,
983 ) -> BoxFuture<'a, Result<GatewayResponse, GatewayError>> {
984 Box::pin(async move {
985 events.emit(GatewayEvent::ContentDelta("stream".into()));
986 self.requests.lock().unwrap().push(request);
987 self.responses
988 .lock()
989 .unwrap()
990 .pop_front()
991 .ok_or_else(|| GatewayError::InvalidResponse("script exhausted".into()))
992 })
993 }
994 }
995
996 struct CountingTool(Arc<AtomicUsize>);
997
998 impl Tool for CountingTool {
999 fn name(&self) -> &str {
1000 "count"
1001 }
1002
1003 fn description(&self) -> &str {
1004 "Count local executions."
1005 }
1006
1007 fn input_schema(&self) -> Value {
1008 json!({"type": "object"})
1009 }
1010
1011 fn effect(&self, _: &Value) -> Result<ToolEffect, ToolError> {
1012 Ok(ToolEffect::Write)
1013 }
1014
1015 fn execute<'a>(
1016 &'a self,
1017 _context: &'a ToolContext,
1018 _arguments: Value,
1019 ) -> BoxFuture<'a, Result<ToolOutput, ToolError>> {
1020 Box::pin(async move {
1021 self.0.fetch_add(1, Ordering::SeqCst);
1022 Ok(ToolOutput {
1023 content: "counted".into(),
1024 is_error: false,
1025 structured: None,
1026 original_bytes: 7,
1027 truncated: false,
1028 durable_content: None,
1029 })
1030 })
1031 }
1032 }
1033
1034 struct CooperativeReadTool {
1035 name: &'static str,
1036 started: Arc<AtomicUsize>,
1037 }
1038
1039 impl Tool for CooperativeReadTool {
1040 fn name(&self) -> &str {
1041 self.name
1042 }
1043
1044 fn description(&self) -> &str {
1045 "Read concurrently."
1046 }
1047
1048 fn input_schema(&self) -> Value {
1049 json!({"type": "object"})
1050 }
1051
1052 fn effect(&self, _: &Value) -> Result<ToolEffect, ToolError> {
1053 Ok(ToolEffect::Read)
1054 }
1055
1056 fn execute<'a>(
1057 &'a self,
1058 _context: &'a ToolContext,
1059 _arguments: Value,
1060 ) -> BoxFuture<'a, Result<ToolOutput, ToolError>> {
1061 let mut announced = false;
1062 Box::pin(std::future::poll_fn(move |task_context| {
1063 if !announced {
1064 self.started.fetch_add(1, Ordering::SeqCst);
1065 announced = true;
1066 }
1067 if self.started.load(Ordering::SeqCst) < 2 {
1068 task_context.waker().wake_by_ref();
1069 return Poll::Pending;
1070 }
1071 Poll::Ready(Ok(ToolOutput {
1072 content: self.name.into(),
1073 is_error: false,
1074 structured: None,
1075 original_bytes: self.name.len(),
1076 truncated: false,
1077 durable_content: None,
1078 }))
1079 }))
1080 }
1081 }
1082
1083 #[derive(Default)]
1084 struct Events(Vec<AgentEvent>);
1085
1086 impl AgentEventSink for Events {
1087 fn emit(&mut self, event: AgentEvent) {
1088 self.0.push(event);
1089 }
1090 }
1091
1092 fn tool_call(provenance: ToolExecutionProvenance) -> ToolCall {
1093 ToolCall {
1094 id: "call_1".into(),
1095 name: "count".into(),
1096 arguments_json: "{}".into(),
1097 argument_integrity: ToolArgumentIntegrity::Valid,
1098 provisional_id: None,
1099 provider_result: (provenance == ToolExecutionProvenance::Provider)
1100 .then(|| "provider result".into()),
1101 provenance,
1102 }
1103 }
1104
1105 fn named_tool_call(id: &str, name: &str) -> ToolCall {
1106 ToolCall {
1107 id: id.into(),
1108 name: name.into(),
1109 arguments_json: "{}".into(),
1110 argument_integrity: ToolArgumentIntegrity::Valid,
1111 provisional_id: None,
1112 provider_result: None,
1113 provenance: ToolExecutionProvenance::FxLocal,
1114 }
1115 }
1116
1117 fn response(content: Option<&str>, tool_calls: Vec<ToolCall>) -> GatewayResponse {
1118 GatewayResponse {
1119 content: content.map(str::to_owned),
1120 tool_calls,
1121 generation_id: Some("generation".into()),
1122 finish_reason: Some(FinishReason::Stop),
1123 usage: Usage {
1124 input_tokens: Some(2),
1125 output_tokens: Some(3),
1126 },
1127 delivery_ambiguous: false,
1128 }
1129 }
1130
1131 fn run_ready(
1132 agent: &Agent,
1133 context: &ToolContext,
1134 permissions: &mut PermissionEngine,
1135 approvals: &mut dyn ApprovalHandler,
1136 events: &mut dyn AgentEventSink,
1137 ) -> Result<AgentResult, AgentError> {
1138 let mut future = agent.run(
1139 AgentRequest {
1140 history: Vec::new(),
1141 prompt: "go".into(),
1142 },
1143 context,
1144 permissions,
1145 approvals,
1146 events,
1147 );
1148 let mut task_context = Context::from_waker(Waker::noop());
1149 match future.as_mut().poll(&mut task_context) {
1150 Poll::Ready(result) => result,
1151 Poll::Pending => panic!("scripted agent unexpectedly yielded"),
1152 }
1153 }
1154
1155 fn run_controlled_ready(
1156 agent: &Agent,
1157 context: &ToolContext,
1158 permissions: &mut PermissionEngine,
1159 approvals: &mut dyn ApprovalHandler,
1160 events: &mut dyn AgentEventSink,
1161 cancellation: Arc<dyn CancellationSignal>,
1162 ) -> Result<AgentResult, AgentError> {
1163 let mut future = agent.run_controlled(
1164 AgentRequest {
1165 history: Vec::new(),
1166 prompt: "go".into(),
1167 },
1168 context,
1169 permissions,
1170 approvals,
1171 events,
1172 cancellation,
1173 );
1174 let mut task_context = Context::from_waker(Waker::noop());
1175 match future.as_mut().poll(&mut task_context) {
1176 Poll::Ready(result) => result,
1177 Poll::Pending => panic!("scripted agent unexpectedly yielded"),
1178 }
1179 }
1180
1181 fn run_eventually_ready(
1182 agent: &Agent,
1183 context: &ToolContext,
1184 permissions: &mut PermissionEngine,
1185 approvals: &mut dyn ApprovalHandler,
1186 events: &mut dyn AgentEventSink,
1187 ) -> Result<AgentResult, AgentError> {
1188 let mut future = agent.run(
1189 AgentRequest {
1190 history: Vec::new(),
1191 prompt: "go".into(),
1192 },
1193 context,
1194 permissions,
1195 approvals,
1196 events,
1197 );
1198 let mut task_context = Context::from_waker(Waker::noop());
1199 for _ in 0..16 {
1200 if let Poll::Ready(result) = future.as_mut().poll(&mut task_context) {
1201 return result;
1202 }
1203 }
1204 panic!("cooperative agent did not complete")
1205 }
1206
1207 #[test]
1208 fn complete_large_tool_output_is_stored_before_model_projection() {
1209 let store = Arc::new(MemoryToolResultStore::default());
1210 let mut context = ToolContext::new(std::env::current_dir().unwrap());
1211 context.tool_results = Some(store.clone());
1212 let call = named_tool_call("call-sensitive", "mcp_remote_search");
1213 let complete = format!(
1214 "evidence\nAPI_KEY=server-private-value\n{}\nneedle",
1215 "你".repeat(6_000)
1216 );
1217 let mut output = ToolOutput {
1218 content: "old truncated projection".into(),
1219 is_error: false,
1220 structured: Some(json!({"large": true})),
1221 original_bytes: complete.len(),
1222 truncated: true,
1223 durable_content: Some(complete.clone()),
1224 };
1225
1226 finalize_tool_output(&context, &call, &mut output);
1227
1228 let stored = store.stored.lock().unwrap();
1229 assert_eq!(stored.len(), 1);
1230 assert!(stored[0].contains("API_KEY=[redacted]"));
1231 assert!(!stored[0].contains("server-private-value"));
1232 assert!(output.content.contains("result-test.txt"));
1233 assert!(output.content.contains("Use read_tool_result"));
1234 assert!(output.content.is_char_boundary(output.content.len()));
1235 assert!(output.structured.is_none());
1236 assert!(output.truncated);
1237 assert!(output.durable_content.is_none());
1238 }
1239
1240 #[test]
1241 fn newly_selected_scoped_context_defers_write_until_next_generation() {
1242 let executions = Arc::new(AtomicUsize::new(0));
1243 let gateway = Arc::new(ScriptedGateway {
1244 responses: Mutex::new(VecDeque::from([
1245 response(None, vec![named_tool_call("call_1", "scoped_write")]),
1246 response(None, vec![named_tool_call("call_2", "scoped_write")]),
1247 response(Some("done"), Vec::new()),
1248 ])),
1249 requests: Mutex::new(Vec::new()),
1250 });
1251 let mut registry = ToolRegistry::default();
1252 registry
1253 .register(ScopedWriteTool(executions.clone()))
1254 .unwrap();
1255 let agent = Agent::new(
1256 gateway.clone(),
1257 Arc::new(registry),
1258 AgentOptions::new("model"),
1259 );
1260 let mut context = ToolContext::new(std::env::current_dir().unwrap());
1261 context.project_context = Some(Arc::new(OneScopedDelta(AtomicUsize::new(0))));
1262 let mut permissions = PermissionEngine::new(PermissionMode::Yolo, Vec::new());
1263 let mut approvals = StaticApprovalHandler::allow_once();
1264 let mut events = Events::default();
1265
1266 let result = run_ready(
1267 &agent,
1268 &context,
1269 &mut permissions,
1270 &mut approvals,
1271 &mut events,
1272 )
1273 .unwrap();
1274
1275 assert_eq!(result.output, "done");
1276 assert_eq!(executions.load(Ordering::SeqCst), 1);
1277 let requests = gateway.requests.lock().unwrap();
1278 assert_eq!(requests.len(), 3);
1279 assert_eq!(requests[1].messages[0].role, Role::System);
1280 assert_eq!(
1281 requests[1].messages[0].content.as_deref(),
1282 Some("NESTED PROJECT RULE")
1283 );
1284 assert!(requests[1].messages.iter().any(|message| {
1285 message
1286 .content
1287 .as_deref()
1288 .is_some_and(|content| content.contains("execution deferred"))
1289 }));
1290 assert!(
1291 requests[2]
1292 .messages
1293 .iter()
1294 .any(|message| { message.content.as_deref() == Some("written") })
1295 );
1296 }
1297
1298 #[test]
1299 fn local_tool_round_trip_preserves_message_order_and_usage() {
1300 let executions = Arc::new(AtomicUsize::new(0));
1301 let gateway = Arc::new(ScriptedGateway {
1302 responses: Mutex::new(VecDeque::from([
1303 response(None, vec![tool_call(ToolExecutionProvenance::FxLocal)]),
1304 response(Some("done"), Vec::new()),
1305 ])),
1306 requests: Mutex::new(Vec::new()),
1307 });
1308 let mut registry = ToolRegistry::default();
1309 registry.register(CountingTool(executions.clone())).unwrap();
1310 let agent = Agent::new(
1311 gateway.clone(),
1312 Arc::new(registry),
1313 AgentOptions::new("model"),
1314 );
1315 let context = ToolContext::new(std::env::current_dir().unwrap());
1316 let mut permissions = PermissionEngine::new(PermissionMode::Auto, Vec::new());
1317 let mut approvals = StaticApprovalHandler::allow_once();
1318 let mut events = Events::default();
1319
1320 let result = run_ready(
1321 &agent,
1322 &context,
1323 &mut permissions,
1324 &mut approvals,
1325 &mut events,
1326 )
1327 .unwrap();
1328 assert_eq!(executions.load(Ordering::SeqCst), 1);
1329 assert_eq!(result.output, "done");
1330 assert_eq!(result.steps, 2);
1331 assert_eq!(result.usage.input_tokens, Some(4));
1332 assert_eq!(result.messages[2].role, Role::Tool);
1333 assert_eq!(result.messages[2].content.as_deref(), Some("counted"));
1334 assert_eq!(gateway.requests.lock().unwrap().len(), 2);
1335 assert!(events.0.iter().any(|event| matches!(
1336 event,
1337 AgentEvent::ToolFinished {
1338 is_error: false,
1339 ..
1340 }
1341 )));
1342 }
1343
1344 #[test]
1345 fn consecutive_read_tools_run_concurrently_and_commit_results_in_call_order() {
1346 let started = Arc::new(AtomicUsize::new(0));
1347 let gateway = Arc::new(ScriptedGateway {
1348 responses: Mutex::new(VecDeque::from([
1349 response(
1350 None,
1351 vec![
1352 named_tool_call("call-a", "read_a"),
1353 named_tool_call("call-b", "read_b"),
1354 ],
1355 ),
1356 response(Some("done"), Vec::new()),
1357 ])),
1358 requests: Mutex::new(Vec::new()),
1359 });
1360 let mut registry = ToolRegistry::default();
1361 registry
1362 .register(CooperativeReadTool {
1363 name: "read_a",
1364 started: started.clone(),
1365 })
1366 .unwrap();
1367 registry
1368 .register(CooperativeReadTool {
1369 name: "read_b",
1370 started: started.clone(),
1371 })
1372 .unwrap();
1373 let agent = Agent::new(gateway, Arc::new(registry), AgentOptions::new("model"));
1374 let context = ToolContext::new(std::env::current_dir().unwrap());
1375 let mut permissions = PermissionEngine::new(PermissionMode::Auto, Vec::new());
1376 let mut approvals = StaticApprovalHandler::deny();
1377 let mut events = Events::default();
1378
1379 let result = run_eventually_ready(
1380 &agent,
1381 &context,
1382 &mut permissions,
1383 &mut approvals,
1384 &mut events,
1385 )
1386 .unwrap();
1387 assert_eq!(started.load(Ordering::SeqCst), 2);
1388 assert_eq!(result.messages[2].tool_call_id.as_deref(), Some("call-a"));
1389 assert_eq!(result.messages[2].content.as_deref(), Some("read_a"));
1390 assert_eq!(result.messages[3].tool_call_id.as_deref(), Some("call-b"));
1391 assert_eq!(result.messages[3].content.as_deref(), Some("read_b"));
1392 let phases = events
1393 .0
1394 .iter()
1395 .filter_map(|event| match event {
1396 AgentEvent::ToolStarted { id, .. } => Some(format!("start:{id}")),
1397 AgentEvent::ToolFinished { id, .. } => Some(format!("finish:{id}")),
1398 AgentEvent::Gateway(_) => None,
1399 })
1400 .collect::<Vec<_>>();
1401 assert_eq!(
1402 phases,
1403 [
1404 "start:call-a",
1405 "start:call-b",
1406 "finish:call-a",
1407 "finish:call-b"
1408 ]
1409 );
1410 }
1411
1412 #[test]
1413 fn provider_tool_result_is_never_dispatched_locally() {
1414 let executions = Arc::new(AtomicUsize::new(0));
1415 let gateway = Arc::new(ScriptedGateway {
1416 responses: Mutex::new(VecDeque::from([
1417 response(None, vec![tool_call(ToolExecutionProvenance::Provider)]),
1418 response(Some("done"), Vec::new()),
1419 ])),
1420 requests: Mutex::new(Vec::new()),
1421 });
1422 let mut registry = ToolRegistry::default();
1423 registry.register(CountingTool(executions.clone())).unwrap();
1424 let agent = Agent::new(
1425 gateway.clone(),
1426 Arc::new(registry),
1427 AgentOptions::new("model"),
1428 );
1429 let context = ToolContext::new(std::env::current_dir().unwrap());
1430 let mut permissions = PermissionEngine::new(PermissionMode::Auto, Vec::new());
1431 let mut approvals = StaticApprovalHandler::deny();
1432 let mut events = Events::default();
1433
1434 let result = run_ready(
1435 &agent,
1436 &context,
1437 &mut permissions,
1438 &mut approvals,
1439 &mut events,
1440 )
1441 .unwrap();
1442 assert_eq!(executions.load(Ordering::SeqCst), 0);
1443 assert_eq!(
1444 result.messages[2].content.as_deref(),
1445 Some("provider result")
1446 );
1447 }
1448
1449 #[test]
1450 fn denied_review_becomes_permission_feedback_for_next_generation() {
1451 let executions = Arc::new(AtomicUsize::new(0));
1452 let gateway = Arc::new(ScriptedGateway {
1453 responses: Mutex::new(VecDeque::from([
1454 response(None, vec![tool_call(ToolExecutionProvenance::FxLocal)]),
1455 response(Some("understood"), Vec::new()),
1456 ])),
1457 requests: Mutex::new(Vec::new()),
1458 });
1459 let mut registry = ToolRegistry::default();
1460 registry.register(CountingTool(executions.clone())).unwrap();
1461 let agent = Agent::new(
1462 gateway.clone(),
1463 Arc::new(registry),
1464 AgentOptions::new("model"),
1465 );
1466 let context = ToolContext::new(std::env::current_dir().unwrap());
1467 let mut permissions = PermissionEngine::new(PermissionMode::Auto, Vec::new());
1468 let mut approvals = StaticApprovalHandler::deny();
1469 let mut events = Events::default();
1470
1471 let result = run_ready(
1472 &agent,
1473 &context,
1474 &mut permissions,
1475 &mut approvals,
1476 &mut events,
1477 )
1478 .unwrap();
1479 assert_eq!(executions.load(Ordering::SeqCst), 0);
1480 assert!(result.messages[2].permission_feedback);
1481 assert!(
1482 result.messages[2]
1483 .content
1484 .as_deref()
1485 .unwrap()
1486 .contains("denied")
1487 );
1488 }
1489
1490 #[test]
1491 fn cancellation_after_gateway_prevents_local_tool_execution() {
1492 let cancellation = Arc::new(AtomicBool::new(false));
1493 let gateway_executions = Arc::new(AtomicUsize::new(0));
1494 let tool_executions = Arc::new(AtomicUsize::new(0));
1495 let gateway = Arc::new(CancellingGateway {
1496 cancellation: cancellation.clone(),
1497 executions: gateway_executions.clone(),
1498 });
1499 let mut registry = ToolRegistry::default();
1500 registry
1501 .register(CountingTool(tool_executions.clone()))
1502 .unwrap();
1503 let agent = Agent::new(gateway, Arc::new(registry), AgentOptions::new("model"));
1504 let context = ToolContext::new(std::env::current_dir().unwrap());
1505 let mut permissions = PermissionEngine::new(PermissionMode::Auto, Vec::new());
1506 let mut approvals = StaticApprovalHandler::allow_once();
1507 let mut events = Events::default();
1508
1509 let result = run_controlled_ready(
1510 &agent,
1511 &context,
1512 &mut permissions,
1513 &mut approvals,
1514 &mut events,
1515 Arc::new(AtomicCancellation(cancellation)),
1516 )
1517 .unwrap();
1518
1519 assert_eq!(result.stop_reason, AgentStopReason::Cancelled);
1520 assert_eq!(result.steps, 1);
1521 assert_eq!(result.output, "partial");
1522 assert_eq!(gateway_executions.load(Ordering::SeqCst), 1);
1523 assert_eq!(tool_executions.load(Ordering::SeqCst), 0);
1524 assert!(!events.0.iter().any(|event| matches!(
1525 event,
1526 AgentEvent::ToolStarted { .. } | AgentEvent::ToolFinished { .. }
1527 )));
1528 }
1529
1530 #[test]
1531 fn cooperative_gateway_cancellation_is_a_normal_agent_stop() {
1532 let cancellation = Arc::new(AtomicBool::new(false));
1533 let agent = Agent::new(
1534 Arc::new(CancelledGateway {
1535 cancellation: cancellation.clone(),
1536 }),
1537 Arc::new(ToolRegistry::default()),
1538 AgentOptions::new("model"),
1539 );
1540 let context = ToolContext::new(std::env::current_dir().unwrap());
1541 let mut permissions = PermissionEngine::new(PermissionMode::Ask, Vec::new());
1542 let mut approvals = StaticApprovalHandler::deny();
1543 let mut events = Events::default();
1544 let result = run_controlled_ready(
1545 &agent,
1546 &context,
1547 &mut permissions,
1548 &mut approvals,
1549 &mut events,
1550 Arc::new(AtomicCancellation(cancellation)),
1551 )
1552 .unwrap();
1553 assert_eq!(result.stop_reason, AgentStopReason::Cancelled);
1554 assert_eq!(result.steps, 1);
1555 }
1556}