1use async_trait::async_trait;
20use futures::StreamExt;
21use serde::{Deserialize, Serialize};
22use serde_json::json;
23use sha2::{Digest, Sha256};
24use std::collections::{BTreeSet, HashMap};
25use std::sync::Arc;
26use std::time::Instant;
27use uuid::Uuid;
28
29use super::{Atom, AtomContext};
30use crate::annotation_hook::{
31 AnnotationProvider, VerifierProvider, collect_annotations, verify_annotations,
32};
33use crate::capabilities::CapabilityRegistry;
34use crate::driver_registry::{
35 DriverRegistry, LlmCompletionMetadata, LlmMessage, LlmMessageContent, LlmMessageRole,
36 LlmStreamEvent,
37};
38use crate::error::{AgentLoopError, Result};
39use crate::events::{
40 CapabilityUsageData, CapabilityUsageKind, CapabilityUsageRecord, EventContext, EventRequest,
41 LlmCompactionInfo, LlmGenerationData, LlmPromptCacheInfo, LlmRequestOptions, LlmRetryInfo,
42 LlmToolSearchInfo, OutputMessageCompletedData, OutputMessageDeltaData,
43 OutputMessageReplacedData, OutputMessageStartedData, ReasonCompletedData, ReasonItemData,
44 ReasonRecoveredData, ReasonStartedData, ReasonThinkingCompletedData, ReasonThinkingDeltaData,
45 ReasonThinkingStartedData, RecoveryMode, TokenUsage, ToolDefinitionSummary,
46 TranscriptRepairAction, TranscriptRepairedData,
47};
48use crate::llm_retry::{
49 LlmRetryConfig, RetryMetadata, is_transient_error_message, is_transient_stream_error,
50};
51use crate::message::{Message, MessageRole};
52use crate::message_retriever::MessageRetriever;
53use crate::openresponses_protocol::{CompactRequest, messages_to_compact_input};
54use crate::output_guardrail::{
55 ArmedGuardrail, OutputGuardrailContext, PostGenerationOutputContext, PostGenerationProvider,
56 TrippedGuardrail, evaluate_guardrails, evaluate_post_generation_guardrails,
57 post_generation_guardrail_text,
58};
59use crate::runtime_context::{AssembledTurnContext, assemble_turn_context};
60use crate::tool_types::{ToolCall, ToolDefinition};
61use crate::traits::{
62 AgentStore, DurableToolCallStatus, DurableToolResultStore, EventEmitter, HarnessStore,
63 ImageResolver, PartialStreamState, PartialStreamStore, ProviderStore, ResolvedImage,
64 ResolvedModel, SessionStore,
65};
66use crate::typed_id::{AgentId, HarnessId, MessageId, SessionId};
67use crate::{ErrorDisclosure, UserFacingError, UserFacingErrorContext, user_facing_error_codes};
68
69#[allow(clippy::too_many_arguments)] async fn apply_tool_call_repair(
89 capability_registry: &CapabilityRegistry,
90 event_emitter: &dyn EventEmitter,
91 session_id: SessionId,
92 context: &AtomContext,
93 resolved_capability_configs: &[crate::AgentCapabilityConfig],
94 tool_definitions: &[ToolDefinition],
95 tool_calls: &mut [ToolCall],
96 iteration: u32,
97) {
98 use crate::capabilities::{
99 RepairOutcome, SalvageResult, TOOL_CALL_REPAIR_CAPABILITY_ID, ToolCallRepairConfig,
100 salvage_tool_arguments,
101 };
102
103 let Some(cfg) = resolved_capability_configs.iter().find(|c| {
105 capability_registry.canonical_id(c.capability_ref.as_str())
106 == Some(TOOL_CALL_REPAIR_CAPABILITY_ID)
107 }) else {
108 return;
109 };
110 let repair_config = ToolCallRepairConfig::from_json(&cfg.config);
111
112 for call in tool_calls.iter_mut() {
113 let schema = tool_definitions
115 .iter()
116 .find(|t| t.name() == call.name)
117 .map(|t| t.full_parameters().clone());
118
119 let outcome = match salvage_tool_arguments(&call.arguments, schema.as_ref()) {
120 SalvageResult::AlreadyValid => continue,
122 SalvageResult::Repaired(fixed) => {
123 call.arguments = fixed;
124 RepairOutcome::LocalSalvage
125 }
126 SalvageResult::Unsalvageable => {
133 repair_config.outcome_after_failed_salvage(iteration.saturating_sub(1))
134 }
135 };
136
137 tracing::info!(
138 session_id = %session_id,
139 turn_id = %context.turn_id,
140 tool_call_id = %call.id,
141 tool_name = %call.name,
142 outcome = outcome.label(),
143 "ReasonAtom: tool-call repair"
144 );
145
146 if let Err(e) = event_emitter
147 .emit(EventRequest::new(
148 session_id,
149 EventContext::from_atom_context(context),
150 crate::events::ToolCallRepairedData {
151 turn_id: context.turn_id,
152 tool_call_id: call.id.clone(),
153 tool_name: call.name.clone(),
154 outcome: outcome.label().to_string(),
155 },
156 ))
157 .await
158 {
159 tracing::warn!(
160 session_id = %session_id,
161 error = %e,
162 "ReasonAtom: failed to emit tool.call_repaired event"
163 );
164 }
165 }
166}
167
168async fn repair_dangling_tool_calls(
180 messages: &[Message],
181 durable_store: Option<&dyn DurableToolResultStore>,
182 event_emitter: &dyn EventEmitter,
183 session_id: crate::typed_id::SessionId,
184 event_context: &EventContext,
185 turn_id: &str,
186) -> Vec<Message> {
187 let mut result = Vec::new();
188
189 for (i, msg) in messages.iter().enumerate() {
190 result.push(msg.clone());
191
192 if msg.role != MessageRole::Agent || !msg.has_tool_calls() {
193 continue;
194 }
195
196 for tc in msg.tool_calls() {
197 let has_result = messages[(i + 1)..]
198 .iter()
199 .any(|m| m.role == MessageRole::ToolResult && m.tool_call_id() == Some(&tc.id));
200
201 if has_result {
202 continue;
203 }
204
205 let (repair_msg, action) = if let Some(store) = durable_store {
207 match store.get_tool_call_status(turn_id, &tc.id).await {
208 Ok(Some(DurableToolCallStatus::Settled { result_json })) => {
209 let repair = match serde_json::from_value::<crate::tool_types::ToolResult>(
211 result_json.clone(),
212 ) {
213 Ok(tr) => Message::tool_result(&tc.id, tr.result, tr.error),
214 Err(_) => Message::tool_result(&tc.id, Some(result_json), None),
215 };
216 (repair, TranscriptRepairAction::Replay)
217 }
218 Ok(Some(DurableToolCallStatus::Interrupted { result_json })) => {
219 let err = result_json
221 .as_ref()
222 .and_then(|v| {
223 serde_json::from_value::<crate::tool_types::ToolResult>(v.clone())
224 .ok()
225 })
226 .and_then(|tr| tr.error)
227 .unwrap_or_else(|| {
228 "tool execution did not complete before recovery; result unknown"
229 .to_string()
230 });
231 (
232 Message::tool_result(&tc.id, None, Some(err)),
233 TranscriptRepairAction::Replay,
234 )
235 }
236 Ok(Some(DurableToolCallStatus::Running)) => {
237 (
239 Message::tool_result(
240 &tc.id,
241 None,
242 Some(
243 "interrupted - tool execution was interrupted by worker \
244 failure and the result is uncertain; do not retry \
245 automatically"
246 .to_string(),
247 ),
248 ),
249 TranscriptRepairAction::Synthesize,
250 )
251 }
252 Ok(None) => {
253 (
255 Message::tool_result(
256 &tc.id,
257 None,
258 Some(
259 "interrupted - tool was not executed before recovery; \
260 safe to retry"
261 .to_string(),
262 ),
263 ),
264 TranscriptRepairAction::Synthesize,
265 )
266 }
267 Err(e) => {
268 tracing::warn!(
270 tool_call_id = %tc.id,
271 error = %e,
272 "transcript repair: durable store error; status unknown"
273 );
274 (
275 Message::tool_result(
276 &tc.id,
277 None,
278 Some(
279 "interrupted - tool execution status unknown due to \
280 store error; do not retry automatically"
281 .to_string(),
282 ),
283 ),
284 TranscriptRepairAction::Synthesize,
285 )
286 }
287 }
288 } else {
289 (
291 Message::tool_result(
292 &tc.id,
293 None,
294 Some(
295 "cancelled - another message came in before it could be completed"
296 .to_string(),
297 ),
298 ),
299 TranscriptRepairAction::Synthesize,
300 )
301 };
302
303 let repair_event = EventRequest::new(
305 session_id,
306 event_context.clone(),
307 TranscriptRepairedData {
308 tool_call_id: tc.id.clone(),
309 tool_name: Some(tc.name.clone()),
310 action,
311 },
312 );
313 if let Err(e) = event_emitter.emit(repair_event).await {
314 tracing::warn!(
315 tool_call_id = %tc.id,
316 error = %e,
317 "transcript repair: failed to emit transcript.repaired event"
318 );
319 }
320
321 result.push(repair_msg);
322 }
323 }
324
325 result
326}
327
328const ERROR_PLACEHOLDER_MESSAGES: &[&str] = &[
331 "I encountered an error while processing your request. Please try again later.",
332 "The AI provider is experiencing issues. Please try again shortly.",
333 "Rate limited by the AI provider. Please wait a moment.",
334 "The conversation has become too long for the model to process. Please start a new session or reduce the context size.",
335 "There is a misconfiguration with the AI provider. Please contact support.",
336 "The AI provider account is out of credits or quota. Add credits or raise the provider account limits to continue.",
337];
338
339fn stream_event_advances_stall_deadline(event: &LlmStreamEvent) -> bool {
341 match event {
342 LlmStreamEvent::TextDelta(delta) | LlmStreamEvent::ThinkingDelta(delta) => {
343 !delta.is_empty()
344 }
345 LlmStreamEvent::ReasonItem {
346 encrypted_content,
347 summary,
348 token_count,
349 ..
350 } => {
351 encrypted_content
352 .as_ref()
353 .is_some_and(|content| !content.is_empty())
354 || summary.iter().any(|item| !item.is_empty())
355 || token_count.is_some_and(|count| count > 0)
356 }
357 LlmStreamEvent::ToolCalls(calls) => !calls.is_empty(),
358 LlmStreamEvent::MessagePhase(_)
360 | LlmStreamEvent::ThinkingSignature(_)
361 | LlmStreamEvent::Done(_)
362 | LlmStreamEvent::Error(_) => false,
363 }
364}
365
366fn should_retry_stream_error(
367 error: &crate::driver_registry::LlmStreamError,
368 retry_attempts: u32,
369 max_retries: u32,
370 has_output: bool,
371) -> bool {
372 retry_attempts < max_retries && !has_output && is_transient_stream_error(error)
373}
374
375fn merge_retry_metadata(
376 existing: Option<RetryMetadata>,
377 additional: &RetryMetadata,
378) -> Option<RetryMetadata> {
379 if !additional.had_retries() {
380 return existing;
381 }
382
383 let mut merged = existing.unwrap_or_default();
384 merged.attempts += additional.attempts;
385 merged.total_retry_wait += additional.total_retry_wait;
386 if additional.last_rate_limit_info.is_some() {
387 merged.last_rate_limit_info = additional.last_rate_limit_info.clone();
388 }
389 Some(merged)
390}
391
392fn unix_now_secs() -> u64 {
393 std::time::SystemTime::now()
394 .duration_since(std::time::UNIX_EPOCH)
395 .unwrap_or_default()
396 .as_secs()
397}
398
399fn is_error_placeholder_message(msg: &Message) -> bool {
402 if msg.role != MessageRole::Agent {
403 return false;
404 }
405 if msg.has_tool_calls() {
407 return false;
408 }
409 if let Some(metadata) = &msg.metadata
410 && let Some(serde_json::Value::String(code)) = metadata.get("error_code")
411 {
412 return matches!(
413 code.as_str(),
414 user_facing_error_codes::BUDGET_EXHAUSTED
415 | user_facing_error_codes::BUDGET_PAUSED
416 | user_facing_error_codes::MODEL_UNAVAILABLE
417 | user_facing_error_codes::REQUEST_TOO_LARGE
418 | user_facing_error_codes::PROVIDER_RATE_LIMITED
419 | user_facing_error_codes::PROVIDER_USAGE_LIMIT_REACHED
420 | user_facing_error_codes::PROVIDER_MISCONFIGURED
421 | user_facing_error_codes::PROVIDER_QUOTA_EXHAUSTED
422 | user_facing_error_codes::PROVIDER_UNAVAILABLE
423 | user_facing_error_codes::DEPENDENCY_UNAVAILABLE
424 | user_facing_error_codes::PROCESSING_ERROR
425 );
426 }
427 let text = msg.text().unwrap_or("");
428 ERROR_PLACEHOLDER_MESSAGES.contains(&text) || is_dynamic_error_placeholder(text)
429}
430
431fn append_guarded_thinking_delta(
432 armed_guardrails: &mut [ArmedGuardrail],
433 thinking: &mut String,
434 pending_thinking_delta: &mut String,
435 delta: &str,
436) -> Option<TrippedGuardrail> {
437 thinking.push_str(delta);
438
439 if let Some(t) = evaluate_guardrails(armed_guardrails, thinking, delta) {
443 pending_thinking_delta.clear();
444 Some(t)
445 } else {
446 pending_thinking_delta.push_str(delta);
447 None
448 }
449}
450
451fn error_disclosure_override(messages: &[Message]) -> Option<String> {
455 messages
456 .iter()
457 .rev()
458 .find(|m| m.role == MessageRole::User)?
459 .controls
460 .as_ref()?
461 .error_disclosure
462 .clone()
463}
464
465fn is_dynamic_error_placeholder(text: &str) -> bool {
466 (text.starts_with("Budget exhausted.") && text.ends_with("Increase the budget to continue."))
467 || (text.starts_with("Budget paused.")
468 && text.ends_with("Increase or resume the budget to continue."))
469 || (text.starts_with("Budget paused with ")
470 && text.ends_with("Increase or resume the budget to continue."))
471 || (text.starts_with("Soft limit reached.") && text.ends_with("soft limit."))
472 || (text.starts_with("The model `") && text.ends_with("Please select a different model."))
473}
474
475#[derive(Debug, Clone, Serialize, Deserialize)]
481pub struct ReasonInput {
482 pub context: AtomContext,
484 pub harness_id: HarnessId,
486 #[serde(skip_serializing_if = "Option::is_none")]
488 pub agent_id: Option<AgentId>,
489 #[serde(default)]
491 pub org_id: i64,
492 #[serde(default)]
496 pub mcp_tool_definitions: Vec<ToolDefinition>,
497 #[serde(skip_serializing_if = "Option::is_none")]
500 pub previous_response_id: Option<String>,
501 #[serde(default = "default_iteration")]
504 pub iteration: u32,
505}
506
507fn default_iteration() -> u32 {
508 1
509}
510
511#[derive(Debug, Clone, Default, Serialize, Deserialize)]
513pub struct ReasonResult {
514 pub success: bool,
516 pub text: String,
518 #[serde(default)]
520 pub tool_calls: Vec<ToolCall>,
521 pub has_tool_calls: bool,
523 #[serde(default)]
525 pub tool_definitions: Vec<ToolDefinition>,
526 #[serde(default = "default_max_iterations")]
528 pub max_iterations: usize,
529 #[serde(skip_serializing_if = "Option::is_none")]
531 pub error: Option<String>,
532 #[serde(default, skip_serializing_if = "Option::is_none")]
536 pub user_facing_error: Option<UserFacingError>,
537 #[serde(default, skip_serializing_if = "Option::is_none")]
539 pub error_disclosure: Option<ErrorDisclosure>,
540 #[serde(skip_serializing_if = "Option::is_none")]
542 pub usage: Option<TokenUsage>,
543 #[serde(skip_serializing_if = "Option::is_none")]
545 pub output_message_id: Option<MessageId>,
546 #[serde(skip_serializing_if = "Option::is_none")]
548 pub time_to_first_token_ms: Option<u64>,
549 #[serde(skip_serializing_if = "Option::is_none")]
551 pub response_id: Option<String>,
552 #[serde(default, skip_serializing_if = "Option::is_none")]
554 pub finish_reason: Option<String>,
555 #[serde(skip_serializing_if = "Option::is_none")]
557 pub locale: Option<String>,
558 #[serde(default, skip_serializing_if = "Option::is_none")]
560 pub network_access: Option<crate::network_access::NetworkAccessList>,
561 #[serde(default, skip_serializing_if = "Option::is_none")]
565 pub parallel_tool_calls: Option<bool>,
566}
567
568fn default_max_iterations() -> usize {
569 500
570}
571
572const CHECKPOINT_REARM_MIN_SUFFIX_MESSAGES: usize = 4;
573const PROACTIVE_RETRY_MIN_TOKEN_GROWTH: u64 = 4_096;
574const PROACTIVE_RETRY_GROWTH_DIVISOR: u64 = 20;
575
576fn proactive_source_fingerprint(
577 provider_opaque_context: Option<&crate::ProviderOpaqueContext>,
578 messages: &[LlmMessage],
579) -> [u8; 32] {
580 let mut input = match provider_opaque_context {
581 Some(crate::ProviderOpaqueContext::OpenResponsesCompact { output }) => {
582 output.iter().map(crate::CompactInputItem::from).collect()
583 }
584 None => Vec::new(),
585 };
586 input.extend(messages_to_compact_input(messages));
587 let bytes = serde_json::to_vec(&input).unwrap_or_default();
588 Sha256::digest(bytes).into()
589}
590
591#[derive(Debug)]
592struct AppliedNativeCompaction {
593 checkpoint_id: Option<String>,
594 input_items_before: usize,
595 output_items_after: usize,
596 tokens_before: Option<u64>,
597 tokens_after: Option<u64>,
598 bytes_before: Option<u64>,
599 bytes_after: Option<u64>,
600 duration_ms: u64,
601}
602
603fn materially_reduced(before: u64, after: u64) -> bool {
604 const MIN_REDUCTION_UNITS: u64 = 32;
605 let five_percent = before.div_ceil(20);
606 let required_reduction = five_percent.max(MIN_REDUCTION_UNITS).min(before);
607 after < before && before - after >= required_reduction
608}
609
610#[allow(clippy::too_many_arguments)]
611async fn try_apply_native_compaction(
612 chat_driver: &dyn crate::ChatDriver,
613 checkpoint_store: Option<&Arc<dyn crate::CompactionCheckpointStore>>,
614 session_id: SessionId,
615 source_sequence: Option<i64>,
616 provider_type: &str,
617 model: &str,
618 system_prompt: Option<&str>,
619 stateful_response_continuation: bool,
620 llm_messages: &mut Vec<LlmMessage>,
621 llm_config: &mut crate::driver_registry::LlmCallConfig,
622) -> Result<Option<AppliedNativeCompaction>> {
623 if !chat_driver.supports_compact() {
624 return Ok(None);
625 }
626
627 let started = Instant::now();
628 let has_system_prompt = system_prompt.is_some();
629 let messages_to_compact = if has_system_prompt {
630 &llm_messages[1..]
631 } else {
632 &llm_messages[..]
633 };
634 let (mut standalone_input, has_prior_opaque_context) =
635 match llm_config.provider_opaque_context.as_ref() {
636 Some(crate::ProviderOpaqueContext::OpenResponsesCompact { output }) => (
637 output.iter().map(crate::CompactInputItem::from).collect(),
638 true,
639 ),
640 None => (Vec::new(), false),
641 };
642 standalone_input.extend(messages_to_compact_input(messages_to_compact));
643 let input_items_before = standalone_input.len();
644 let local_tokens_before = (!stateful_response_continuation && !has_prior_opaque_context)
645 .then(|| crate::capabilities::estimate_total_tokens(messages_to_compact) as u64);
646 let bytes_before = (!stateful_response_continuation || has_prior_opaque_context)
647 .then(|| {
648 serde_json::to_vec(&standalone_input)
649 .ok()
650 .map(|value| value.len() as u64)
651 })
652 .flatten();
653 let (input, compact_previous_response_id) = (standalone_input, None);
657
658 let compact_response = match chat_driver
659 .compact(CompactRequest {
660 model: model.to_string(),
661 input,
662 previous_response_id: compact_previous_response_id,
663 instructions: system_prompt.map(str::to_string),
664 })
665 .await
666 {
667 Ok(Some(response)) => response,
668 Ok(None) => return Ok(None),
669 Err(error) => {
670 tracing::warn!(
671 session_id = %session_id,
672 error = %error,
673 "ReasonAtom: native compaction failed"
674 );
675 return Ok(None);
676 }
677 };
678
679 let tokens_before = compact_response
680 .usage
681 .as_ref()
682 .and_then(|usage| usage.input_tokens)
683 .map(u64::from)
684 .or(local_tokens_before);
685 let tokens_after = compact_response
686 .usage
687 .as_ref()
688 .and_then(|usage| usage.output_tokens)
689 .map(u64::from);
690 let bytes_after = serde_json::to_vec(&compact_response.output)
691 .ok()
692 .map(|value| value.len() as u64);
693 let effective = match (tokens_before, tokens_after) {
694 (Some(before), Some(after)) => materially_reduced(before, after),
695 _ => match (bytes_before, bytes_after) {
696 (Some(before), Some(after)) => materially_reduced(before, after),
697 _ => false,
698 },
699 };
700 if !effective {
701 tracing::info!(
702 session_id = %session_id,
703 ?tokens_before,
704 ?tokens_after,
705 ?bytes_before,
706 ?bytes_after,
707 "ReasonAtom: native compaction produced no material reduction"
708 );
709 return Ok(None);
710 }
711
712 let output_items_after = compact_response.output.len();
713 let opaque_context = crate::driver_registry::ProviderOpaqueContext::OpenResponsesCompact {
714 output: compact_response.output,
715 };
716 let checkpoint_id =
717 if let (Some(store), Some(source_sequence)) = (checkpoint_store, source_sequence) {
718 let id = Uuid::now_v7();
719 let installed = store
720 .install(crate::CompactionCheckpoint {
721 id,
722 session_id,
723 source_sequence,
724 provider_type: provider_type.to_string(),
725 model: model.to_string(),
726 format_version: crate::COMPACTION_CHECKPOINT_FORMAT_VERSION,
727 payload: crate::CompactionCheckpointPayload::ProviderOpaque {
728 context: opaque_context.clone(),
729 },
730 })
731 .await?;
732 if !installed {
733 return Err(AgentLoopError::store(
734 "a newer compaction checkpoint was installed concurrently",
735 ));
736 }
737 Some(id.to_string())
738 } else {
739 None
740 };
741
742 llm_config.previous_response_id = None;
745 llm_config.provider_opaque_context = Some(opaque_context);
746 llm_messages.retain(|message| message.role == LlmMessageRole::System);
747
748 Ok(Some(AppliedNativeCompaction {
749 checkpoint_id,
750 input_items_before,
751 output_items_after,
752 tokens_before,
753 tokens_after,
754 bytes_before,
755 bytes_after,
756 duration_ms: started.elapsed().as_millis() as u64,
757 }))
758}
759
760fn build_request_options(
761 config: &crate::driver_registry::LlmCallConfig,
762 provider: &str,
763) -> Option<LlmRequestOptions> {
764 let prompt_cache = config
765 .prompt_cache
766 .as_ref()
767 .filter(|cfg| cfg.enabled)
768 .map(|cfg| LlmPromptCacheInfo {
769 enabled: true,
770 strategy: cfg.strategy,
771 provider_mode: match provider {
772 "openai" => Some("prompt_cache_key".to_string()),
773 "anthropic" => Some("cache_control".to_string()),
774 "gemini" => Some(
775 if cfg.gemini_cached_content.is_some() {
776 "cached_content"
777 } else {
778 "implicit"
779 }
780 .to_string(),
781 ),
782 _ => None,
783 },
784 });
785
786 let tool_search = config
787 .tool_search
788 .as_ref()
789 .filter(|cfg| cfg.enabled)
790 .map(|cfg| LlmToolSearchInfo {
791 enabled: true,
792 threshold: cfg.threshold,
793 });
794
795 let mut provider_options = HashMap::new();
796 if provider == "openai" && config.previous_response_id.is_some() {
797 provider_options.insert(
798 "openai".to_string(),
799 json!({ "previous_response_id": true }),
800 );
801 }
802 if provider == "gemini"
803 && config
804 .prompt_cache
805 .as_ref()
806 .filter(|cfg| cfg.enabled)
807 .and_then(|cfg| cfg.gemini_cached_content.as_ref())
808 .is_some()
809 {
810 provider_options.insert("gemini".to_string(), json!({ "cached_content": true }));
811 }
812
813 let request_options = LlmRequestOptions {
814 prompt_cache,
815 tool_search,
816 provider_options,
817 metadata: config.metadata.clone(),
818 };
819
820 (!request_options.is_empty()).then_some(request_options)
821}
822
823fn capability_name_snapshot(registry: &CapabilityRegistry, capability_id: &str) -> Option<String> {
824 registry
825 .get(capability_id)
826 .map(|capability| capability.name().to_string())
827}
828
829fn capability_usage_snapshot_records(
830 registry: &CapabilityRegistry,
831 resolved_capability_configs: &[crate::AgentCapabilityConfig],
832 tool_definitions: &[ToolDefinition],
833) -> Vec<CapabilityUsageRecord> {
834 let mut records = Vec::new();
835 let mut seen = BTreeSet::new();
836
837 for config in resolved_capability_configs {
838 let capability_id = config.capability_id().to_string();
839 if seen.insert((
840 "resolved".to_string(),
841 capability_id.clone(),
842 None::<String>,
843 )) {
844 records.push(CapabilityUsageRecord {
845 capability_name: capability_name_snapshot(registry, &capability_id),
846 capability_id,
847 usage_kind: CapabilityUsageKind::Resolved,
848 tool_name: None,
849 usage_count: Some(1),
850 duration_ms: None,
851 });
852 }
853 }
854
855 for tool in tool_definitions {
856 let Some((capability_id, capability_name)) = tool.capability_attribution() else {
857 continue;
858 };
859 let capability_id = capability_id.to_string();
860 let tool_name = tool.name().to_string();
861 if seen.insert((
862 "exposed".to_string(),
863 capability_id.clone(),
864 Some(tool_name.clone()),
865 )) {
866 records.push(CapabilityUsageRecord {
867 capability_name: capability_name
868 .map(str::to_string)
869 .or_else(|| capability_name_snapshot(registry, &capability_id)),
870 capability_id,
871 usage_kind: CapabilityUsageKind::Exposed,
872 tool_name: Some(tool_name),
873 usage_count: Some(1),
874 duration_ms: None,
875 });
876 }
877 }
878
879 records
880}
881
882pub struct ReasonAtom {
901 harness_store: Arc<dyn HarnessStore>,
902 agent_store: Arc<dyn AgentStore>,
903 session_store: Arc<dyn SessionStore>,
904 message_retriever: Arc<dyn MessageRetriever>,
905 provider_store: Arc<dyn ProviderStore>,
906 capability_registry: CapabilityRegistry,
907 driver_registry: DriverRegistry,
908 event_emitter: Arc<dyn EventEmitter>,
909 image_resolver: Option<Arc<dyn ImageResolver>>,
911 file_store: Option<Arc<dyn crate::traits::SessionFileSystem>>,
914 stream_heartbeater: Option<Arc<dyn crate::traits::StreamHeartbeater>>,
916 provider_stall_timeout: Option<std::time::Duration>,
918 durable_tool_result_store: Option<Arc<dyn DurableToolResultStore>>,
920 partial_stream_store: Option<Arc<dyn PartialStreamStore>>,
922 reasoning_effort_handle: Option<crate::traits::ReasoningEffortHandle>,
926 utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
930 schedule_store: Option<Arc<dyn crate::traits::SessionScheduleStore>>,
936 compaction_checkpoint_store: Option<Arc<dyn crate::CompactionCheckpointStore>>,
938}
939
940impl ReasonAtom {
941 #[allow(clippy::too_many_arguments)]
943 pub fn new(
944 harness_store: impl HarnessStore + 'static,
945 agent_store: impl AgentStore + 'static,
946 session_store: impl SessionStore + 'static,
947 message_retriever: impl MessageRetriever + 'static,
948 provider_store: impl ProviderStore + 'static,
949 capability_registry: CapabilityRegistry,
950 driver_registry: DriverRegistry,
951 event_emitter: impl EventEmitter + 'static,
952 ) -> Self {
953 Self {
954 harness_store: Arc::new(harness_store),
955 agent_store: Arc::new(agent_store),
956 session_store: Arc::new(session_store),
957 message_retriever: Arc::new(message_retriever),
958 provider_store: Arc::new(provider_store),
959 capability_registry,
960 driver_registry,
961 event_emitter: Arc::new(event_emitter),
962 image_resolver: None,
963 file_store: None,
964 stream_heartbeater: None,
965 provider_stall_timeout: None,
966 durable_tool_result_store: None,
967 partial_stream_store: None,
968 reasoning_effort_handle: None,
969 utility_llm_service: None,
970 schedule_store: None,
971 compaction_checkpoint_store: None,
972 }
973 }
974
975 pub fn with_schedule_store(
978 mut self,
979 store: Arc<dyn crate::traits::SessionScheduleStore>,
980 ) -> Self {
981 self.schedule_store = Some(store);
982 self
983 }
984
985 pub fn with_compaction_checkpoint_store(
986 mut self,
987 store: Arc<dyn crate::CompactionCheckpointStore>,
988 ) -> Self {
989 self.compaction_checkpoint_store = Some(store);
990 self
991 }
992
993 fn collect_llm_error_hooks(
999 &self,
1000 resolved_capability_configs: &[crate::AgentCapabilityConfig],
1001 ) -> Vec<(
1002 Arc<dyn crate::llm_error_hook::LlmErrorHook>,
1003 serde_json::Value,
1004 )> {
1005 resolved_capability_configs
1006 .iter()
1007 .filter_map(|cfg| {
1008 let cap = self.capability_registry.get(cfg.capability_ref.as_str())?;
1009 let hook = cap.llm_error_hook()?;
1010 Some((hook, cfg.config.clone()))
1011 })
1012 .collect()
1013 }
1014
1015 pub fn with_file_store(
1022 mut self,
1023 file_store: Arc<dyn crate::traits::SessionFileSystem>,
1024 ) -> Self {
1025 self.file_store = Some(file_store);
1026 self
1027 }
1028
1029 pub fn with_image_resolver(mut self, resolver: Arc<dyn ImageResolver>) -> Self {
1042 self.image_resolver = Some(resolver);
1043 self
1044 }
1045
1046 pub fn with_stream_heartbeater(
1048 mut self,
1049 heartbeater: Arc<dyn crate::traits::StreamHeartbeater>,
1050 ) -> Self {
1051 self.stream_heartbeater = Some(heartbeater);
1052 self
1053 }
1054
1055 pub fn with_provider_stall_timeout(mut self, timeout: std::time::Duration) -> Self {
1058 self.provider_stall_timeout = Some(timeout);
1059 self
1060 }
1061
1062 pub fn with_durable_tool_result_store(
1068 mut self,
1069 store: Arc<dyn DurableToolResultStore>,
1070 ) -> Self {
1071 self.durable_tool_result_store = Some(store);
1072 self
1073 }
1074
1075 pub fn with_partial_stream_store(mut self, store: Arc<dyn PartialStreamStore>) -> Self {
1077 self.partial_stream_store = Some(store);
1078 self
1079 }
1080
1081 pub fn with_reasoning_effort_handle(
1088 mut self,
1089 handle: crate::traits::ReasoningEffortHandle,
1090 ) -> Self {
1091 self.reasoning_effort_handle = Some(handle);
1092 self
1093 }
1094
1095 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1098 self.utility_llm_service = Some(service);
1099 self
1100 }
1101}
1102
1103#[async_trait]
1104impl Atom for ReasonAtom {
1105 type Input = ReasonInput;
1106 type Output = ReasonResult;
1107
1108 fn name(&self) -> &'static str {
1109 "reason"
1110 }
1111
1112 async fn execute(&self, input: Self::Input) -> Result<Self::Output> {
1113 self.execute_inner(input, None).await
1114 }
1115}
1116
1117impl ReasonAtom {
1118 pub async fn execute_with_assembled_context(
1123 &self,
1124 input: ReasonInput,
1125 assembled: AssembledTurnContext,
1126 ) -> Result<ReasonResult> {
1127 self.execute_inner(input, Some(assembled)).await
1128 }
1129
1130 async fn emit_capability_usage_snapshot(
1131 &self,
1132 session_id: SessionId,
1133 context: &AtomContext,
1134 resolved_capability_configs: &[crate::AgentCapabilityConfig],
1135 tool_definitions: &[ToolDefinition],
1136 ) {
1137 let records = capability_usage_snapshot_records(
1138 &self.capability_registry,
1139 resolved_capability_configs,
1140 tool_definitions,
1141 );
1142 if records.is_empty() {
1143 return;
1144 }
1145
1146 if let Err(error) = self
1147 .event_emitter
1148 .emit(EventRequest::new(
1149 session_id,
1150 EventContext::from_atom_context(context),
1151 CapabilityUsageData { records },
1152 ))
1153 .await
1154 {
1155 tracing::warn!(
1156 session_id = %session_id,
1157 error = %error,
1158 "ReasonAtom: failed to emit capability.usage event"
1159 );
1160 }
1161 }
1162
1163 async fn repair_malformed_tool_calls(
1172 &self,
1173 session_id: SessionId,
1174 context: &AtomContext,
1175 resolved_capability_configs: &[crate::AgentCapabilityConfig],
1176 tool_definitions: &[ToolDefinition],
1177 tool_calls: &mut [ToolCall],
1178 iteration: u32,
1179 ) {
1180 apply_tool_call_repair(
1181 &self.capability_registry,
1182 self.event_emitter.as_ref(),
1183 session_id,
1184 context,
1185 resolved_capability_configs,
1186 tool_definitions,
1187 tool_calls,
1188 iteration,
1189 )
1190 .await;
1191 }
1192
1193 async fn execute_inner(
1194 &self,
1195 input: ReasonInput,
1196 assembled: Option<AssembledTurnContext>,
1197 ) -> Result<ReasonResult> {
1198 let ReasonInput {
1199 context,
1200 harness_id,
1201 agent_id,
1202 org_id,
1203 mcp_tool_definitions,
1204 previous_response_id,
1205 iteration,
1206 } = input;
1207
1208 tracing::info!(
1209 session_id = %context.session_id,
1210 turn_id = %context.turn_id,
1211 exec_id = %context.exec_id,
1212 harness_id = %harness_id,
1213 agent_id = ?agent_id,
1214 mcp_tools_count = %mcp_tool_definitions.len(),
1215 "ReasonAtom: starting LLM call"
1216 );
1217
1218 let trace_id = context.turn_id.to_string();
1226 let reason_span_id = Uuid::now_v7().to_string();
1227 let parent_span_id = trace_id.clone(); let event_context = EventContext::from_atom_context(&context).with_span(
1231 trace_id.clone(),
1232 reason_span_id.clone(),
1233 Some(parent_span_id.clone()),
1234 );
1235
1236 let reason_start = Instant::now();
1238
1239 if let Err(e) = self
1241 .event_emitter
1242 .emit(EventRequest::new(
1243 context.session_id,
1244 event_context.clone(),
1245 ReasonStartedData {
1246 harness_id,
1247 agent_id,
1248 metadata: None, },
1250 ))
1251 .await
1252 {
1253 tracing::warn!(
1254 session_id = %context.session_id,
1255 error = %e,
1256 "ReasonAtom: failed to emit reason.started event"
1257 );
1258 }
1259
1260 let assembled = match assembled {
1264 Some(assembled) => Ok(assembled),
1265 None => {
1266 assemble_turn_context(
1267 self.harness_store.as_ref(),
1268 self.agent_store.as_ref(),
1269 self.session_store.as_ref(),
1270 self.message_retriever.as_ref(),
1271 self.provider_store.as_ref(),
1272 &self.capability_registry,
1273 context.session_id,
1274 harness_id,
1275 agent_id,
1276 &mcp_tool_definitions,
1277 self.file_store.clone(),
1278 )
1279 .await
1280 }
1281 };
1282
1283 let (error_disclosure, error_context, error_hooks, call_result) = match assembled {
1284 Ok(assembled) => {
1285 let error_disclosure = crate::capabilities::resolve_error_disclosure(
1286 &assembled.resolved_capability_configs,
1287 error_disclosure_override(&assembled.messages).as_deref(),
1288 );
1289 let error_hooks =
1293 self.collect_llm_error_hooks(&assembled.resolved_capability_configs);
1294 let error_context = UserFacingErrorContext::default()
1295 .with_provider(assembled.model_with_provider.provider_type.to_string())
1296 .with_model_id(assembled.model_with_provider.model.clone());
1297 let call_result = self
1298 .execute_llm_call(
1299 context.session_id,
1300 harness_id,
1301 agent_id,
1302 org_id,
1303 &context,
1304 &trace_id,
1305 &reason_span_id,
1306 previous_response_id,
1307 iteration,
1308 assembled,
1309 )
1310 .await;
1311 (error_disclosure, error_context, error_hooks, call_result)
1312 }
1313 Err(error) => (
1314 ErrorDisclosure::default(),
1315 UserFacingErrorContext::default(),
1316 Vec::new(),
1317 Err(error),
1318 ),
1319 };
1320
1321 let result = match call_result {
1323 Ok(result) => {
1324 let reason_duration_ms = reason_start.elapsed().as_millis() as u64;
1326
1327 let completed_context = EventContext::from_atom_context(&context).with_span(
1329 trace_id.clone(),
1330 reason_span_id.clone(), Some(parent_span_id.clone()),
1332 );
1333 if let Err(e) = self
1334 .event_emitter
1335 .emit(EventRequest::new(
1336 context.session_id,
1337 completed_context,
1338 ReasonCompletedData::success(
1339 &result.text,
1340 result.has_tool_calls,
1341 result.tool_calls.len() as u32,
1342 Some(reason_duration_ms),
1343 result.usage.clone(),
1344 ),
1345 ))
1346 .await
1347 {
1348 tracing::warn!(
1349 session_id = %context.session_id,
1350 error = %e,
1351 "ReasonAtom: failed to emit reason.completed event"
1352 );
1353 }
1354 result
1355 }
1356 Err(e) => {
1357 let reason_duration_ms = reason_start.elapsed().as_millis() as u64;
1359
1360 tracing::warn!(
1363 session_id = %context.session_id,
1364 turn_id = %context.turn_id,
1365 error = %e,
1366 "ReasonAtom: LLM call failed"
1367 );
1368
1369 let error_msg = e.to_string();
1370 let mut source_error = e.user_facing_error(error_context);
1371
1372 let is_transient = e.is_transient_llm_error()
1379 || (e.llm_error_kind().is_none() && is_transient_error_message(&error_msg));
1380
1381 if !is_transient && !error_hooks.is_empty() {
1387 let services = crate::llm_error_hook::LlmErrorHookServices {
1388 schedule_store: self.schedule_store.clone(),
1389 };
1390 for (hook, config) in &error_hooks {
1391 let outcome = {
1392 let ctx = crate::llm_error_hook::LlmErrorContext {
1393 session_id: context.session_id,
1394 error_code: &source_error.code,
1395 error_fields: &source_error.fields,
1396 config,
1397 services: &services,
1398 };
1399 hook.on_llm_error(&ctx).await
1400 };
1401 for (key, value) in outcome.extra_error_fields {
1402 source_error = source_error.with_field(key, value);
1403 }
1404 }
1405 }
1406
1407 let user_error = source_error.apply_disclosure(error_disclosure, Some(&error_msg));
1408 let user_error_text = user_error.fallback_message();
1409
1410 let mut output_message_id = None;
1411
1412 if !is_transient {
1413 let mut error_message = Message::assistant(&user_error_text);
1415 let mut metadata = std::collections::HashMap::new();
1416 user_error.apply_to_message_metadata(&mut metadata);
1417 UserFacingError::apply_disclosure_to_message_metadata(
1418 &mut metadata,
1419 error_disclosure,
1420 &source_error.code,
1421 );
1422 error_message.metadata = Some(metadata);
1423
1424 output_message_id = Some(error_message.id);
1425
1426 let error_msg_context = EventContext::from_atom_context(&context).with_span(
1429 trace_id.clone(),
1430 Uuid::now_v7().to_string(), Some(reason_span_id.clone()), );
1433 if let Err(emit_err) = self
1434 .event_emitter
1435 .emit(EventRequest::new(
1436 context.session_id,
1437 error_msg_context,
1438 OutputMessageCompletedData::new(error_message)
1439 .with_user_facing_error(&user_error)
1440 .with_error_disclosure(error_disclosure),
1441 ))
1442 .await
1443 {
1444 tracing::warn!(
1445 session_id = %context.session_id,
1446 error = %emit_err,
1447 "ReasonAtom: failed to emit output.message.completed event for error"
1448 );
1449 }
1450 } else {
1451 tracing::info!(
1452 session_id = %context.session_id,
1453 "ReasonAtom: skipping error event for transient LLM error (will be retried)"
1454 );
1455 }
1456
1457 let completed_context = EventContext::from_atom_context(&context).with_span(
1459 trace_id.clone(),
1460 reason_span_id.clone(), Some(parent_span_id.clone()),
1462 );
1463 if let Err(emit_err) = self
1464 .event_emitter
1465 .emit(EventRequest::new(
1466 context.session_id,
1467 completed_context,
1468 ReasonCompletedData::failure(error_msg.clone(), Some(reason_duration_ms)),
1469 ))
1470 .await
1471 {
1472 tracing::warn!(
1473 session_id = %context.session_id,
1474 error = %emit_err,
1475 "ReasonAtom: failed to emit reason.completed event"
1476 );
1477 }
1478
1479 ReasonResult {
1480 success: false,
1481 text: user_error_text,
1482 tool_calls: vec![],
1483 has_tool_calls: false,
1484 tool_definitions: vec![],
1485 max_iterations: default_max_iterations(),
1486 error: Some(error_msg.clone()),
1487 user_facing_error: Some(user_error),
1488 error_disclosure: Some(error_disclosure),
1489 usage: None,
1490 output_message_id,
1491 time_to_first_token_ms: None,
1492 response_id: None,
1493 finish_reason: error_msg
1494 .to_ascii_lowercase()
1495 .contains("model refused")
1496 .then(|| "refusal".to_string()),
1497 locale: None,
1498 network_access: None,
1499 parallel_tool_calls: None,
1500 }
1501 }
1502 };
1503
1504 Ok(result)
1505 }
1506
1507 #[allow(clippy::too_many_arguments)]
1509 async fn execute_llm_call(
1510 &self,
1511 session_id: SessionId,
1512 harness_id: HarnessId,
1513 agent_id: Option<AgentId>,
1514 org_id: i64,
1515 context: &AtomContext,
1516 trace_id: &str,
1517 reason_span_id: &str,
1518 previous_response_id: Option<String>,
1519 iteration: u32,
1520 assembled: AssembledTurnContext,
1521 ) -> Result<ReasonResult> {
1522 let mut messages = assembled.messages;
1523 let mut message_source_sequence = assembled.message_source_sequence;
1524 let prior_usage = assembled.session.usage.clone();
1525 let model_with_provider = assembled.model_with_provider;
1526 let resolved_model_id = assembled.resolved_model_id;
1527 let resolved_locale = assembled.resolved_locale;
1528 let compaction_config = assembled.compaction_config;
1529 let resolved_capability_configs = assembled.resolved_capability_configs;
1530 let runtime_agent = assembled.runtime_agent;
1531 let embedder_metadata = assembled.embedder_metadata;
1532
1533 self.emit_capability_usage_snapshot(
1534 session_id,
1535 context,
1536 &resolved_capability_configs,
1537 &runtime_agent.tools,
1538 )
1539 .await;
1540
1541 let guardrail_providers: Vec<(
1548 &str,
1549 &serde_json::Value,
1550 Arc<dyn crate::output_guardrail::OutputGuardrail>,
1551 )> = resolved_capability_configs
1552 .iter()
1553 .filter_map(|cfg| {
1554 let cap_id = cfg.capability_ref.as_str();
1555 let cap = self.capability_registry.get(cap_id)?;
1556 let guards = cap.output_guardrails();
1557 if guards.is_empty() {
1558 return None;
1559 }
1560 Some(
1561 guards
1562 .into_iter()
1563 .map(move |g| (cap_id, &cfg.config, g))
1564 .collect::<Vec<_>>(),
1565 )
1566 })
1567 .flatten()
1568 .collect();
1569
1570 let post_output_providers: Vec<PostGenerationProvider> = resolved_capability_configs
1576 .iter()
1577 .filter_map(|cfg| {
1578 let cap_id = cfg.capability_ref.as_str();
1579 let cap = self.capability_registry.get(cap_id)?;
1580 let providers = cap.post_output_guardrails_with_config(&cfg.config);
1581 if providers.is_empty() {
1582 return None;
1583 }
1584 Some(
1585 providers
1586 .into_iter()
1587 .map(move |provider| PostGenerationProvider {
1588 capability_id: cap_id.to_string(),
1589 provider,
1590 })
1591 .collect::<Vec<_>>(),
1592 )
1593 })
1594 .flatten()
1595 .collect();
1596
1597 let annotation_providers: Vec<AnnotationProvider> = resolved_capability_configs
1602 .iter()
1603 .filter_map(|cfg| {
1604 let cap_id = cfg.capability_ref.as_str();
1605 let cap = self.capability_registry.get(cap_id)?;
1606 let providers = cap.post_output_annotation_hooks_with_config(&cfg.config);
1607 if providers.is_empty() {
1608 return None;
1609 }
1610 Some(
1611 providers
1612 .into_iter()
1613 .map(move |provider| AnnotationProvider {
1614 capability_id: cap_id.to_string(),
1615 provider,
1616 })
1617 .collect::<Vec<_>>(),
1618 )
1619 })
1620 .flatten()
1621 .collect();
1622
1623 let citation_verifiers: Vec<VerifierProvider> = resolved_capability_configs
1627 .iter()
1628 .filter_map(|cfg| {
1629 let cap_id = cfg.capability_ref.as_str();
1630 let cap = self.capability_registry.get(cap_id)?;
1631 let verifier = cap.citation_verifier_with_config(&cfg.config)?;
1632 Some(VerifierProvider {
1633 capability_id: cap_id.to_string(),
1634 provider: verifier,
1635 })
1636 })
1637 .collect();
1638
1639 let chat_driver = self.create_chat_driver(&model_with_provider)?;
1641 let stateful_response_continuation =
1642 previous_response_id.is_some() && chat_driver.supports_stateful_responses();
1643 let mut restored_checkpoint: Option<crate::CompactionCheckpoint> = None;
1644 let mut checkpoint_suffix_message_count = 0usize;
1645
1646 if compaction_config.is_some()
1647 && let Some(store) = self.compaction_checkpoint_store.as_ref()
1648 && let Some(checkpoint) = store
1649 .get_latest(
1650 session_id,
1651 model_with_provider.provider_type.as_str(),
1652 &model_with_provider.model,
1653 )
1654 .await?
1655 && checkpoint.is_compatible(
1656 model_with_provider.provider_type.as_str(),
1657 &model_with_provider.model,
1658 )
1659 {
1660 let filters = crate::capabilities::collect_message_filters_only(
1661 &resolved_capability_configs,
1662 &self.capability_registry,
1663 );
1664 let mut query =
1665 crate::MessageQuery::new(session_id).after_sequence(checkpoint.source_sequence);
1666 filters.apply_message_filters(&mut query);
1667 let history = self.message_retriever.load_filtered_history(query).await?;
1668 messages = history.messages;
1669 checkpoint_suffix_message_count = messages.len();
1670 filters.apply_post_load_filters(&mut messages);
1671 if let crate::CompactionCheckpointPayload::Summary { text } = &checkpoint.payload {
1672 messages.insert(
1673 0,
1674 Message::system(format!(
1675 "[CONVERSATION_SUMMARY]\n{text}\n[/CONVERSATION_SUMMARY]"
1676 )),
1677 );
1678 }
1679 message_source_sequence = history.source_sequence.or(message_source_sequence);
1680 restored_checkpoint = Some(checkpoint);
1681 }
1682
1683 let handle_effort = self
1694 .reasoning_effort_handle
1695 .as_ref()
1696 .and_then(|handle| handle.get());
1697 let raw_reasoning_effort = handle_effort.or_else(|| {
1698 messages
1699 .iter()
1700 .rev()
1701 .find(|m| m.role == MessageRole::User)
1702 .and_then(|m| m.controls.as_ref())
1703 .and_then(|c| c.reasoning.as_ref())
1704 .and_then(|r| r.effort.clone())
1705 });
1706 let reasoning_effort = raw_reasoning_effort.filter(|effort| {
1707 if effort.eq_ignore_ascii_case("none") {
1709 return false;
1710 }
1711 let profile = crate::model_profiles::get_model_profile(
1714 &model_with_provider.provider_type,
1715 &model_with_provider.model,
1716 );
1717 match profile {
1718 Some(p) if !p.reasoning => {
1719 tracing::warn!(
1720 model = %model_with_provider.model,
1721 effort = %effort,
1722 "Stripping reasoning_effort: model does not support reasoning"
1723 );
1724 false
1725 }
1726 _ => true,
1727 }
1728 });
1729
1730 let speed = messages
1735 .iter()
1736 .rev()
1737 .find(|m| m.role == MessageRole::User)
1738 .and_then(|m| m.controls.as_ref())
1739 .and_then(|c| c.speed.clone())
1740 .filter(|speed| {
1741 let profile = crate::model_profiles::get_model_profile(
1742 &model_with_provider.provider_type,
1743 &model_with_provider.model,
1744 );
1745 let Some(p) = profile else {
1746 return true;
1747 };
1748 let Some(speed_config) = p.speed else {
1749 tracing::warn!(
1750 model = %model_with_provider.model,
1751 speed = %speed,
1752 "Stripping speed: model does not support service tiers"
1753 );
1754 return false;
1755 };
1756 let supported = speed_config.values.iter().any(|value| {
1757 matches!(
1758 (&value.value, speed.as_str()),
1759 (crate::model::Speed::Flex, "flex")
1760 | (crate::model::Speed::Default, "default")
1761 | (crate::model::Speed::Priority, "priority")
1762 )
1763 });
1764 if !supported {
1765 tracing::warn!(
1766 model = %model_with_provider.model,
1767 speed = %speed,
1768 "Stripping speed: model does not support requested service tier"
1769 );
1770 }
1771 supported
1772 });
1773
1774 let verbosity = messages
1779 .iter()
1780 .rev()
1781 .find(|m| m.role == MessageRole::User)
1782 .and_then(|m| m.controls.as_ref())
1783 .and_then(|c| c.verbosity.clone())
1784 .filter(|verbosity| {
1785 let profile = crate::model_profiles::get_model_profile(
1786 &model_with_provider.provider_type,
1787 &model_with_provider.model,
1788 );
1789 match profile {
1790 Some(p) if p.verbosity.is_none() => {
1791 tracing::warn!(
1792 model = %model_with_provider.model,
1793 verbosity = %verbosity,
1794 "Stripping verbosity: model does not support verbosity control"
1795 );
1796 false
1797 }
1798 _ => true,
1799 }
1800 });
1801
1802 if let Some(ref store) = self.partial_stream_store {
1806 let turn_id_str = context.turn_id.to_string();
1807 match store.get_partial_stream(session_id, &turn_id_str).await {
1808 Ok(Some(partial)) if !partial.accumulated.is_empty() => {
1809 return self
1811 .finalize_partial_stream(
1812 session_id,
1813 context,
1814 partial,
1815 iteration,
1816 runtime_agent.max_iterations,
1817 &runtime_agent.tools,
1818 )
1819 .await;
1820 }
1821 Ok(Some(_)) => {
1822 let recovery_ctx = EventContext::from_atom_context(context);
1825 let _ = self
1826 .event_emitter
1827 .emit(EventRequest::new(
1828 session_id,
1829 recovery_ctx,
1830 ReasonRecoveredData {
1831 turn_id: context.turn_id,
1832 mode: RecoveryMode::Restart,
1833 accumulated_len: 0,
1834 },
1835 ))
1836 .await;
1837 tracing::info!(
1838 session_id = %session_id,
1839 turn_id = %context.turn_id,
1840 "ReasonAtom: partial stream detected with empty accumulated; restarting clean"
1841 );
1842 }
1843 Ok(None) => {} Err(e) => {
1845 tracing::warn!(
1847 session_id = %session_id,
1848 turn_id = %context.turn_id,
1849 error = %e,
1850 "ReasonAtom: partial-stream store error; proceeding with normal execution"
1851 );
1852 }
1853 }
1854 }
1855
1856 let repair_event_context = EventContext::from_atom_context(context);
1860 let patched_messages = repair_dangling_tool_calls(
1861 &messages,
1862 self.durable_tool_result_store.as_deref(),
1863 self.event_emitter.as_ref(),
1864 session_id,
1865 &repair_event_context,
1866 &context.turn_id.to_string(),
1867 )
1868 .await;
1869 let raw_tool_result_bytes = crate::capabilities::total_tool_result_bytes(&patched_messages);
1870
1871 let model_view_providers = crate::capabilities::collect_model_view_providers(
1874 &resolved_capability_configs,
1875 &self.capability_registry,
1876 Some(model_with_provider.model.as_str()),
1877 );
1878 let model_view_context = crate::capabilities::ModelViewContext {
1879 session_id,
1880 prior_usage: prior_usage.as_ref(),
1881 };
1882 let mut context_messages =
1883 model_view_providers.apply_model_view(patched_messages, &model_view_context);
1884 context_messages = crate::tool_call_integrity::retain_complete_message_tool_exchanges(
1885 &context_messages,
1886 stateful_response_continuation || restored_checkpoint.is_some(),
1887 );
1888
1889 let mut volatile_suffix_len = 0usize;
1896 {
1897 let facts_ctx = crate::capabilities::FactsContext::new(session_id);
1898 let dynamic_facts = crate::capabilities::collect_dynamic_facts(
1899 &resolved_capability_configs,
1900 &self.capability_registry,
1901 Some(model_with_provider.model.as_str()),
1902 &facts_ctx,
1903 );
1904 if let Some(block) = crate::capabilities::render_facts_block(&dynamic_facts) {
1905 context_messages.push(Message::user(block));
1906 volatile_suffix_len = 1;
1907 }
1908 }
1909
1910 let resolved_images = self.resolve_images(&context_messages).await;
1915
1916 let mut llm_messages = Vec::new();
1918
1919 let has_system_prompt = !runtime_agent.system_prompt.is_empty();
1921 if has_system_prompt {
1922 llm_messages.push(LlmMessage {
1923 role: LlmMessageRole::System,
1924 content: LlmMessageContent::Text(runtime_agent.system_prompt.clone()),
1925 tool_calls: None,
1926 tool_call_id: None,
1927 phase: None,
1928 thinking: None,
1929 thinking_signature: None,
1930 });
1931 }
1932
1933 let messages_for_event: Vec<Message> = if has_system_prompt {
1935 std::iter::once(Message::system(&runtime_agent.system_prompt))
1936 .chain(context_messages.iter().cloned())
1937 .collect()
1938 } else {
1939 context_messages.clone()
1940 };
1941
1942 let mut stripped_error_count = 0u32;
1948 for msg in &context_messages {
1949 if is_error_placeholder_message(msg) {
1950 stripped_error_count += 1;
1951 continue;
1952 }
1953 let mut llm_msg =
1954 crate::llm_conversions::llm_message_from_message_with_images(msg, &resolved_images);
1955 if msg.role == MessageRole::User
1956 && let Some(ref actor) = msg.external_actor
1957 {
1958 llm_msg.prepend_text_prefix(&format!("[{}] ", actor.display_label()));
1959 }
1960 llm_messages.push(llm_msg);
1961 }
1962 if stripped_error_count > 0 {
1963 tracing::info!(
1964 session_id = %session_id,
1965 stripped_error_count,
1966 "ReasonAtom: stripped error placeholder messages from LLM input"
1967 );
1968 }
1969
1970 llm_messages = crate::tool_call_integrity::retain_complete_llm_tool_exchanges_for_request(
1975 llm_messages,
1976 stateful_response_continuation || restored_checkpoint.is_some(),
1977 );
1978
1979 let mut llm_config_builder =
1981 crate::llm_conversions::llm_call_config_builder_from_agent(&runtime_agent);
1982 if let Some(effort) = reasoning_effort.clone() {
1983 llm_config_builder = llm_config_builder.reasoning_effort(effort);
1984 }
1985 if let Some(speed) = speed {
1986 llm_config_builder = llm_config_builder.speed(speed);
1987 }
1988 if let Some(verbosity) = verbosity {
1989 llm_config_builder = llm_config_builder.verbosity(verbosity);
1990 }
1991
1992 for (k, v) in &embedder_metadata {
1994 llm_config_builder = llm_config_builder.with_metadata(k, v.clone());
1995 }
1996
1997 llm_config_builder = llm_config_builder
2001 .with_metadata("session_id", session_id.to_string())
2002 .with_metadata("harness_id", harness_id.to_string())
2003 .with_metadata("turn_id", context.turn_id.to_string())
2004 .with_metadata("exec_id", context.exec_id.to_string())
2005 .with_metadata("org_id", format!("org_{:032x}", org_id));
2006 if let Some(agent_id) = agent_id {
2007 llm_config_builder = llm_config_builder.with_metadata("agent_id", agent_id.to_string());
2008 }
2009
2010 if let Some(model_id) = &resolved_model_id {
2012 llm_config_builder = llm_config_builder.with_metadata("model_id", model_id.to_string());
2013 }
2014
2015 let mut llm_config = llm_config_builder
2016 .previous_response_id(previous_response_id.clone())
2017 .volatile_suffix_len(volatile_suffix_len)
2018 .build();
2019 if let Some(checkpoint) = restored_checkpoint.as_ref()
2020 && let crate::CompactionCheckpointPayload::ProviderOpaque { context } =
2021 &checkpoint.payload
2022 {
2023 llm_config.previous_response_id = None;
2024 llm_config.provider_opaque_context = Some(context.clone());
2025 }
2026
2027 tracing::debug!(
2028 session_id = %session_id,
2029 turn_id = %context.turn_id,
2030 model = %runtime_agent.model,
2031 message_count = %llm_messages.len(),
2032 "ReasonAtom: calling LLM"
2033 );
2034
2035 let streaming_event_context = EventContext::from_atom_context(context);
2038
2039 let mut armed_guardrails: Vec<ArmedGuardrail> = Vec::new();
2046 for (cap_id, cfg, provider) in &guardrail_providers {
2047 let ctx = OutputGuardrailContext {
2048 system_prompt: &runtime_agent.system_prompt,
2049 config: cfg,
2050 };
2051 let guardrail_id = provider.id().to_string();
2052 if let Some(run) = provider.arm(&ctx) {
2053 armed_guardrails.push(ArmedGuardrail {
2054 capability_id: (*cap_id).to_string(),
2055 guardrail_id,
2056 run,
2057 });
2058 }
2059 }
2060 let mut tripped: Option<TrippedGuardrail> = None;
2061 let buffer_output_deltas = !post_output_providers.is_empty();
2066 let output_message_id = MessageId::new();
2070 tracing::info!(
2071 session_id = %session_id,
2072 turn_id = %context.turn_id,
2073 "ReasonAtom: emitting output.message.started event"
2074 );
2075 if let Err(e) = self
2076 .event_emitter
2077 .emit(EventRequest::new(
2078 session_id,
2079 streaming_event_context.clone(),
2080 OutputMessageStartedData {
2081 turn_id: context.turn_id,
2082 message_id: output_message_id,
2083 model: Some(runtime_agent.model.clone()),
2084 iteration: Some(iteration),
2085 phase: None,
2088 },
2089 ))
2090 .await
2091 {
2092 tracing::warn!(
2093 session_id = %session_id,
2094 error = %e,
2095 "ReasonAtom: failed to emit output.message.started event"
2096 );
2097 } else {
2098 tracing::info!(
2099 session_id = %session_id,
2100 "ReasonAtom: output.message.started event emitted successfully"
2101 );
2102 }
2103
2104 let thinking_enabled = reasoning_effort.is_some();
2106 if thinking_enabled {
2107 tracing::info!(
2108 session_id = %session_id,
2109 turn_id = %context.turn_id,
2110 "ReasonAtom: emitting reason.thinking.started event"
2111 );
2112 if let Err(e) = self
2113 .event_emitter
2114 .emit(EventRequest::new(
2115 session_id,
2116 streaming_event_context.clone(),
2117 ReasonThinkingStartedData {
2118 turn_id: context.turn_id,
2119 model: Some(runtime_agent.model.clone()),
2120 },
2121 ))
2122 .await
2123 {
2124 tracing::warn!(
2125 session_id = %session_id,
2126 error = %e,
2127 "ReasonAtom: failed to emit reason.thinking.started event"
2128 );
2129 } else {
2130 tracing::info!(
2131 session_id = %session_id,
2132 "ReasonAtom: reason.thinking.started event emitted successfully"
2133 );
2134 }
2135 }
2136
2137 let llm_start = Instant::now();
2139
2140 let mut compaction_info: Option<LlmCompactionInfo> = None;
2144 let mut llm_messages_for_call = llm_messages.clone();
2145
2146 if let Some(ref config) = compaction_config {
2151 use crate::capabilities::CompactionStrategy;
2152 let context_window = chat_driver
2153 .effective_context_window(&model_with_provider.model)
2154 .or_else(|| {
2155 crate::model_profiles::get_model_profile(
2156 &model_with_provider.provider_type,
2157 &model_with_provider.model,
2158 )
2159 .and_then(|profile| profile.limits.map(|limits| limits.context as usize))
2160 })
2161 .unwrap_or(128_000);
2162 let checkpoint_rearmed = restored_checkpoint.is_none()
2163 || checkpoint_suffix_message_count >= CHECKPOINT_REARM_MIN_SUFFIX_MESSAGES;
2164 let native_strategy = matches!(
2165 config.strategy,
2166 CompactionStrategy::Auto | CompactionStrategy::Native
2167 );
2168 let durable_source_available =
2169 self.compaction_checkpoint_store.is_some() && message_source_sequence.is_some();
2170 let estimated_tokens_before =
2171 crate::capabilities::estimate_total_tokens(&llm_messages_for_call) as u64;
2172 let native_attempt_rearmed = if let (Some(store), Some(source_sequence)) = (
2173 self.compaction_checkpoint_store.as_ref(),
2174 message_source_sequence,
2175 ) {
2176 match store
2177 .get_proactive_attempt(
2178 session_id,
2179 model_with_provider.provider_type.as_str(),
2180 &model_with_provider.model,
2181 )
2182 .await
2183 {
2184 Ok(attempt) => attempt.is_none_or(|attempt| {
2185 if source_sequence < attempt.source_sequence
2186 || llm_messages_for_call.len() < attempt.input_message_count
2187 {
2188 return true;
2189 }
2190 let same_source_lineage = proactive_source_fingerprint(
2191 llm_config.provider_opaque_context.as_ref(),
2192 &llm_messages_for_call[..attempt.input_message_count],
2193 ) == attempt.source_fingerprint;
2194 if !same_source_lineage {
2195 return true;
2196 }
2197 if source_sequence == attempt.source_sequence {
2198 return false;
2199 }
2200 let required_growth = PROACTIVE_RETRY_MIN_TOKEN_GROWTH
2201 .max(attempt.estimated_input_tokens / PROACTIVE_RETRY_GROWTH_DIVISOR);
2202 estimated_tokens_before.saturating_sub(attempt.estimated_input_tokens)
2203 >= required_growth
2204 }),
2205 Err(error) => {
2206 tracing::warn!(
2207 session_id = %session_id,
2208 error = %error,
2209 "ReasonAtom: proactive compaction attempt watermark lookup failed"
2210 );
2211 true
2212 }
2213 }
2214 } else {
2215 true
2216 };
2217 let window_pressure = crate::capabilities::should_compact_proactively(
2218 &llm_messages_for_call,
2219 config,
2220 context_window,
2221 );
2222 let cost_pressure = crate::capabilities::should_compact_for_cost(
2223 estimated_tokens_before as usize,
2224 raw_tool_result_bytes,
2225 config,
2226 prior_usage.as_ref(),
2227 );
2228 let local_pressure = !stateful_response_continuation
2229 && checkpoint_rearmed
2230 && (window_pressure || cost_pressure);
2231 let should_attempt = native_strategy
2232 && chat_driver.supports_compact()
2233 && durable_source_available
2234 && native_attempt_rearmed
2235 && local_pressure;
2236 let mut native_applied = false;
2237
2238 if should_attempt {
2239 use crate::events::{
2240 CompactionReason, CompactionStepData, ContextCompactedData,
2241 ContextCompactingData,
2242 };
2243 let messages_before = llm_messages_for_call.len();
2244 let input_message_count = llm_messages_for_call.len();
2245 let source_fingerprint = proactive_source_fingerprint(
2246 llm_config.provider_opaque_context.as_ref(),
2247 &llm_messages_for_call,
2248 );
2249 if let Err(error) = self
2250 .compaction_checkpoint_store
2251 .as_ref()
2252 .expect("durable source availability checked above")
2253 .record_proactive_attempt(
2254 session_id,
2255 model_with_provider.provider_type.as_str(),
2256 &model_with_provider.model,
2257 crate::ProactiveCompactionAttempt {
2258 source_sequence: message_source_sequence
2259 .expect("durable source availability checked above"),
2260 estimated_input_tokens: estimated_tokens_before,
2261 input_message_count,
2262 source_fingerprint,
2263 },
2264 )
2265 .await
2266 {
2267 tracing::warn!(
2268 session_id = %session_id,
2269 error = %error,
2270 "ReasonAtom: proactive compaction attempt watermark write failed"
2271 );
2272 }
2273 let _ = self
2274 .event_emitter
2275 .emit(EventRequest::new(
2276 session_id,
2277 streaming_event_context.clone(),
2278 ContextCompactingData {
2279 reason: CompactionReason::ProactiveBudget,
2280 strategy: config.strategy.to_string(),
2281 messages_before,
2282 tokens_before: Some(estimated_tokens_before),
2283 bytes_before: None,
2284 },
2285 ))
2286 .await;
2287
2288 if let Some(applied) = try_apply_native_compaction(
2289 chat_driver.as_ref(),
2290 self.compaction_checkpoint_store.as_ref(),
2291 session_id,
2292 message_source_sequence,
2293 model_with_provider.provider_type.as_str(),
2294 &model_with_provider.model,
2295 has_system_prompt.then_some(runtime_agent.system_prompt.as_str()),
2296 false,
2297 &mut llm_messages_for_call,
2298 &mut llm_config,
2299 )
2300 .await?
2301 {
2302 native_applied = true;
2303 compaction_info = Some(LlmCompactionInfo::new(
2304 Some(applied.input_items_before as u32),
2305 applied
2306 .tokens_after
2307 .and_then(|value| u32::try_from(value).ok()),
2308 Some(applied.duration_ms),
2309 ));
2310 let steps = vec![CompactionStepData {
2311 strategy: "native".to_string(),
2312 messages_after: applied.output_items_after,
2313 duration_ms: applied.duration_ms,
2314 }];
2315 let _ = self
2316 .event_emitter
2317 .emit(EventRequest::new(
2318 session_id,
2319 streaming_event_context.clone(),
2320 ContextCompactedData {
2321 checkpoint_id: applied.checkpoint_id,
2322 strategy_used: "native".to_string(),
2323 messages_before,
2324 messages_after: applied.output_items_after,
2325 tokens_before: applied.tokens_before,
2326 tokens_after: applied.tokens_after,
2327 bytes_before: applied.bytes_before,
2328 bytes_after: applied.bytes_after,
2329 duration_ms: applied.duration_ms,
2330 steps,
2331 },
2332 ))
2333 .await;
2334 }
2335 }
2336
2337 if local_pressure && !native_applied {
2341 if matches!(
2342 config.strategy,
2343 CompactionStrategy::Auto | CompactionStrategy::ObservationMasking
2344 ) {
2345 let conversation = if has_system_prompt {
2346 &llm_messages_for_call[1..]
2347 } else {
2348 &llm_messages_for_call[..]
2349 };
2350 let masked = crate::capabilities::apply_observation_masking(
2351 conversation,
2352 &config.observation_masking,
2353 );
2354 if masked.masked_count > 0 {
2355 let mut model_view = Vec::new();
2356 if has_system_prompt {
2357 model_view.push(llm_messages_for_call[0].clone());
2358 }
2359 model_view.extend(masked.messages);
2360 llm_messages_for_call = model_view;
2361 }
2362 }
2363 let budget_tokens = (context_window as f32 * config.budget_percent) as usize;
2364 if crate::capabilities::estimate_total_tokens(&llm_messages_for_call)
2365 > budget_tokens
2366 {
2367 llm_messages_for_call = crate::capabilities::aggressive_trim(
2368 &llm_messages_for_call,
2369 budget_tokens,
2370 has_system_prompt,
2371 );
2372 }
2373 }
2374 }
2375
2376 const DELTA_BATCH_INTERVAL_MS: u64 = 100;
2379 let retry_config = LlmRetryConfig::default();
2380 let mut stream_retry_metadata = RetryMetadata::default();
2381 let mut streamed_phase: Option<crate::message::ExecutionPhase> = None;
2386 let (
2387 text,
2388 thinking,
2389 thinking_signature,
2390 tool_calls,
2391 completion_metadata,
2392 time_to_first_token_ms,
2393 pending_delta,
2394 ) = 'stream_attempt: loop {
2395 let mut stream = match chat_driver
2396 .chat_completion_stream(llm_messages_for_call.clone(), &llm_config)
2397 .await
2398 {
2399 Ok(stream) => stream,
2400 Err(e) if e.is_request_too_large() => {
2401 use crate::capabilities::{CompactionStrategy, apply_observation_masking};
2403 use crate::events::{
2404 CompactionReason, CompactionStepData, ContextCompactedData,
2405 ContextCompactingData,
2406 };
2407
2408 let Some(config) = compaction_config.clone() else {
2409 tracing::warn!(
2410 session_id = %session_id,
2411 turn_id = %context.turn_id,
2412 "ReasonAtom: context too large and compaction capability is not enabled"
2413 );
2414 return Err(e);
2415 };
2416 let messages_before = llm_messages_for_call.len();
2417
2418 tracing::info!(
2419 session_id = %session_id,
2420 turn_id = %context.turn_id,
2421 strategy = %config.strategy,
2422 messages = messages_before,
2423 "ReasonAtom: context too large, attempting compaction"
2424 );
2425
2426 let _ = self
2428 .event_emitter
2429 .emit(EventRequest::new(
2430 session_id,
2431 streaming_event_context.clone(),
2432 ContextCompactingData {
2433 reason: CompactionReason::RequestTooLarge,
2434 strategy: config.strategy.to_string(),
2435 messages_before,
2436 tokens_before: Some(crate::capabilities::estimate_total_tokens(
2437 &llm_messages_for_call,
2438 ) as u64),
2439 bytes_before: None,
2440 },
2441 ))
2442 .await;
2443
2444 let cascade_start = Instant::now();
2445 let mut steps: Vec<CompactionStepData> = Vec::new();
2446 let mut strategies_used: Vec<String> = Vec::new();
2447 let mut checkpoint_id: Option<String> = None;
2448 let mut tokens_before = Some(crate::capabilities::estimate_total_tokens(
2449 &llm_messages_for_call,
2450 ) as u64);
2451 let mut tokens_after = None;
2452 let mut bytes_before = None;
2453 let mut bytes_after = None;
2454
2455 let run_masking = matches!(
2457 config.strategy,
2458 CompactionStrategy::Auto | CompactionStrategy::ObservationMasking
2459 );
2460 let run_native = matches!(
2461 config.strategy,
2462 CompactionStrategy::Auto | CompactionStrategy::Native
2463 ) && chat_driver.supports_compact();
2464 let run_summarization = matches!(
2465 config.strategy,
2466 CompactionStrategy::Auto | CompactionStrategy::Summarization
2467 );
2468
2469 if run_masking {
2471 let step_start = Instant::now();
2472 let conversation_msgs = if has_system_prompt {
2473 &llm_messages_for_call[1..]
2474 } else {
2475 &llm_messages_for_call[..]
2476 };
2477
2478 let masking_result = apply_observation_masking(
2479 conversation_msgs,
2480 &config.observation_masking,
2481 );
2482
2483 if masking_result.masked_count > 0 {
2484 let mut new_messages = Vec::new();
2485 if has_system_prompt {
2486 new_messages.push(llm_messages_for_call[0].clone());
2487 }
2488 new_messages.extend(masking_result.messages);
2489 llm_messages_for_call = new_messages;
2490
2491 let step_duration = step_start.elapsed().as_millis() as u64;
2492 strategies_used.push("observation_masking".to_string());
2493 steps.push(CompactionStepData {
2494 strategy: "observation_masking".to_string(),
2495 messages_after: llm_messages_for_call.len(),
2496 duration_ms: step_duration,
2497 });
2498
2499 tracing::info!(
2500 session_id = %session_id,
2501 masked_count = masking_result.masked_count,
2502 duration_ms = step_duration,
2503 "ReasonAtom: observation masking applied"
2504 );
2505 }
2506 }
2507
2508 if run_native
2510 && let Some(applied) = try_apply_native_compaction(
2511 chat_driver.as_ref(),
2512 self.compaction_checkpoint_store.as_ref(),
2513 session_id,
2514 message_source_sequence,
2515 model_with_provider.provider_type.as_str(),
2516 &model_with_provider.model,
2517 has_system_prompt.then_some(runtime_agent.system_prompt.as_str()),
2518 stateful_response_continuation,
2519 &mut llm_messages_for_call,
2520 &mut llm_config,
2521 )
2522 .await?
2523 {
2524 compaction_info = Some(LlmCompactionInfo::new(
2525 Some(applied.input_items_before as u32),
2526 applied
2527 .tokens_after
2528 .and_then(|value| u32::try_from(value).ok()),
2529 Some(applied.duration_ms),
2530 ));
2531 checkpoint_id = applied.checkpoint_id;
2532 tokens_before = applied.tokens_before;
2533 tokens_after = applied.tokens_after;
2534 bytes_before = applied.bytes_before;
2535 bytes_after = applied.bytes_after;
2536 strategies_used.push("native".to_string());
2537 steps.push(CompactionStepData {
2538 strategy: "native".to_string(),
2539 messages_after: applied.output_items_after,
2540 duration_ms: applied.duration_ms,
2541 });
2542 }
2543
2544 if run_summarization && !strategies_used.contains(&"native".to_string()) {
2547 use crate::capabilities::{
2548 build_summarization_prompt, compose_summary_with_recent,
2549 format_messages_for_summarization,
2550 };
2551
2552 let step_start = Instant::now();
2553 let conversation_msgs = if has_system_prompt {
2554 &llm_messages_for_call[1..]
2555 } else {
2556 &llm_messages_for_call[..]
2557 };
2558
2559 let keep_recent = 10.min(conversation_msgs.len());
2561 let to_summarize =
2562 &conversation_msgs[..conversation_msgs.len() - keep_recent];
2563 let recent = &conversation_msgs[conversation_msgs.len() - keep_recent..];
2564
2565 if !to_summarize.is_empty() {
2566 let summary_prompt = build_summarization_prompt(&config.summarization);
2567 let messages_text = format_messages_for_summarization(to_summarize);
2568
2569 let summary_messages = vec![
2571 LlmMessage {
2572 role: LlmMessageRole::System,
2573 content: LlmMessageContent::Text(summary_prompt),
2574 tool_calls: None,
2575 tool_call_id: None,
2576 phase: None,
2577 thinking: None,
2578 thinking_signature: None,
2579 },
2580 LlmMessage {
2581 role: LlmMessageRole::User,
2582 content: LlmMessageContent::Text(messages_text),
2583 tool_calls: None,
2584 tool_call_id: None,
2585 phase: None,
2586 thinking: None,
2587 thinking_signature: None,
2588 },
2589 ];
2590
2591 let summary_config = crate::driver_registry::LlmCallConfig {
2592 speed: None,
2593 verbosity: None,
2594 model: config
2595 .summarization
2596 .model
2597 .clone()
2598 .unwrap_or_else(|| runtime_agent.model.clone()),
2599 temperature: Some(0.0),
2600 max_tokens: Some(2000),
2601 tools: vec![],
2602 reasoning_effort: None,
2603 metadata: HashMap::new(),
2604 previous_response_id: None,
2605 provider_opaque_context: None,
2606 tool_search: None,
2607 prompt_cache: None,
2608 openrouter_routing: None,
2609 parallel_tool_calls: None,
2610 volatile_suffix_len: 0,
2611 };
2612
2613 match chat_driver
2614 .chat_completion(summary_messages, &summary_config)
2615 .await
2616 {
2617 Ok(response) => {
2618 let summary_text = response.text;
2619 let system_message =
2620 has_system_prompt.then(|| llm_messages_for_call[0].clone());
2621 llm_messages_for_call = compose_summary_with_recent(
2622 system_message,
2623 &summary_text,
2624 recent,
2625 );
2626
2627 let step_duration = step_start.elapsed().as_millis() as u64;
2628 strategies_used.push("summarization".to_string());
2629 steps.push(CompactionStepData {
2630 strategy: "summarization".to_string(),
2631 messages_after: llm_messages_for_call.len(),
2632 duration_ms: step_duration,
2633 });
2634
2635 tracing::info!(
2636 session_id = %session_id,
2637 duration_ms = step_duration,
2638 messages_after = llm_messages_for_call.len(),
2639 "ReasonAtom: summarization applied"
2640 );
2641 }
2642 Err(e) => {
2643 tracing::warn!(
2644 session_id = %session_id,
2645 error = %e,
2646 "ReasonAtom: summarization failed, continuing"
2647 );
2648 }
2649 }
2650 }
2651 }
2652
2653 if strategies_used.is_empty()
2657 || llm_messages_for_call.len() > messages_before / 2
2658 {
2659 use crate::capabilities::aggressive_trim;
2660 let step_start = Instant::now();
2661 let estimated_total =
2663 crate::capabilities::estimate_total_tokens(&llm_messages_for_call);
2664 let target = estimated_total / 2;
2665 let trimmed =
2666 aggressive_trim(&llm_messages_for_call, target, has_system_prompt);
2667 if trimmed.len() < llm_messages_for_call.len() {
2668 llm_messages_for_call = trimmed;
2669 let step_duration = step_start.elapsed().as_millis() as u64;
2670 strategies_used.push("aggressive_trim".to_string());
2671 steps.push(CompactionStepData {
2672 strategy: "aggressive_trim".to_string(),
2673 messages_after: llm_messages_for_call.len(),
2674 duration_ms: step_duration,
2675 });
2676 tracing::info!(
2677 session_id = %session_id,
2678 messages_after = llm_messages_for_call.len(),
2679 "ReasonAtom: aggressive trim applied (last resort)"
2680 );
2681 }
2682 }
2683
2684 let cascade_duration = cascade_start.elapsed().as_millis() as u64;
2685 let messages_after = llm_messages_for_call.len();
2686 if tokens_after.is_none() {
2687 tokens_after = Some(crate::capabilities::estimate_total_tokens(
2688 &llm_messages_for_call,
2689 ) as u64);
2690 }
2691 let effective = match (tokens_before, tokens_after) {
2692 (Some(before), Some(after)) => materially_reduced(before, after),
2693 _ => match (bytes_before, bytes_after) {
2694 (Some(before), Some(after)) => materially_reduced(before, after),
2695 _ => false,
2696 },
2697 };
2698 if !effective {
2699 tracing::warn!(
2700 session_id = %session_id,
2701 ?tokens_before,
2702 ?tokens_after,
2703 "ReasonAtom: compaction cascade made no material reduction"
2704 );
2705 return Err(e);
2706 }
2707
2708 let strategy_used = strategies_used.join("+");
2709 let durable_semantic_compaction = strategies_used
2710 .iter()
2711 .any(|strategy| strategy != "observation_masking");
2712 if durable_semantic_compaction {
2713 let _ = self
2714 .event_emitter
2715 .emit(EventRequest::new(
2716 session_id,
2717 streaming_event_context.clone(),
2718 ContextCompactedData {
2719 checkpoint_id,
2720 strategy_used: strategy_used.clone(),
2721 messages_before,
2722 messages_after,
2723 tokens_before,
2724 tokens_after,
2725 bytes_before,
2726 bytes_after,
2727 duration_ms: cascade_duration,
2728 steps,
2729 },
2730 ))
2731 .await;
2732 }
2733
2734 tracing::info!(
2735 session_id = %session_id,
2736 strategy = %strategy_used,
2737 messages_before,
2738 messages_after,
2739 duration_ms = cascade_duration,
2740 "ReasonAtom: compaction cascade completed, retrying LLM call"
2741 );
2742
2743 chat_driver
2744 .chat_completion_stream(llm_messages_for_call.clone(), &llm_config)
2745 .await?
2746 }
2747 Err(e) => return Err(e),
2748 };
2749
2750 let mut text = String::new();
2751 let mut thinking = String::new();
2752 let mut thinking_signature: Option<String> = None;
2753 let mut tool_calls = Vec::new();
2754 let mut completion_metadata: Option<LlmCompletionMetadata> = None;
2755 let mut stream_has_output = false;
2756 let mut pending_delta = String::new();
2757 let mut pending_thinking_delta = String::new();
2758 let mut last_delta_emit = Instant::now();
2759 let mut last_thinking_delta_emit = Instant::now();
2760 let mut time_to_first_token_ms: Option<u64> = None;
2761
2762 let stall_timeout = self
2764 .provider_stall_timeout
2765 .unwrap_or(std::time::Duration::from_secs(120));
2766 let mut stall_sleep = Box::pin(tokio::time::sleep(stall_timeout));
2767 let mut keepalive_ticker = tokio::time::interval(std::time::Duration::from_secs(12));
2768 keepalive_ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
2769 keepalive_ticker.tick().await; let mut last_stream_heartbeat = Instant::now();
2771 let mut last_token_at_unix: u64 = unix_now_secs();
2776
2777 loop {
2778 let event = tokio::select! {
2779 biased;
2780 next = stream.next() => match next {
2781 Some(e) => e,
2782 None => break,
2783 },
2784 _ = &mut stall_sleep => {
2785 let stall_error =
2795 crate::driver_registry::LlmStreamError::new(format!(
2796 "provider stream stall: no tokens for {}s",
2797 stall_timeout.as_secs()
2798 ));
2799 tracing::warn!(
2800 session_id = %session_id,
2801 turn_id = %context.turn_id,
2802 stall_secs = stall_timeout.as_secs(),
2803 "ReasonAtom: provider stream stall timeout"
2804 );
2805 if should_retry_stream_error(
2806 &stall_error,
2807 stream_retry_metadata.attempts,
2808 retry_config.max_retries,
2809 stream_has_output,
2810 ) {
2811 let wait_duration = retry_config
2812 .calculate_backoff(stream_retry_metadata.attempts);
2813 tracing::warn!(
2814 session_id = %session_id,
2815 turn_id = %context.turn_id,
2816 attempt = stream_retry_metadata.attempts + 1,
2817 max_retries = retry_config.max_retries,
2818 wait_secs = wait_duration.as_secs_f64(),
2819 "ReasonAtom: provider stream stall, retrying"
2820 );
2821 stream_retry_metadata.record_retry(wait_duration, None);
2822 tokio::time::sleep(wait_duration).await;
2823 continue 'stream_attempt;
2824 }
2825 return Err(AgentLoopError::llm(stall_error.message));
2826 },
2827 _ = keepalive_ticker.tick() => {
2828 if let Some(ref hb) = self.stream_heartbeater {
2829 hb.heartbeat(crate::traits::StreamProgress {
2830 accumulated_len: text.len() + thinking.len(),
2831 last_delta_at: last_token_at_unix,
2832 })
2833 .await;
2834 last_stream_heartbeat = Instant::now();
2835 }
2836 continue;
2837 },
2838 };
2839 let event = event?;
2840 let advances_stall_deadline = stream_event_advances_stall_deadline(&event);
2841 if advances_stall_deadline {
2842 stall_sleep
2843 .as_mut()
2844 .reset(tokio::time::Instant::now() + stall_timeout);
2845 last_token_at_unix = unix_now_secs();
2846 }
2847 match event {
2848 LlmStreamEvent::TextDelta(delta) => {
2849 if delta.is_empty() {
2850 continue;
2851 }
2852 stream_has_output = true;
2853 if time_to_first_token_ms.is_none() {
2855 let ttft = llm_start.elapsed().as_millis() as u64;
2856 time_to_first_token_ms = Some(ttft);
2857 tracing::info!(
2858 session_id = %session_id,
2859 time_to_first_token_ms = ttft,
2860 "ReasonAtom: received first token from LLM"
2861 );
2862 }
2863 text.push_str(&delta);
2864 pending_delta.push_str(&delta);
2865
2866 if !armed_guardrails.is_empty()
2873 && let Some(t) =
2874 evaluate_guardrails(&mut armed_guardrails, &text, &delta)
2875 {
2876 tracing::warn!(
2877 session_id = %session_id,
2878 turn_id = %context.turn_id,
2879 guardrail_capability_id = %t.capability_id,
2880 guardrail_id = %t.guardrail_id,
2881 reason_code = %t.block.reason_code,
2882 "ReasonAtom: output guardrail tripped, replacing assistant message"
2883 );
2884 pending_delta.clear();
2885 tripped = Some(t);
2886 break;
2887 }
2888
2889 if !buffer_output_deltas
2891 && last_delta_emit.elapsed().as_millis() as u64
2892 >= DELTA_BATCH_INTERVAL_MS
2893 && !pending_delta.is_empty()
2894 {
2895 if let Err(e) = self
2896 .event_emitter
2897 .emit(EventRequest::new(
2898 session_id,
2899 streaming_event_context.clone(),
2900 OutputMessageDeltaData {
2901 turn_id: context.turn_id,
2902 message_id: output_message_id,
2903 delta: pending_delta.clone(),
2904 accumulated: text.clone(),
2905 phase: streamed_phase,
2906 },
2907 ))
2908 .await
2909 {
2910 tracing::warn!(
2911 session_id = %session_id,
2912 error = %e,
2913 "ReasonAtom: failed to emit output.message.delta event"
2914 );
2915 }
2916 pending_delta.clear();
2917 last_delta_emit = Instant::now();
2918 }
2919 }
2920 LlmStreamEvent::ThinkingDelta(delta) => {
2921 if delta.is_empty() {
2922 continue;
2923 }
2924 stream_has_output = true;
2925 if let Some(t) = append_guarded_thinking_delta(
2926 &mut armed_guardrails,
2927 &mut thinking,
2928 &mut pending_thinking_delta,
2929 &delta,
2930 ) {
2931 tracing::warn!(
2932 session_id = %session_id,
2933 guardrail_capability_id = %t.capability_id,
2934 guardrail_id = %t.guardrail_id,
2935 "ReasonAtom: output guardrail tripped on thinking stream, replacing assistant message"
2936 );
2937 tripped = Some(t);
2938 break;
2939 }
2940 tracing::debug!(
2941 session_id = %session_id,
2942 delta_len = delta.len(),
2943 total_thinking_len = thinking.len(),
2944 "ReasonAtom: received ThinkingDelta from LLM"
2945 );
2946
2947 if last_thinking_delta_emit.elapsed().as_millis() as u64
2949 >= DELTA_BATCH_INTERVAL_MS
2950 && !pending_thinking_delta.is_empty()
2951 {
2952 if let Err(e) = self
2953 .event_emitter
2954 .emit(EventRequest::new(
2955 session_id,
2956 streaming_event_context.clone(),
2957 ReasonThinkingDeltaData {
2958 turn_id: context.turn_id,
2959 delta: pending_thinking_delta.clone(),
2960 accumulated: thinking.clone(),
2961 },
2962 ))
2963 .await
2964 {
2965 tracing::warn!(
2966 session_id = %session_id,
2967 error = %e,
2968 "ReasonAtom: failed to emit reason.thinking.delta event"
2969 );
2970 }
2971 pending_thinking_delta.clear();
2972 last_thinking_delta_emit = Instant::now();
2973 }
2974 }
2975 LlmStreamEvent::ThinkingSignature(signature) => {
2976 stream_has_output = true;
2977 tracing::debug!(
2979 session_id = %session_id,
2980 signature_len = signature.len(),
2981 "ReasonAtom: received ThinkingSignature from LLM"
2982 );
2983 thinking_signature = Some(signature);
2984 }
2985 LlmStreamEvent::ReasonItem {
2986 provider,
2987 model,
2988 item_id,
2989 encrypted_content,
2990 summary,
2991 token_count,
2992 } => {
2993 stream_has_output = true;
2994 if let Some(sig) = encrypted_content.as_ref() {
3000 tracing::debug!(
3001 session_id = %session_id,
3002 signature_len = sig.len(),
3003 provider = %provider,
3004 item_id = %item_id,
3005 "ReasonAtom: captured encrypted reasoning content from ReasonItem"
3006 );
3007 thinking_signature = Some(sig.clone());
3008 }
3009 if let Err(e) = self
3010 .event_emitter
3011 .emit(EventRequest::new(
3012 session_id,
3013 streaming_event_context.clone(),
3014 ReasonItemData {
3015 turn_id: context.turn_id,
3016 provider,
3017 model,
3018 item_id,
3019 encrypted_content,
3020 summary,
3021 token_count,
3022 },
3023 ))
3024 .await
3025 {
3026 tracing::warn!(
3027 session_id = %session_id,
3028 error = %e,
3029 "ReasonAtom: failed to emit reason.item event"
3030 );
3031 }
3032 }
3033 LlmStreamEvent::ToolCalls(calls) => {
3034 stream_has_output |= !calls.is_empty();
3035 tool_calls = calls;
3036 }
3037 LlmStreamEvent::MessagePhase(phase) => {
3038 streamed_phase = crate::message::ExecutionPhase::refine_streamed_hint(
3047 streamed_phase,
3048 phase,
3049 );
3050 }
3051 LlmStreamEvent::Done(metadata) => {
3052 if !buffer_output_deltas
3056 && !pending_delta.is_empty()
3057 && let Err(e) = self
3058 .event_emitter
3059 .emit(EventRequest::new(
3060 session_id,
3061 streaming_event_context.clone(),
3062 OutputMessageDeltaData {
3063 turn_id: context.turn_id,
3064 message_id: output_message_id,
3065 delta: pending_delta.clone(),
3066 accumulated: text.clone(),
3067 phase: streamed_phase,
3068 },
3069 ))
3070 .await
3071 {
3072 tracing::warn!(
3073 session_id = %session_id,
3074 error = %e,
3075 "ReasonAtom: failed to emit final output.message.delta event"
3076 );
3077 }
3078
3079 if !pending_thinking_delta.is_empty()
3081 && let Err(e) = self
3082 .event_emitter
3083 .emit(EventRequest::new(
3084 session_id,
3085 streaming_event_context.clone(),
3086 ReasonThinkingDeltaData {
3087 turn_id: context.turn_id,
3088 delta: pending_thinking_delta.clone(),
3089 accumulated: thinking.clone(),
3090 },
3091 ))
3092 .await
3093 {
3094 tracing::warn!(
3095 session_id = %session_id,
3096 error = %e,
3097 "ReasonAtom: failed to emit final reason.thinking.delta event"
3098 );
3099 }
3100
3101 if !thinking.is_empty()
3103 && let Err(e) = self
3104 .event_emitter
3105 .emit(EventRequest::new(
3106 session_id,
3107 streaming_event_context.clone(),
3108 ReasonThinkingCompletedData {
3109 turn_id: context.turn_id,
3110 thinking: thinking.clone(),
3111 },
3112 ))
3113 .await
3114 {
3115 tracing::warn!(
3116 session_id = %session_id,
3117 error = %e,
3118 "ReasonAtom: failed to emit reason.thinking.completed event"
3119 );
3120 }
3121 completion_metadata = Some(*metadata);
3122 break;
3123 }
3124 LlmStreamEvent::Error(err) => {
3125 let has_partial_output = !tool_calls.is_empty() || !text.is_empty();
3130
3131 if has_partial_output {
3132 tracing::warn!(
3133 session_id = %session_id,
3134 error = %err,
3135 tool_call_count = tool_calls.len(),
3136 text_len = text.len(),
3137 "ReasonAtom: trailing stream error after valid output — treating as partial success"
3138 );
3139 break;
3143 }
3144
3145 if should_retry_stream_error(
3146 &err,
3147 stream_retry_metadata.attempts,
3148 retry_config.max_retries,
3149 stream_has_output,
3150 ) {
3151 let wait_duration =
3152 retry_config.calculate_backoff(stream_retry_metadata.attempts);
3153 tracing::warn!(
3154 session_id = %session_id,
3155 turn_id = %context.turn_id,
3156 attempt = stream_retry_metadata.attempts + 1,
3157 max_retries = retry_config.max_retries,
3158 wait_secs = wait_duration.as_secs_f64(),
3159 error_code = err.code.as_deref().unwrap_or("none"),
3160 error_status = err.status,
3161 error = %err,
3162 "ReasonAtom: transient stream error before output, retrying"
3163 );
3164 stream_retry_metadata.record_retry(wait_duration, None);
3165 tokio::time::sleep(wait_duration).await;
3166 continue 'stream_attempt;
3167 }
3168
3169 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
3171 let event_context = EventContext::from_atom_context(context).with_span(
3172 trace_id.to_string(),
3173 Uuid::now_v7().to_string(),
3174 Some(reason_span_id.to_string()),
3175 );
3176 let tools_summary: Vec<ToolDefinitionSummary> =
3177 runtime_agent.tools.iter().map(|t| t.into()).collect();
3178 let generation_data = LlmGenerationData::failure(
3179 messages_for_event.clone(),
3180 tools_summary,
3181 runtime_agent.model.clone(),
3182 Some(model_with_provider.provider_type.to_string()),
3183 err.to_string(),
3184 Some(llm_duration_ms),
3185 time_to_first_token_ms,
3186 );
3187 let _ = self
3188 .event_emitter
3189 .emit(EventRequest::new(
3190 session_id,
3191 event_context,
3192 generation_data,
3193 ))
3194 .await;
3195 return Err(AgentLoopError::llm_kind(err.kind(), err.to_string()));
3196 }
3197 }
3198 if last_stream_heartbeat.elapsed().as_millis() as u64 >= 5_000
3201 && let Some(ref hb) = self.stream_heartbeater
3202 {
3203 hb.heartbeat(crate::traits::StreamProgress {
3204 accumulated_len: text.len() + thinking.len(),
3205 last_delta_at: last_token_at_unix,
3206 })
3207 .await;
3208 last_stream_heartbeat = Instant::now();
3209 }
3210 }
3211 if let Some(metadata) = completion_metadata.as_mut() {
3212 metadata.retry_metadata =
3213 merge_retry_metadata(metadata.retry_metadata.take(), &stream_retry_metadata);
3214 }
3215
3216 break 'stream_attempt (
3217 text,
3218 thinking,
3219 thinking_signature,
3220 tool_calls,
3221 completion_metadata,
3222 time_to_first_token_ms,
3223 pending_delta,
3224 );
3225 };
3226 let (mut text, mut thinking, thinking_signature, mut tool_calls) =
3227 (text, thinking, thinking_signature, tool_calls);
3228
3229 let mut citation_annotations: Vec<crate::message::TextAnnotation> = Vec::new();
3240 if tripped.is_none()
3241 && !annotation_providers.is_empty()
3242 && !text.is_empty()
3243 && tool_calls.is_empty()
3244 {
3245 text = crate::capabilities::strip_leading_timestamp_annotations(&text);
3249 let collected = collect_annotations(
3250 &annotation_providers,
3251 &runtime_agent.system_prompt,
3252 &text,
3253 &messages,
3254 self.utility_llm_service.as_ref(),
3255 )
3256 .await;
3257 text = collected.text;
3258 citation_annotations = collected.annotations;
3259
3260 if !citation_annotations.is_empty() && !post_output_providers.is_empty() {
3263 let guarded_output = post_generation_guardrail_text(&text, &citation_annotations);
3264 let ctx = PostGenerationOutputContext {
3265 system_prompt: &runtime_agent.system_prompt,
3266 message_text: &guarded_output,
3267 utility_llm_service: self.utility_llm_service.as_ref(),
3268 };
3269 tripped = evaluate_post_generation_guardrails(&post_output_providers, &ctx).await;
3270 }
3271
3272 if tripped.is_none()
3275 && !citation_annotations.is_empty()
3276 && !citation_verifiers.is_empty()
3277 {
3278 citation_annotations = verify_annotations(
3279 &citation_verifiers,
3280 &text,
3281 self.utility_llm_service.as_ref(),
3282 citation_annotations,
3283 )
3284 .await;
3285 }
3286 }
3287
3288 if tripped.is_none()
3291 && citation_annotations.is_empty()
3292 && !post_output_providers.is_empty()
3293 && !text.is_empty()
3294 {
3295 let ctx = PostGenerationOutputContext {
3296 system_prompt: &runtime_agent.system_prompt,
3297 message_text: &text,
3298 utility_llm_service: self.utility_llm_service.as_ref(),
3299 };
3300 tripped = evaluate_post_generation_guardrails(&post_output_providers, &ctx).await;
3301 }
3302
3303 if tripped.is_some() {
3304 citation_annotations.clear();
3305 }
3306
3307 if buffer_output_deltas
3310 && tripped.is_none()
3311 && !pending_delta.is_empty()
3312 && let Err(e) = self
3313 .event_emitter
3314 .emit(EventRequest::new(
3315 session_id,
3316 streaming_event_context.clone(),
3317 OutputMessageDeltaData {
3318 turn_id: context.turn_id,
3319 message_id: output_message_id,
3320 delta: pending_delta.clone(),
3321 accumulated: text.clone(),
3322 phase: streamed_phase,
3323 },
3324 ))
3325 .await
3326 {
3327 tracing::warn!(
3328 session_id = %session_id,
3329 error = %e,
3330 "ReasonAtom: failed to emit guarded output.message.delta event"
3331 );
3332 }
3333
3334 if let Some(ref t) = tripped {
3340 let replaced_event_context = EventContext::from_atom_context(context).with_span(
3341 trace_id.to_string(),
3342 Uuid::now_v7().to_string(),
3343 Some(reason_span_id.to_string()),
3344 );
3345 if let Err(e) = self
3346 .event_emitter
3347 .emit(EventRequest::new(
3348 session_id,
3349 replaced_event_context,
3350 OutputMessageReplacedData {
3351 turn_id: context.turn_id,
3352 message_id: output_message_id,
3353 guardrail_capability_id: t.capability_id.clone(),
3354 guardrail_id: t.guardrail_id.clone(),
3355 reason_code: t.block.reason_code.clone(),
3356 replacement: t.block.replacement.clone(),
3357 },
3358 ))
3359 .await
3360 {
3361 tracing::warn!(
3362 session_id = %session_id,
3363 error = %e,
3364 "ReasonAtom: failed to emit output.message.replaced event"
3365 );
3366 }
3367 text = t.block.replacement.clone();
3368 tool_calls.clear();
3369 thinking.clear();
3370 }
3371
3372 if !tool_calls.is_empty() {
3380 self.repair_malformed_tool_calls(
3381 session_id,
3382 context,
3383 &resolved_capability_configs,
3384 &runtime_agent.tools,
3385 &mut tool_calls,
3386 iteration,
3387 )
3388 .await;
3389 }
3390
3391 let llm_duration_ms = llm_start.elapsed().as_millis() as u64;
3392
3393 let response_id = completion_metadata
3395 .as_ref()
3396 .and_then(|meta| meta.response_id.clone());
3397 let finish_reason = completion_metadata
3398 .as_ref()
3399 .and_then(|meta| meta.finish_reason.clone());
3400
3401 let usage = completion_metadata.as_ref().and_then(|meta| {
3409 match (meta.prompt_tokens, meta.completion_tokens) {
3410 (Some(input), Some(output)) => {
3411 let actual_cost_usd = meta.provider_cost_usd;
3412 let estimated_cost_usd = crate::model_profiles::estimate_cost_usd(
3413 &model_with_provider.provider_type,
3414 &runtime_agent.model,
3415 input,
3416 output,
3417 meta.cache_read_tokens.unwrap_or(0),
3418 meta.cache_creation_tokens.unwrap_or(0),
3419 );
3420 Some(
3421 TokenUsage::with_cache(
3422 input,
3423 output,
3424 meta.cache_read_tokens,
3425 meta.cache_creation_tokens,
3426 )
3427 .with_cost(actual_cost_usd, estimated_cost_usd),
3428 )
3429 }
3430 _ => None,
3431 }
3432 });
3433
3434 let event_context = EventContext::from_atom_context(context).with_span(
3436 trace_id.to_string(),
3437 Uuid::now_v7().to_string(),
3438 Some(reason_span_id.to_string()),
3439 );
3440 let tools_summary: Vec<ToolDefinitionSummary> =
3441 runtime_agent.tools.iter().map(|t| t.into()).collect();
3442 let finish_reasons = Some(vec![finish_reason.clone().unwrap_or_else(|| {
3443 if tool_calls.is_empty() {
3444 "stop".to_string()
3445 } else {
3446 "tool_calls".to_string()
3447 }
3448 })]);
3449 let retry_info = completion_metadata
3451 .as_ref()
3452 .and_then(|meta| meta.retry_metadata.as_ref())
3453 .filter(|rm| rm.had_retries())
3454 .map(|rm| LlmRetryInfo {
3455 attempts: rm.attempts,
3456 total_wait_ms: rm.total_retry_wait.as_millis() as u64,
3457 });
3458 let mut generation_data = LlmGenerationData::success_with_retry(
3460 messages_for_event.clone(),
3461 tools_summary,
3462 Some(text.clone()).filter(|s| !s.is_empty()),
3463 tool_calls.clone(),
3464 runtime_agent.model.clone(),
3465 Some(model_with_provider.provider_type.to_string()),
3466 usage.clone(),
3467 Some(llm_duration_ms),
3468 time_to_first_token_ms,
3469 finish_reasons,
3470 response_id.clone(),
3471 retry_info,
3472 );
3473
3474 if let Some(info) = compaction_info {
3476 generation_data = generation_data.with_compaction(info);
3477 }
3478
3479 if let Some(request_options) =
3480 build_request_options(&llm_config, &model_with_provider.provider_type.to_string())
3481 {
3482 generation_data = generation_data.with_request_options(request_options);
3483 }
3484
3485 if let Err(e) = self
3486 .event_emitter
3487 .emit(EventRequest::new(
3488 session_id,
3489 event_context,
3490 generation_data,
3491 ))
3492 .await
3493 {
3494 tracing::warn!(
3495 session_id = %session_id,
3496 error = %e,
3497 "ReasonAtom: failed to emit llm.generation event"
3498 );
3499 }
3500
3501 let mut metadata = std::collections::HashMap::new();
3503 metadata.insert(
3504 "model".to_string(),
3505 serde_json::Value::String(runtime_agent.model.clone()),
3506 );
3507 if let Some(ref effort) = reasoning_effort {
3508 metadata.insert(
3509 "reasoning_effort".to_string(),
3510 serde_json::Value::String(effort.clone()),
3511 );
3512 }
3513 metadata.insert(
3520 "provider".to_string(),
3521 serde_json::Value::String(model_with_provider.provider_type.to_string()),
3522 );
3523 if let Some(ref rid) = response_id {
3524 metadata.insert(
3525 "response_id".to_string(),
3526 serde_json::Value::String(rid.clone()),
3527 );
3528 }
3529
3530 let text = crate::capabilities::strip_leading_timestamp_annotations(&text);
3535 let has_tool_calls = !tool_calls.is_empty();
3536 let mut assistant_message = if has_tool_calls {
3537 Message::assistant_with_tools(&text, tool_calls.clone())
3538 } else {
3539 Message::assistant(&text)
3540 }
3541 .with_id(output_message_id);
3542 if !citation_annotations.is_empty() {
3545 for part in assistant_message.content.iter_mut() {
3546 if let crate::message::ContentPart::Text(t) = part {
3547 t.annotations = std::mem::take(&mut citation_annotations);
3548 break;
3549 }
3550 }
3551 }
3552 assistant_message.phase = completion_metadata
3556 .as_ref()
3557 .and_then(|meta| meta.phase.as_deref())
3558 .and_then(crate::message::ExecutionPhase::from_provider_str)
3559 .or_else(|| {
3560 Some(crate::message::ExecutionPhase::from_has_tool_calls(
3561 has_tool_calls,
3562 ))
3563 });
3564 assistant_message.metadata = Some(metadata);
3565 if !thinking.is_empty() {
3568 assistant_message.thinking = Some(thinking.clone());
3569 assistant_message.thinking_signature = thinking_signature.clone();
3570 }
3571 let message_event_context = EventContext::from_atom_context(context).with_span(
3574 trace_id.to_string(),
3575 Uuid::now_v7().to_string(),
3576 Some(reason_span_id.to_string()),
3577 );
3578 let mut output_message_data = OutputMessageCompletedData::new(assistant_message);
3579 if let Some(ref u) = usage {
3580 output_message_data = output_message_data.with_usage(u.clone());
3581 }
3582 self.event_emitter
3583 .emit(EventRequest::new(
3584 session_id,
3585 message_event_context,
3586 output_message_data,
3587 ))
3588 .await?;
3589
3590 tracing::info!(
3591 session_id = %session_id,
3592 turn_id = %context.turn_id,
3593 has_tool_calls = %has_tool_calls,
3594 tool_count = %tool_calls.len(),
3595 "ReasonAtom: LLM call completed"
3596 );
3597
3598 Ok(ReasonResult {
3599 success: true,
3600 text,
3601 tool_calls,
3602 has_tool_calls,
3603 tool_definitions: runtime_agent.tools.clone(),
3604 max_iterations: runtime_agent.max_iterations,
3605 error: None,
3606 user_facing_error: None,
3607 error_disclosure: None,
3608 usage,
3609 output_message_id: Some(output_message_id),
3610 time_to_first_token_ms,
3611 response_id,
3612 finish_reason,
3613 locale: resolved_locale,
3614 network_access: runtime_agent.network_access.clone(),
3615 parallel_tool_calls: runtime_agent.parallel_tool_calls,
3616 })
3617 }
3618
3619 async fn finalize_partial_stream(
3624 &self,
3625 session_id: SessionId,
3626 context: &AtomContext,
3627 partial: PartialStreamState,
3628 iteration: u32,
3629 max_iterations: usize,
3630 tool_definitions: &[ToolDefinition],
3631 ) -> Result<ReasonResult> {
3632 let event_context = EventContext::from_atom_context(context);
3633 let turn_id = context.turn_id;
3634 let message_id = partial.message_id;
3635
3636 let _ = self
3638 .event_emitter
3639 .emit(EventRequest::new(
3640 session_id,
3641 event_context.clone(),
3642 OutputMessageStartedData {
3643 turn_id,
3644 message_id,
3645 model: None,
3646 iteration: Some(iteration),
3647 phase: None,
3650 },
3651 ))
3652 .await;
3653
3654 let accumulated =
3657 crate::capabilities::strip_leading_timestamp_annotations(&partial.accumulated);
3658 let assistant_message = Message::assistant(&accumulated).with_id(message_id);
3659 let output_message_id = message_id;
3660 self.event_emitter
3661 .emit(EventRequest::new(
3662 session_id,
3663 event_context.clone(),
3664 OutputMessageCompletedData::new(assistant_message),
3665 ))
3666 .await?;
3667
3668 let accumulated_len = accumulated.len();
3670 let _ = self
3671 .event_emitter
3672 .emit(EventRequest::new(
3673 session_id,
3674 event_context.clone(),
3675 ReasonRecoveredData {
3676 turn_id,
3677 mode: RecoveryMode::Finalize,
3678 accumulated_len,
3679 },
3680 ))
3681 .await;
3682
3683 tracing::info!(
3684 session_id = %session_id,
3685 turn_id = %turn_id,
3686 accumulated_len,
3687 "ReasonAtom: finalized partial stream from persisted accumulated text"
3688 );
3689
3690 Ok(ReasonResult {
3691 success: true,
3692 text: accumulated,
3693 tool_calls: vec![],
3694 has_tool_calls: false,
3695 tool_definitions: tool_definitions.to_vec(),
3696 max_iterations,
3697 error: None,
3698 user_facing_error: None,
3699 error_disclosure: None,
3700 usage: None,
3701 output_message_id: Some(output_message_id),
3702 time_to_first_token_ms: None,
3703 response_id: None,
3704 finish_reason: Some("stop".to_string()),
3705 locale: None,
3706 network_access: None,
3707 parallel_tool_calls: None,
3709 })
3710 }
3711
3712 fn create_chat_driver(
3715 &self,
3716 model: &ResolvedModel,
3717 ) -> Result<crate::driver_registry::BoxedChatDriver> {
3718 self.driver_registry
3719 .create_chat_driver(&crate::llm_conversions::provider_config_from_resolved_model(model))
3720 }
3721
3722 async fn resolve_images(&self, messages: &[Message]) -> HashMap<Uuid, ResolvedImage> {
3733 let mut resolved = HashMap::new();
3734
3735 let resolver = match &self.image_resolver {
3737 Some(r) => r,
3738 None => return resolved,
3739 };
3740
3741 let image_ids: Vec<Uuid> = messages
3743 .iter()
3744 .flat_map(crate::llm_conversions::extract_image_file_ids)
3745 .collect::<std::collections::HashSet<_>>()
3746 .into_iter()
3747 .collect();
3748
3749 if image_ids.is_empty() {
3750 return resolved;
3751 }
3752
3753 tracing::debug!(
3754 image_count = image_ids.len(),
3755 "ReasonAtom: resolving image_file references"
3756 );
3757
3758 for image_id in image_ids {
3760 match resolver.resolve_image(image_id).await {
3761 Ok(Some(image)) => {
3762 resolved.insert(image_id, image);
3763 }
3764 Ok(None) => {
3765 tracing::warn!(
3766 image_id = %image_id,
3767 "ReasonAtom: image not found during resolution"
3768 );
3769 }
3770 Err(e) => {
3771 tracing::warn!(
3772 image_id = %image_id,
3773 error = %e,
3774 "ReasonAtom: failed to resolve image"
3775 );
3776 }
3777 }
3778 }
3779
3780 tracing::debug!(
3781 resolved_count = resolved.len(),
3782 "ReasonAtom: image resolution complete"
3783 );
3784
3785 resolved
3786 }
3787}
3788
3789#[cfg(test)]
3794mod tests {
3795 use super::*;
3796 use crate::driver_registry::{LlmCallConfig, PromptCacheConfig, PromptCacheStrategy};
3797 use std::collections::HashMap;
3798
3799 #[test]
3800 fn material_reduction_requires_five_percent_at_normal_sizes() {
3801 assert!(!materially_reduced(1_000, 951));
3802 assert!(materially_reduced(1_000, 950));
3803 }
3804
3805 #[test]
3806 fn material_reduction_uses_absolute_floor_for_small_sizes() {
3807 assert!(!materially_reduced(0, 0));
3808 assert!(!materially_reduced(100, 69));
3809 assert!(materially_reduced(100, 68));
3810 }
3811
3812 struct BlockWhenDeltaContains {
3813 needle: &'static str,
3814 }
3815
3816 impl crate::output_guardrail::OutputGuardrailRun for BlockWhenDeltaContains {
3817 fn check(
3818 &mut self,
3819 _accumulated: &str,
3820 delta: &str,
3821 ) -> crate::output_guardrail::GuardrailDecision {
3822 if delta.contains(self.needle) {
3823 crate::output_guardrail::GuardrailDecision::block("test_leak", "[blocked]")
3824 } else {
3825 crate::output_guardrail::GuardrailDecision::Pass
3826 }
3827 }
3828 }
3829
3830 fn test_armed_guardrail() -> ArmedGuardrail {
3831 ArmedGuardrail {
3832 capability_id: "test_capability".to_string(),
3833 guardrail_id: "test_guardrail".to_string(),
3834 run: Box::new(BlockWhenDeltaContains { needle: "secret" }),
3835 }
3836 }
3837
3838 #[test]
3839 fn test_append_guarded_thinking_delta_blocks_before_pending_emit() {
3840 let mut guardrails = vec![test_armed_guardrail()];
3841 let mut thinking = "safe ".to_string();
3842 let mut pending = "safe ".to_string();
3843
3844 let tripped = append_guarded_thinking_delta(
3845 &mut guardrails,
3846 &mut thinking,
3847 &mut pending,
3848 "secret instructions",
3849 )
3850 .expect("thinking delta should trip guardrail");
3851
3852 assert_eq!(tripped.capability_id, "test_capability");
3853 assert_eq!(tripped.guardrail_id, "test_guardrail");
3854 assert_eq!(tripped.block.reason_code, "test_leak");
3855 assert_eq!(tripped.block.replacement, "[blocked]");
3856 assert_eq!(thinking, "safe secret instructions");
3857 assert!(pending.is_empty());
3858 }
3859
3860 #[test]
3861 fn test_append_guarded_thinking_delta_allows_safe_pending_emit() {
3862 let mut guardrails = vec![test_armed_guardrail()];
3863 let mut thinking = String::new();
3864 let mut pending = String::new();
3865
3866 let tripped = append_guarded_thinking_delta(
3867 &mut guardrails,
3868 &mut thinking,
3869 &mut pending,
3870 "ordinary reasoning",
3871 );
3872
3873 assert!(tripped.is_none());
3874 assert_eq!(thinking, "ordinary reasoning");
3875 assert_eq!(pending, "ordinary reasoning");
3876 }
3877
3878 #[test]
3879 fn test_reason_result_default() {
3880 let result = ReasonResult::default();
3881 assert!(!result.success);
3882 assert!(result.text.is_empty());
3883 assert!(result.tool_calls.is_empty());
3884 assert!(!result.has_tool_calls);
3885 assert_eq!(result.max_iterations, 0);
3887 }
3888
3889 #[test]
3890 fn test_reason_result_serde_default() {
3891 let json = r#"{"success":true,"text":"","has_tool_calls":false}"#;
3893 let result: ReasonResult = serde_json::from_str(json).unwrap();
3894 assert_eq!(result.max_iterations, 500);
3895 }
3896
3897 #[test]
3898 fn test_capability_usage_snapshot_keeps_resolved_and_exposed_separate() {
3899 let registry = CapabilityRegistry::with_builtins();
3900 let tool = ToolDefinition::Builtin(crate::tool_types::BuiltinTool {
3901 name: "demo_tool".to_string(),
3902 display_name: None,
3903 description: "demo".to_string(),
3904 parameters: json!({"type": "object"}),
3905 policy: crate::tool_types::ToolPolicy::Auto,
3906 category: None,
3907 deferrable: crate::tool_types::DeferrablePolicy::default(),
3908 hints: crate::tool_types::ToolHints::default(),
3909 full_parameters: None,
3910 })
3911 .with_capability_attribution("cap:demo", Some("Demo Capability"));
3912
3913 let records = capability_usage_snapshot_records(
3914 ®istry,
3915 &[crate::AgentCapabilityConfig::new("current_time")],
3916 &[tool],
3917 );
3918
3919 assert!(records.iter().any(|record| {
3920 matches!(record.usage_kind, CapabilityUsageKind::Resolved)
3921 && record.capability_id == "current_time"
3922 && record.tool_name.is_none()
3923 }));
3924 assert!(records.iter().any(|record| {
3925 matches!(record.usage_kind, CapabilityUsageKind::Exposed)
3926 && record.capability_id == "cap:demo"
3927 && record.tool_name.as_deref() == Some("demo_tool")
3928 }));
3929 }
3930
3931 #[test]
3932 fn stream_stall_deadline_ignores_empty_keepalive_events() {
3933 assert!(!stream_event_advances_stall_deadline(
3934 &LlmStreamEvent::TextDelta(String::new())
3935 ));
3936 assert!(!stream_event_advances_stall_deadline(
3937 &LlmStreamEvent::ThinkingDelta(String::new())
3938 ));
3939 assert!(!stream_event_advances_stall_deadline(
3940 &LlmStreamEvent::ThinkingSignature("signature".to_string())
3941 ));
3942 assert!(!stream_event_advances_stall_deadline(
3943 &LlmStreamEvent::ReasonItem {
3944 provider: "openai".to_string(),
3945 model: None,
3946 item_id: "item_1".to_string(),
3947 encrypted_content: None,
3948 summary: vec![String::new()],
3949 token_count: Some(0),
3950 }
3951 ));
3952 }
3953
3954 #[test]
3955 fn stream_stall_deadline_advances_on_output_progress() {
3956 assert!(stream_event_advances_stall_deadline(
3957 &LlmStreamEvent::TextDelta("hello".to_string())
3958 ));
3959 assert!(stream_event_advances_stall_deadline(
3960 &LlmStreamEvent::ThinkingDelta("thinking".to_string())
3961 ));
3962 assert!(stream_event_advances_stall_deadline(
3963 &LlmStreamEvent::ReasonItem {
3964 provider: "openai".to_string(),
3965 model: Some("gpt-5.4".to_string()),
3966 item_id: "item_1".to_string(),
3967 encrypted_content: Some("encrypted".to_string()),
3968 summary: vec![],
3969 token_count: None,
3970 }
3971 ));
3972 assert!(stream_event_advances_stall_deadline(
3973 &LlmStreamEvent::ReasonItem {
3974 provider: "openai".to_string(),
3975 model: None,
3976 item_id: "item_2".to_string(),
3977 encrypted_content: None,
3978 summary: vec!["summary".to_string()],
3979 token_count: None,
3980 }
3981 ));
3982 assert!(stream_event_advances_stall_deadline(
3983 &LlmStreamEvent::ReasonItem {
3984 provider: "openai".to_string(),
3985 model: None,
3986 item_id: "item_3".to_string(),
3987 encrypted_content: None,
3988 summary: vec![],
3989 token_count: Some(1),
3990 }
3991 ));
3992 assert!(stream_event_advances_stall_deadline(
3993 &LlmStreamEvent::ToolCalls(vec![ToolCall {
3994 id: "call_1".to_string(),
3995 name: "demo".to_string(),
3996 arguments: json!({}),
3997 }])
3998 ));
3999 }
4000
4001 #[tokio::test]
4002 async fn test_repair_dangling_tool_calls_no_tool_calls() {
4003 use crate::events::EventContext;
4004 use crate::typed_id::SessionId;
4005 let messages = vec![Message::user("Hello"), Message::assistant("Hi there!")];
4006 let emitter = crate::traits::NoopEventEmitter;
4007 let session_id = SessionId::new();
4008 let ctx = EventContext::empty();
4009 let patched =
4010 repair_dangling_tool_calls(&messages, None, &emitter, session_id, &ctx, "turn_01")
4011 .await;
4012 assert_eq!(patched.len(), 2);
4013 }
4014
4015 #[tokio::test]
4016 async fn test_repair_dangling_tool_calls_with_result() {
4017 use crate::events::EventContext;
4018 use crate::typed_id::SessionId;
4019 let tool_call = ToolCall {
4020 id: "call_123".to_string(),
4021 name: "get_weather".to_string(),
4022 arguments: serde_json::json!({"city": "NYC"}),
4023 };
4024
4025 let messages = vec![
4026 Message::user("What's the weather?"),
4027 Message::assistant_with_tools("Let me check", vec![tool_call]),
4028 Message::tool_result("call_123", Some(serde_json::json!({"temp": 72})), None),
4029 ];
4030
4031 let emitter = crate::traits::NoopEventEmitter;
4032 let session_id = SessionId::new();
4033 let ctx = EventContext::empty();
4034 let patched =
4035 repair_dangling_tool_calls(&messages, None, &emitter, session_id, &ctx, "turn_01")
4036 .await;
4037 assert_eq!(patched.len(), 3);
4038 }
4039
4040 #[tokio::test]
4041 async fn test_repair_dangling_tool_calls_missing_result_no_store() {
4042 use crate::events::EventContext;
4043 use crate::typed_id::SessionId;
4044 let tool_call = ToolCall {
4045 id: "call_456".to_string(),
4046 name: "search_web".to_string(),
4047 arguments: serde_json::json!({"query": "rust"}),
4048 };
4049
4050 let messages = vec![
4051 Message::user("Search for rust"),
4052 Message::assistant_with_tools("Searching...", vec![tool_call]),
4053 Message::user("Actually, never mind"),
4054 ];
4055
4056 let emitter = crate::traits::NoopEventEmitter;
4057 let session_id = SessionId::new();
4058 let ctx = EventContext::empty();
4059 let patched =
4060 repair_dangling_tool_calls(&messages, None, &emitter, session_id, &ctx, "turn_01")
4061 .await;
4062 assert_eq!(patched.len(), 4);
4064 assert_eq!(patched[2].role, MessageRole::ToolResult);
4065 assert_eq!(patched[2].tool_call_id(), Some("call_456"));
4066 }
4067
4068 #[tokio::test]
4069 async fn test_repair_dangling_tool_calls_settled_result_replayed() {
4070 use crate::events::EventContext;
4071 use crate::traits::{DurableToolCallStatus, DurableToolResultStore, ToolCallClaimResult};
4072 use crate::typed_id::SessionId;
4073
4074 struct MockSettledStore;
4075 #[async_trait::async_trait]
4076 impl DurableToolResultStore for MockSettledStore {
4077 async fn try_claim_tool_call(
4078 &self,
4079 _: &str,
4080 _: &str,
4081 _: &str,
4082 _: &str,
4083 ) -> crate::error::Result<ToolCallClaimResult> {
4084 Ok(ToolCallClaimResult::Claimed {
4085 claim_token: uuid::Uuid::new_v4(),
4086 })
4087 }
4088 async fn settle_tool_call(
4089 &self,
4090 _: &str,
4091 _: &str,
4092 _: serde_json::Value,
4093 _: &str,
4094 _: uuid::Uuid,
4095 ) -> crate::error::Result<bool> {
4096 Ok(true)
4097 }
4098 async fn get_tool_call_status(
4099 &self,
4100 _turn_id: &str,
4101 _tool_call_id: &str,
4102 ) -> crate::error::Result<Option<DurableToolCallStatus>> {
4103 Ok(Some(DurableToolCallStatus::Settled {
4104 result_json: serde_json::json!({
4105 "tool_call_id": "call_789",
4106 "result": {"answer": 42},
4107 "error": null,
4108 "images": null,
4109 "connection_required": null,
4110 "raw_output": null
4111 }),
4112 }))
4113 }
4114 }
4115
4116 let tool_call = ToolCall {
4117 id: "call_789".to_string(),
4118 name: "compute".to_string(),
4119 arguments: serde_json::json!({"x": 21}),
4120 };
4121 let messages = vec![
4122 Message::user("Compute"),
4123 Message::assistant_with_tools("Computing...", vec![tool_call]),
4124 ];
4125
4126 let store = MockSettledStore;
4127 let emitter = crate::traits::NoopEventEmitter;
4128 let session_id = SessionId::new();
4129 let ctx = EventContext::empty();
4130 let patched = repair_dangling_tool_calls(
4131 &messages,
4132 Some(&store as &dyn DurableToolResultStore),
4133 &emitter,
4134 session_id,
4135 &ctx,
4136 "turn_01",
4137 )
4138 .await;
4139
4140 assert_eq!(patched.len(), 3);
4142 assert_eq!(patched[2].role, MessageRole::ToolResult);
4143 assert_eq!(patched[2].tool_call_id(), Some("call_789"));
4144 }
4145
4146 #[tokio::test]
4147 async fn test_repair_dangling_tool_calls_interrupted_result_replayed() {
4148 use crate::events::EventContext;
4149 use crate::traits::{DurableToolCallStatus, DurableToolResultStore, ToolCallClaimResult};
4150 use crate::typed_id::SessionId;
4151
4152 struct MockInterruptedStore;
4153 #[async_trait::async_trait]
4154 impl DurableToolResultStore for MockInterruptedStore {
4155 async fn try_claim_tool_call(
4156 &self,
4157 _: &str,
4158 _: &str,
4159 _: &str,
4160 _: &str,
4161 ) -> crate::error::Result<ToolCallClaimResult> {
4162 Ok(ToolCallClaimResult::Claimed {
4163 claim_token: uuid::Uuid::new_v4(),
4164 })
4165 }
4166 async fn settle_tool_call(
4167 &self,
4168 _: &str,
4169 _: &str,
4170 _: serde_json::Value,
4171 _: &str,
4172 _: uuid::Uuid,
4173 ) -> crate::error::Result<bool> {
4174 Ok(true)
4175 }
4176 async fn get_tool_call_status(
4177 &self,
4178 _turn_id: &str,
4179 _tool_call_id: &str,
4180 ) -> crate::error::Result<Option<DurableToolCallStatus>> {
4181 Ok(Some(DurableToolCallStatus::Interrupted {
4182 result_json: None,
4183 }))
4184 }
4185 }
4186
4187 let tool_call = ToolCall {
4188 id: "call_int".to_string(),
4189 name: "slow_op".to_string(),
4190 arguments: serde_json::json!({}),
4191 };
4192 let messages = vec![
4193 Message::user("Do it"),
4194 Message::assistant_with_tools("Doing...", vec![tool_call]),
4195 ];
4196
4197 let store = MockInterruptedStore;
4198 let emitter = crate::traits::NoopEventEmitter;
4199 let session_id = SessionId::new();
4200 let ctx = EventContext::empty();
4201 let patched = repair_dangling_tool_calls(
4202 &messages,
4203 Some(&store as &dyn DurableToolResultStore),
4204 &emitter,
4205 session_id,
4206 &ctx,
4207 "turn_01",
4208 )
4209 .await;
4210
4211 assert_eq!(patched.len(), 3);
4212 let repair = &patched[2];
4213 assert_eq!(repair.role, MessageRole::ToolResult);
4214 assert_eq!(repair.tool_call_id(), Some("call_int"));
4215 let content = format!("{:?}", repair);
4217 assert!(
4218 content.contains("interrupted") || content.contains("not complete"),
4219 "expected interrupted message, got: {content}"
4220 );
4221 }
4222
4223 #[tokio::test]
4224 async fn test_repair_dangling_tool_calls_running_synthesized() {
4225 use crate::events::EventContext;
4226 use crate::traits::{DurableToolCallStatus, DurableToolResultStore, ToolCallClaimResult};
4227 use crate::typed_id::SessionId;
4228
4229 struct MockRunningStore;
4230 #[async_trait::async_trait]
4231 impl DurableToolResultStore for MockRunningStore {
4232 async fn try_claim_tool_call(
4233 &self,
4234 _: &str,
4235 _: &str,
4236 _: &str,
4237 _: &str,
4238 ) -> crate::error::Result<ToolCallClaimResult> {
4239 Ok(ToolCallClaimResult::Claimed {
4240 claim_token: uuid::Uuid::new_v4(),
4241 })
4242 }
4243 async fn settle_tool_call(
4244 &self,
4245 _: &str,
4246 _: &str,
4247 _: serde_json::Value,
4248 _: &str,
4249 _: uuid::Uuid,
4250 ) -> crate::error::Result<bool> {
4251 Ok(true)
4252 }
4253 async fn get_tool_call_status(
4254 &self,
4255 _turn_id: &str,
4256 _tool_call_id: &str,
4257 ) -> crate::error::Result<Option<DurableToolCallStatus>> {
4258 Ok(Some(DurableToolCallStatus::Running))
4259 }
4260 }
4261
4262 let tool_call = ToolCall {
4263 id: "call_run".to_string(),
4264 name: "long_job".to_string(),
4265 arguments: serde_json::json!({}),
4266 };
4267 let messages = vec![
4268 Message::user("Start job"),
4269 Message::assistant_with_tools("Starting...", vec![tool_call]),
4270 ];
4271
4272 let store = MockRunningStore;
4273 let emitter = crate::traits::NoopEventEmitter;
4274 let session_id = SessionId::new();
4275 let ctx = EventContext::empty();
4276 let patched = repair_dangling_tool_calls(
4277 &messages,
4278 Some(&store as &dyn DurableToolResultStore),
4279 &emitter,
4280 session_id,
4281 &ctx,
4282 "turn_01",
4283 )
4284 .await;
4285
4286 assert_eq!(patched.len(), 3);
4287 let repair = &patched[2];
4288 assert_eq!(repair.role, MessageRole::ToolResult);
4289 assert_eq!(repair.tool_call_id(), Some("call_run"));
4290 let content = format!("{:?}", repair);
4292 assert!(
4293 content.contains("uncertain") || content.contains("do not retry"),
4294 "expected uncertain/do-not-retry message, got: {content}"
4295 );
4296 }
4297
4298 #[tokio::test]
4299 async fn test_repair_dangling_tool_calls_store_error_unknown() {
4300 use crate::error::AgentLoopError;
4301 use crate::events::EventContext;
4302 use crate::traits::{DurableToolResultStore, ToolCallClaimResult};
4303 use crate::typed_id::SessionId;
4304
4305 struct MockErrorStore;
4306 #[async_trait::async_trait]
4307 impl DurableToolResultStore for MockErrorStore {
4308 async fn try_claim_tool_call(
4309 &self,
4310 _: &str,
4311 _: &str,
4312 _: &str,
4313 _: &str,
4314 ) -> crate::error::Result<ToolCallClaimResult> {
4315 Ok(ToolCallClaimResult::Claimed {
4316 claim_token: uuid::Uuid::new_v4(),
4317 })
4318 }
4319 async fn settle_tool_call(
4320 &self,
4321 _: &str,
4322 _: &str,
4323 _: serde_json::Value,
4324 _: &str,
4325 _: uuid::Uuid,
4326 ) -> crate::error::Result<bool> {
4327 Ok(true)
4328 }
4329 async fn get_tool_call_status(
4330 &self,
4331 _turn_id: &str,
4332 _tool_call_id: &str,
4333 ) -> crate::error::Result<Option<crate::traits::DurableToolCallStatus>> {
4334 Err(AgentLoopError::tool("simulated store failure"))
4335 }
4336 }
4337
4338 let tool_call = ToolCall {
4339 id: "call_err".to_string(),
4340 name: "risky_op".to_string(),
4341 arguments: serde_json::json!({}),
4342 };
4343 let messages = vec![
4344 Message::user("Do risky op"),
4345 Message::assistant_with_tools("On it...", vec![tool_call]),
4346 ];
4347
4348 let store = MockErrorStore;
4349 let emitter = crate::traits::NoopEventEmitter;
4350 let session_id = SessionId::new();
4351 let ctx = EventContext::empty();
4352 let patched = repair_dangling_tool_calls(
4353 &messages,
4354 Some(&store as &dyn DurableToolResultStore),
4355 &emitter,
4356 session_id,
4357 &ctx,
4358 "turn_01",
4359 )
4360 .await;
4361
4362 assert_eq!(patched.len(), 3);
4363 let repair = &patched[2];
4364 assert_eq!(repair.role, MessageRole::ToolResult);
4365 assert_eq!(repair.tool_call_id(), Some("call_err"));
4366 let content = format!("{:?}", repair);
4368 assert!(
4369 !content.contains("safe to retry"),
4370 "store error must not say 'safe to retry', got: {content}"
4371 );
4372 assert!(
4373 content.contains("do not retry") || content.contains("status unknown"),
4374 "expected do-not-retry/status-unknown message, got: {content}"
4375 );
4376 }
4377
4378 #[test]
4379 fn test_build_request_options_for_openai_prompt_cache() {
4380 let config = LlmCallConfig {
4381 speed: None,
4382 verbosity: None,
4383 model: "gpt-5.4".to_string(),
4384 temperature: None,
4385 max_tokens: None,
4386 tools: vec![],
4387 reasoning_effort: None,
4388 metadata: HashMap::new(),
4389 previous_response_id: Some("resp_123".to_string()),
4390 provider_opaque_context: None,
4391 tool_search: None,
4392 prompt_cache: Some(PromptCacheConfig {
4393 enabled: true,
4394 strategy: PromptCacheStrategy::Auto,
4395 gemini_cached_content: None,
4396 }),
4397 openrouter_routing: None,
4398 parallel_tool_calls: None,
4399 volatile_suffix_len: 0,
4400 };
4401
4402 let request_options = build_request_options(&config, "openai").unwrap();
4403 assert_eq!(
4404 request_options
4405 .prompt_cache
4406 .and_then(|info| info.provider_mode),
4407 Some("prompt_cache_key".to_string())
4408 );
4409 assert_eq!(
4410 request_options.provider_options.get("openai"),
4411 Some(&json!({ "previous_response_id": true }))
4412 );
4413 }
4414
4415 #[test]
4416 fn test_build_request_options_for_gemini_explicit_cache() {
4417 let config = LlmCallConfig {
4418 speed: None,
4419 verbosity: None,
4420 model: "gemini-2.5-pro".to_string(),
4421 temperature: None,
4422 max_tokens: None,
4423 tools: vec![],
4424 reasoning_effort: None,
4425 metadata: HashMap::new(),
4426 previous_response_id: None,
4427 provider_opaque_context: None,
4428 tool_search: None,
4429 prompt_cache: Some(PromptCacheConfig {
4430 enabled: true,
4431 strategy: PromptCacheStrategy::Auto,
4432 gemini_cached_content: Some("cachedContents/demo-cache".to_string()),
4433 }),
4434 openrouter_routing: None,
4435 parallel_tool_calls: None,
4436 volatile_suffix_len: 0,
4437 };
4438
4439 let request_options = build_request_options(&config, "gemini").unwrap();
4440 assert_eq!(
4441 request_options
4442 .prompt_cache
4443 .and_then(|info| info.provider_mode),
4444 Some("cached_content".to_string())
4445 );
4446 assert_eq!(
4447 request_options.provider_options.get("gemini"),
4448 Some(&json!({ "cached_content": true }))
4449 );
4450 }
4451
4452 #[test]
4453 fn test_build_request_options_omits_gemini_cache_flag_when_disabled() {
4454 let config = LlmCallConfig {
4455 speed: None,
4456 verbosity: None,
4457 model: "gemini-2.5-pro".to_string(),
4458 temperature: None,
4459 max_tokens: None,
4460 tools: vec![],
4461 reasoning_effort: None,
4462 metadata: HashMap::new(),
4463 previous_response_id: None,
4464 provider_opaque_context: None,
4465 tool_search: None,
4466 prompt_cache: Some(PromptCacheConfig {
4467 enabled: false,
4468 strategy: PromptCacheStrategy::Auto,
4469 gemini_cached_content: Some("cachedContents/demo-cache".to_string()),
4470 }),
4471 openrouter_routing: None,
4472 parallel_tool_calls: None,
4473 volatile_suffix_len: 0,
4474 };
4475
4476 assert!(build_request_options(&config, "gemini").is_none());
4477 }
4478
4479 #[test]
4480 fn system_keys_override_embedder_keys_in_metadata() {
4481 let embedder_metadata: HashMap<String, String> = [
4485 ("session_id".to_string(), "attacker_value".to_string()),
4486 ("custom_key".to_string(), "custom_value".to_string()),
4487 ]
4488 .into();
4489
4490 let mut metadata: HashMap<String, String> = HashMap::new();
4491
4492 for (k, v) in &embedder_metadata {
4494 metadata.insert(k.clone(), v.clone());
4495 }
4496
4497 metadata.insert("session_id".to_string(), "real_session_id".to_string());
4499
4500 assert_eq!(
4501 metadata.get("session_id").map(String::as_str),
4502 Some("real_session_id"),
4503 "system key must overwrite embedder key with same name"
4504 );
4505 assert_eq!(
4506 metadata.get("custom_key").map(String::as_str),
4507 Some("custom_value"),
4508 "non-colliding embedder key must be preserved"
4509 );
4510 }
4511
4512 use crate::traits::{NoopPartialStreamStore, PartialStreamState, PartialStreamStore};
4517
4518 struct MockPartialStore(Option<PartialStreamState>);
4519
4520 #[async_trait::async_trait]
4521 impl PartialStreamStore for MockPartialStore {
4522 async fn get_partial_stream(
4523 &self,
4524 _session_id: crate::typed_id::SessionId,
4525 _turn_id: &str,
4526 ) -> crate::error::Result<Option<PartialStreamState>> {
4527 Ok(self.0.clone())
4528 }
4529 }
4530
4531 #[tokio::test]
4532 async fn test_noop_partial_stream_store_returns_none() {
4533 let store = NoopPartialStreamStore;
4534 let result = store
4535 .get_partial_stream(crate::typed_id::SessionId::new(), "turn_01")
4536 .await
4537 .unwrap();
4538 assert!(result.is_none());
4539 }
4540
4541 #[tokio::test]
4542 async fn test_partial_stream_store_returns_accumulated_when_partial_exists() {
4543 let message_id = MessageId::new();
4544 let store = MockPartialStore(Some(PartialStreamState {
4545 message_id,
4546 accumulated: "partial text so far".to_string(),
4547 }));
4548 let result = store
4549 .get_partial_stream(crate::typed_id::SessionId::new(), "turn_01")
4550 .await
4551 .unwrap();
4552 let partial = result.unwrap();
4553 assert_eq!(partial.message_id, message_id);
4554 assert_eq!(partial.accumulated, "partial text so far");
4555 }
4556
4557 #[tokio::test]
4558 async fn test_partial_stream_store_returns_empty_when_started_no_delta() {
4559 let store = MockPartialStore(Some(PartialStreamState {
4560 message_id: MessageId::new(),
4561 accumulated: String::new(),
4562 }));
4563 let result = store
4564 .get_partial_stream(crate::typed_id::SessionId::new(), "turn_01")
4565 .await
4566 .unwrap();
4567 assert!(result.unwrap().accumulated.is_empty());
4568 }
4569
4570 mod tool_call_repair_integration {
4575 use super::*;
4576 use crate::capabilities::{CapabilityRegistry, TOOL_CALL_REPAIR_CAPABILITY_ID};
4577 use crate::events::{EventData, EventRequest};
4578 use crate::tool_types::{BuiltinTool, ToolDefinition, ToolPolicy};
4579 use crate::traits::EventEmitter;
4580 use crate::typed_id::{MessageId, SessionId, TurnId};
4581 use std::sync::Mutex;
4582
4583 #[derive(Default)]
4585 struct RecordingEmitter {
4586 events: Mutex<Vec<EventRequest>>,
4587 }
4588
4589 #[async_trait::async_trait]
4590 impl EventEmitter for RecordingEmitter {
4591 async fn emit(
4592 &self,
4593 request: EventRequest,
4594 ) -> crate::error::Result<crate::events::Event> {
4595 let event = request
4596 .clone()
4597 .into_event(crate::typed_id::EventId::new(), 0);
4598 self.events.lock().unwrap().push(request);
4599 Ok(event)
4600 }
4601 }
4602
4603 impl RecordingEmitter {
4604 fn repaired_events(&self) -> Vec<(String, String)> {
4605 self.events
4606 .lock()
4607 .unwrap()
4608 .iter()
4609 .filter_map(|r| match &r.data {
4610 EventData::ToolCallRepaired(d) => {
4611 Some((d.tool_call_id.clone(), d.outcome.clone()))
4612 }
4613 _ => None,
4614 })
4615 .collect()
4616 }
4617 }
4618
4619 fn ctx() -> AtomContext {
4620 AtomContext::new(SessionId::new(), TurnId::new(), MessageId::new())
4621 }
4622
4623 fn read_file_tool() -> ToolDefinition {
4624 ToolDefinition::Builtin(BuiltinTool {
4625 name: "read_file".to_string(),
4626 display_name: None,
4627 description: "read".to_string(),
4628 parameters: json!({
4629 "type": "object",
4630 "properties": { "path": { "type": "string" } },
4631 "required": ["path"]
4632 }),
4633 policy: ToolPolicy::Auto,
4634 category: None,
4635 deferrable: Default::default(),
4636 hints: Default::default(),
4637 full_parameters: None,
4638 })
4639 }
4640
4641 fn malformed_call() -> ToolCall {
4642 ToolCall {
4645 id: "call_1".to_string(),
4646 name: "read_file".to_string(),
4647 arguments: json!("here you go: {'path': '/foo'}"),
4648 }
4649 }
4650
4651 fn enabled_configs() -> Vec<crate::AgentCapabilityConfig> {
4652 vec![crate::AgentCapabilityConfig::new(
4653 TOOL_CALL_REPAIR_CAPABILITY_ID,
4654 )]
4655 }
4656
4657 async fn run(
4658 configs: &[crate::AgentCapabilityConfig],
4659 tools: &[ToolDefinition],
4660 calls: &mut [ToolCall],
4661 emitter: &RecordingEmitter,
4662 ) {
4663 run_at(configs, tools, calls, emitter, 1).await;
4664 }
4665
4666 async fn run_at(
4667 configs: &[crate::AgentCapabilityConfig],
4668 tools: &[ToolDefinition],
4669 calls: &mut [ToolCall],
4670 emitter: &RecordingEmitter,
4671 iteration: u32,
4672 ) {
4673 let registry = CapabilityRegistry::with_builtins();
4674 apply_tool_call_repair(
4675 ®istry,
4676 emitter,
4677 SessionId::new(),
4678 &ctx(),
4679 configs,
4680 tools,
4681 calls,
4682 iteration,
4683 )
4684 .await;
4685 }
4686
4687 #[tokio::test]
4688 async fn malformed_call_is_repaired_and_turn_proceeds() {
4689 let emitter = RecordingEmitter::default();
4690 let tools = vec![read_file_tool()];
4691 let mut calls = vec![malformed_call()];
4692
4693 run(&enabled_configs(), &tools, &mut calls, &emitter).await;
4694
4695 assert_eq!(calls[0].arguments, json!({ "path": "/foo" }));
4697 assert_eq!(
4698 emitter.repaired_events(),
4699 vec![("call_1".to_string(), "local-salvage".to_string())]
4700 );
4701 }
4702
4703 #[tokio::test]
4704 async fn observability_event_fires_with_outcome_label() {
4705 let emitter = RecordingEmitter::default();
4706 let tools = vec![read_file_tool()];
4707 let mut calls = vec![malformed_call()];
4708
4709 run(&enabled_configs(), &tools, &mut calls, &emitter).await;
4710
4711 let repaired = emitter.repaired_events();
4712 assert_eq!(repaired.len(), 1, "exactly one repair event");
4713 assert_eq!(repaired[0].1, "local-salvage");
4714 }
4715
4716 #[tokio::test]
4717 async fn attempt_cap_honored_and_falls_through_to_error_path() {
4718 let emitter = RecordingEmitter::default();
4720 let tools = vec![read_file_tool()];
4721 let original = json!("totally not json and no braces at all");
4722 let mut calls = vec![ToolCall {
4723 id: "call_2".to_string(),
4724 name: "read_file".to_string(),
4725 arguments: original.clone(),
4726 }];
4727 let configs = vec![crate::AgentCapabilityConfig::with_config(
4728 TOOL_CALL_REPAIR_CAPABILITY_ID,
4729 json!({ "max_reprompts": 0 }),
4730 )];
4731
4732 run(&configs, &tools, &mut calls, &emitter).await;
4733
4734 assert_eq!(calls[0].arguments, original);
4736 assert_eq!(
4737 emitter.repaired_events(),
4738 vec![("call_2".to_string(), "gave-up".to_string())]
4739 );
4740 }
4741
4742 #[tokio::test]
4743 async fn unsalvageable_with_remaining_attempts_labels_reprompt() {
4744 let emitter = RecordingEmitter::default();
4745 let tools = vec![read_file_tool()];
4746 let mut calls = vec![ToolCall {
4747 id: "call_3".to_string(),
4748 name: "read_file".to_string(),
4749 arguments: json!("no json here"),
4750 }];
4751 run(&enabled_configs(), &tools, &mut calls, &emitter).await;
4753 assert_eq!(
4754 emitter.repaired_events(),
4755 vec![("call_3".to_string(), "re-prompt".to_string())]
4756 );
4757 }
4758
4759 #[tokio::test]
4760 async fn attempt_cap_transitions_to_gave_up_on_later_iteration() {
4761 let tools = vec![read_file_tool()];
4766 let garbage = json!("no json here");
4767
4768 let first = RecordingEmitter::default();
4769 let mut calls1 = vec![ToolCall {
4770 id: "call_a".to_string(),
4771 name: "read_file".to_string(),
4772 arguments: garbage.clone(),
4773 }];
4774 run_at(&enabled_configs(), &tools, &mut calls1, &first, 1).await;
4775 assert_eq!(
4776 first.repaired_events(),
4777 vec![("call_a".to_string(), "re-prompt".to_string())]
4778 );
4779
4780 let later = RecordingEmitter::default();
4781 let mut calls2 = vec![ToolCall {
4782 id: "call_a".to_string(),
4783 name: "read_file".to_string(),
4784 arguments: garbage,
4785 }];
4786 run_at(&enabled_configs(), &tools, &mut calls2, &later, 2).await;
4787 assert_eq!(
4788 later.repaired_events(),
4789 vec![("call_a".to_string(), "gave-up".to_string())]
4790 );
4791 }
4792
4793 #[tokio::test]
4794 async fn disabled_capability_produces_no_drift() {
4795 let emitter = RecordingEmitter::default();
4798 let tools = vec![read_file_tool()];
4799 let original = malformed_call().arguments;
4800 let mut calls = vec![malformed_call()];
4801
4802 run(&[], &tools, &mut calls, &emitter).await;
4803
4804 assert_eq!(calls[0].arguments, original);
4805 assert!(emitter.repaired_events().is_empty());
4806 }
4807
4808 #[tokio::test]
4809 async fn well_formed_call_is_untouched_and_silent() {
4810 let emitter = RecordingEmitter::default();
4811 let tools = vec![read_file_tool()];
4812 let mut calls = vec![ToolCall {
4813 id: "call_4".to_string(),
4814 name: "read_file".to_string(),
4815 arguments: json!({ "path": "/already/good" }),
4816 }];
4817
4818 run(&enabled_configs(), &tools, &mut calls, &emitter).await;
4819
4820 assert_eq!(calls[0].arguments, json!({ "path": "/already/good" }));
4821 assert!(emitter.repaired_events().is_empty());
4822 }
4823 }
4824}