use super::{
AgentOutputSink, AgentRunOutput, session_persistence::SessionPersistence,
turn_state::AgentTurnState, util::AssistantChunkBatch,
};
use crate::{
agent::cancellation::{
AgentCancellation, AgentCancellationHandle, AgentRunCanceled, is_run_canceled,
},
agent::ttsr::{TtsrInterrupted, TtsrMatch, TtsrRuleSet},
context::{project_text_tokens, usage_input_tokens},
output::{ContextUsageSource, OutputEvent, redact_sensitive_text},
providers::{
ANTHROPIC_PROVIDER, Provider, ProviderEvent, ProviderRequest, ReasoningSummary, ToolCall,
Usage,
error::{
ProviderStreamTrace, incomplete_semantic_progress_timeout_error,
provider_stream_trace_from_error,
},
},
sessions::SessionEventKind,
};
fn sanitize_provider_response_item(mut item: serde_json::Value) -> serde_json::Value {
fn remove_encrypted_content(value: &mut serde_json::Value) {
match value {
serde_json::Value::Object(object) => {
object.remove("encrypted_content");
for child in object.values_mut() {
remove_encrypted_content(child);
}
}
serde_json::Value::Array(items) => {
for item in items {
remove_encrypted_content(item);
}
}
_ => {}
}
}
remove_encrypted_content(&mut item);
item
}
use serde_json::json;
pub(super) struct ProviderEventCollector<'a, 'sink, 'run> {
pub(super) cancellation: &'a AgentCancellation,
pub(super) output_sink: &'a mut Option<&'sink mut dyn AgentOutputSink>,
pub(super) output: &'a mut AgentRunOutput,
pub(super) turn_state: &'a mut AgentTurnState,
pub(super) assistant_chunk_batch: &'a mut AssistantChunkBatch<'run>,
pub(super) session_persistence: &'a mut SessionPersistence<'run>,
pub(super) context_budget_max_tokens: usize,
pub(super) provider_id: &'a str,
pub(super) model: &'a str,
pub(super) request_sequence: u64,
pub(super) request_input_tokens: usize,
pub(super) ttsr: Option<&'a mut TtsrRuleSet>,
pub(super) ttsr_cancel_handle: Option<&'a AgentCancellationHandle>,
}
pub(super) struct ProviderStreamResult {
pub(super) tool_calls: Vec<ToolCall>,
pub(super) reasoning_summaries: Vec<String>,
pub(super) provider_input_tokens: Option<usize>,
}
pub(super) struct ProviderStreamFailure {
pub(super) error: anyhow::Error,
pub(super) partial_assistant_text: bool,
pub(super) observed_tool_call: bool,
pub(super) observed_function_response_item: bool,
pub(super) observed_reasoning_delta: bool,
pub(super) unsafe_tool_call_progress: bool,
pub(super) reasoning_preview: Option<String>,
pub(super) partial_tool_call_summary: Option<String>,
pub(super) provider_stream_trace: Option<Box<ProviderStreamTrace>>,
pub(super) cancelled: bool,
pub(super) ttsr_match: Option<Box<TtsrMatch>>,
}
impl ProviderStreamFailure {
pub(super) fn recovery_payload(&self) -> Option<serde_json::Value> {
if !self.observed_reasoning_delta && !self.unsafe_tool_call_progress {
return None;
}
let mut payload = json!({
"observed_reasoning_delta": self.observed_reasoning_delta,
"unsafe_tool_call_progress": self.unsafe_tool_call_progress,
});
if let Some(preview) = &self.reasoning_preview
&& !preview.trim().is_empty()
{
payload["reasoning_preview"] = json!(sanitize_recovery_preview(preview));
}
if let Some(summary) = &self.partial_tool_call_summary
&& !summary.trim().is_empty()
{
payload["partial_tool_call_summary"] = json!(summary);
}
Some(payload)
}
pub(super) fn provider_stream_trace_payload(&self) -> Option<serde_json::Value> {
self.provider_stream_trace
.as_deref()
.and_then(|trace| serde_json::to_value(trace).ok())
}
pub(super) fn eligible_for_auto_continue(&self) -> bool {
self.auto_continue_block_reason().is_none()
}
pub(super) fn recovery_diagnostic_relevant(&self) -> bool {
self.unsafe_tool_call_progress
|| self.provider_stream_trace.is_some()
|| incomplete_semantic_progress_timeout_error(&self.error)
}
pub(super) fn auto_continue_block_reason(&self) -> Option<&'static str> {
if self.cancelled {
return Some("run was cancelled");
}
if self.unsafe_tool_call_progress {
return Some(
"partial provider tool-call progress was observed; automatic recovery cannot safely reconstruct arguments",
);
}
if !incomplete_semantic_progress_timeout_error(&self.error) {
return Some("failure was not semantic-progress timeout");
}
if !self.partial_assistant_text {
return Some("no partial assistant text was available to continue from");
}
if self.observed_tool_call {
return Some(
"complete tool call was observed; automatic recovery must not rerun tools",
);
}
if self.observed_function_response_item {
return Some(
"provider function response item was observed; automatic recovery must not infer missing protocol state",
);
}
None
}
}
impl From<anyhow::Error> for ProviderStreamFailure {
fn from(error: anyhow::Error) -> Self {
let cancelled = is_run_canceled(&error);
let provider_stream_trace = provider_stream_trace_from_error(&error).map(Box::new);
let ttsr_match = ttsr_match_from_error(&error);
Self {
error,
partial_assistant_text: false,
observed_tool_call: false,
observed_function_response_item: false,
observed_reasoning_delta: false,
unsafe_tool_call_progress: false,
reasoning_preview: None,
partial_tool_call_summary: None,
provider_stream_trace,
cancelled,
ttsr_match,
}
}
}
pub(super) fn collect_provider_events<P: Provider + ?Sized>(
provider: &P,
request: ProviderRequest,
mut collector: ProviderEventCollector<'_, '_, '_>,
) -> Result<ProviderStreamResult, ProviderStreamFailure> {
let mut tool_calls = Vec::new();
let mut reasoning_summaries = Vec::new();
let mut reasoning_completion_sequence = 0u64;
let mut observed_function_response_item = false;
let mut observed_reasoning_delta = false;
let mut reasoning_preview = String::new();
let mut provider_input_tokens = None;
let mut projection_tracker = ContextProjectionTracker::new(
collector.provider_id,
collector.model,
collector.request_sequence,
collector.request_input_tokens,
collector.context_budget_max_tokens,
);
let stream_result =
provider.stream_cancellable(request, collector.cancellation, &mut |event| {
if collector.cancellation.is_canceled() {
return Err(AgentRunCanceled.into());
}
match event {
ProviderEvent::TextDelta(delta) => {
let should_emit_separator = collector.turn_state.take_segment_separator()
&& !collector.output.text.is_empty();
let mut emit_delta = |text: &str,
record_turn_state: bool|
-> anyhow::Result<()> {
if let Some(sink) = collector.output_sink.as_deref_mut() {
sink.output_event(OutputEvent::AssistantDelta {
text: text.to_string(),
})?;
}
collector.output.text.push_str(text);
if record_turn_state {
collector.turn_state.push_assistant_delta(text);
}
collector.assistant_chunk_batch.push(text);
if let Some(event) = projection_tracker.update_after_delta(text)
&& let Some(sink) = collector.output_sink.as_deref_mut()
&& let Err(error) = sink.output_event(event)
{
if let Err(persistence_error) = collector.assistant_chunk_batch.flush()
{
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
return Err(error);
}
if let Err(error) = collector.assistant_chunk_batch.flush_if_full() {
collector
.session_persistence
.warn_once(collector.output_sink, &error)?;
}
Ok(())
};
if let Some(ttsr) = collector.ttsr.as_deref_mut() {
let candidate = if should_emit_separator {
format!("\n\n{delta}")
} else {
delta.clone()
};
if let Some(ttsr_match) = ttsr.check_text(&candidate) {
if let Some(handle) = collector.ttsr_cancel_handle {
handle.cancel();
}
if let Err(persistence_error) = collector.assistant_chunk_batch.flush()
{
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
return Err(TtsrInterrupted { ttsr_match }.into());
}
}
if should_emit_separator {
emit_delta("\n\n", false)?;
}
emit_delta(&delta, true)?;
}
ProviderEvent::Usage(usage) => {
let input_tokens = provider_usage_input_tokens(collector.provider_id, &usage);
if let Some(sink) = collector.output_sink.as_deref_mut()
&& let Err(error) = sink.output_event(OutputEvent::ContextUsage {
current_tokens: input_tokens,
max_tokens: collector.context_budget_max_tokens,
reasoning_tokens: usage
.reasoning_tokens
.and_then(|tokens| usize::try_from(tokens).ok()),
source: ContextUsageSource::ProviderExact,
request_sequence: collector.request_sequence,
})
{
if let Err(persistence_error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
return Err(error);
}
provider_input_tokens = Some(input_tokens);
accumulate_total_tokens(collector.output, &usage);
collector.output.usage = Some(usage);
}
ProviderEvent::UsagePartial(usage) => {
let input_tokens = provider_usage_input_tokens(collector.provider_id, &usage);
if let Some(sink) = collector.output_sink.as_deref_mut()
&& let Err(error) = sink.output_event(OutputEvent::ContextUsage {
current_tokens: input_tokens,
max_tokens: collector.context_budget_max_tokens,
reasoning_tokens: usage
.reasoning_tokens
.and_then(|tokens| usize::try_from(tokens).ok()),
source: ContextUsageSource::ProviderPartial,
request_sequence: collector.request_sequence,
})
{
if let Err(persistence_error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
return Err(error);
}
collector.output.usage = Some(usage);
}
ProviderEvent::ReasoningSummaryDelta(text) => {
observed_reasoning_delta = true;
push_recovery_preview(&mut reasoning_preview, &text);
if let Err(error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &error)?;
}
if let Some(sink) = collector.output_sink.as_deref_mut()
&& let Err(error) = sink
.output_event(OutputEvent::ThinkingSummaryDelta { text: text.clone() })
{
return Err(error);
}
}
ProviderEvent::ReasoningSummaryComplete(text) => {
reasoning_completion_sequence = reasoning_completion_sequence.saturating_add(1);
let item_id = Some(format!("legacy-{}-{reasoning_completion_sequence}", collector.request_sequence));
let turn_id = Some(collector.request_sequence.to_string());
handle_reasoning_summary_completion(&mut collector, &mut reasoning_summaries, &mut observed_reasoning_delta, &mut reasoning_preview, text, Some((item_id, turn_id)))?;
}
ProviderEvent::ReasoningSummaryCompleteIdentified(ReasoningSummary { text, item_id, turn_id }) => {
let turn_id = turn_id.or_else(|| Some(collector.request_sequence.to_string()));
handle_reasoning_summary_completion(&mut collector, &mut reasoning_summaries, &mut observed_reasoning_delta, &mut reasoning_preview, text, Some((item_id, turn_id)))?;
}
ProviderEvent::ToolCall(call) => {
if let Some(event) = projection_tracker.force_event_if_changed()
&& let Some(sink) = collector.output_sink.as_deref_mut()
&& let Err(error) = sink.output_event(event)
{
if let Err(persistence_error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
return Err(error);
}
if let Some(ttsr) = collector.ttsr.as_deref_mut() {
let candidate =
json!({"name": call.name, "arguments": call.arguments}).to_string();
if let Some(ttsr_match) = ttsr.check_standalone(&candidate) {
if let Some(handle) = collector.ttsr_cancel_handle {
handle.cancel();
}
if let Err(persistence_error) = collector.assistant_chunk_batch.flush()
{
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
return Err(TtsrInterrupted { ttsr_match }.into());
}
}
tool_calls.push(call);
}
ProviderEvent::ResponseItem(item) => {
if item.get("type").and_then(serde_json::Value::as_str) == Some("function_call")
{
observed_function_response_item = true;
}
collector.turn_state.push_response_item(item.clone());
if let Err(error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &error)?;
}
collector.session_persistence.try_record(
SessionEventKind::ProviderResponseItem,
json!({"item": sanitize_provider_response_item(item.clone())}),
collector.output_sink,
)?;
}
ProviderEvent::ResponseIdentity(identity) => {
collector.session_persistence.try_record(
SessionEventKind::ProviderStreamTrace,
json!({"response_identity": identity}),
collector.output_sink,
)?;
if let Some(sink) = collector.output_sink.as_deref_mut() {
sink.output_event(OutputEvent::Diagnostic {
level: "info".to_string(),
message: format!("OpenAI Responses attempt {} observed provider-reported model metadata", identity.attempt),
})?;
}
}
ProviderEvent::Done => {
if let Some(event) = projection_tracker.force_event_if_changed()
&& let Some(sink) = collector.output_sink.as_deref_mut()
&& let Err(error) = sink.output_event(event)
{
return Err(error);
}
}
}
collector.cancellation.check()
});
if let Err(error) = stream_result {
if let Err(persistence_error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
let cancelled = is_run_canceled(&error);
let unsafe_tool_call_progress = unsafe_tool_call_progress_from_error(&error);
let provider_stream_trace = provider_stream_trace_from_error(&error).map(Box::new);
let ttsr_match = ttsr_match_from_error(&error);
return Err(ProviderStreamFailure {
error,
partial_assistant_text: !collector.output.text.trim().is_empty(),
observed_tool_call: !tool_calls.is_empty(),
observed_function_response_item,
observed_reasoning_delta,
unsafe_tool_call_progress,
reasoning_preview: (!reasoning_preview.trim().is_empty()).then_some(reasoning_preview),
partial_tool_call_summary: unsafe_tool_call_progress.then(|| {
"partial provider tool-call progress observed; raw arguments omitted; complete-tool-call recovery not attempted without a finished parser event".to_string()
}),
provider_stream_trace,
cancelled,
ttsr_match,
});
}
if let Err(error) = collector.cancellation.check() {
if let Err(persistence_error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &persistence_error)?;
}
let unsafe_tool_call_progress = false;
return Err(ProviderStreamFailure {
error,
partial_assistant_text: !collector.output.text.trim().is_empty(),
observed_tool_call: !tool_calls.is_empty(),
observed_function_response_item,
observed_reasoning_delta,
unsafe_tool_call_progress,
reasoning_preview: (!reasoning_preview.trim().is_empty()).then_some(reasoning_preview),
partial_tool_call_summary: None,
provider_stream_trace: None,
cancelled: true,
ttsr_match: None,
});
}
Ok(ProviderStreamResult {
tool_calls,
reasoning_summaries,
provider_input_tokens,
})
}
struct ContextProjectionTracker<'a> {
provider_id: &'a str,
model: &'a str,
request_sequence: u64,
request_input_tokens: usize,
max_tokens: usize,
has_output: bool,
pending_delta_text: String,
projected_output_tokens: usize,
projection_source: ContextUsageSource,
last_emitted_tokens: usize,
}
impl<'a> ContextProjectionTracker<'a> {
const FLUSH_BYTES: usize = 512;
const EMIT_TOKEN_DELTA: usize = 16;
fn new(
provider_id: &'a str,
model: &'a str,
request_sequence: u64,
request_input_tokens: usize,
max_tokens: usize,
) -> Self {
Self {
provider_id,
model,
request_sequence,
request_input_tokens,
max_tokens,
has_output: false,
pending_delta_text: String::new(),
projected_output_tokens: 0,
projection_source: ContextUsageSource::FallbackProjection,
last_emitted_tokens: request_input_tokens,
}
}
fn update_after_delta(&mut self, delta: &str) -> Option<OutputEvent> {
self.has_output = true;
self.pending_delta_text.push_str(delta);
if self.pending_delta_text.len() < Self::FLUSH_BYTES {
return None;
}
let pending_delta = std::mem::take(&mut self.pending_delta_text);
let projection = project_text_tokens(self.provider_id, self.model, &pending_delta);
self.projected_output_tokens = self
.projected_output_tokens
.saturating_add(projection.tokens);
self.projection_source = projection.source;
if self.projected_output_tokens.abs_diff(
self.last_emitted_tokens
.saturating_sub(self.request_input_tokens),
) < Self::EMIT_TOKEN_DELTA
{
return None;
}
let event = self.project_event_from_incremental();
let current_tokens = event_current_tokens(&event);
(current_tokens.abs_diff(self.last_emitted_tokens) >= Self::EMIT_TOKEN_DELTA).then(|| {
self.last_emitted_tokens = current_tokens;
event
})
}
fn force_event_if_changed(&mut self) -> Option<OutputEvent> {
if !self.has_output {
return None;
}
if !self.pending_delta_text.is_empty() {
let pending_delta = std::mem::take(&mut self.pending_delta_text);
let projection = project_text_tokens(self.provider_id, self.model, &pending_delta);
self.projected_output_tokens = self
.projected_output_tokens
.saturating_add(projection.tokens);
self.projection_source = projection.source;
}
let event = self.project_event_from_incremental();
let current_tokens = event_current_tokens(&event);
(current_tokens != self.last_emitted_tokens).then(|| {
self.last_emitted_tokens = current_tokens;
event
})
}
fn project_event_from_incremental(&self) -> OutputEvent {
OutputEvent::ContextUsage {
current_tokens: self
.request_input_tokens
.saturating_add(self.projected_output_tokens),
max_tokens: self.max_tokens,
reasoning_tokens: None,
source: self.projection_source,
request_sequence: self.request_sequence,
}
}
}
fn event_current_tokens(event: &OutputEvent) -> usize {
match event {
OutputEvent::ContextUsage { current_tokens, .. } => *current_tokens,
_ => unreachable!("context projection tracker only emits context usage"),
}
}
fn provider_usage_input_tokens(provider_id: &str, usage: &Usage) -> usize {
if matches!(
provider_id,
ANTHROPIC_PROVIDER | crate::providers::CLAUDE_CODE_PROVIDER
) {
let tokens = usage
.input
.saturating_add(usage.cache_read)
.saturating_add(usage.cache_write);
return usize::try_from(tokens).unwrap_or(usize::MAX);
}
usage_input_tokens(usage)
}
fn accumulate_total_tokens(output: &mut AgentRunOutput, usage: &crate::providers::Usage) {
let tokens = if usage.total > 0 {
usage.total
} else {
usage.input.saturating_add(usage.output)
};
output.total_tokens = Some(
output
.total_tokens
.unwrap_or_default()
.saturating_add(tokens),
);
}
fn push_recovery_preview(preview: &mut String, delta: &str) {
const RECOVERY_REASONING_PREVIEW_LIMIT: usize = 4_000;
if preview.chars().count() >= RECOVERY_REASONING_PREVIEW_LIMIT {
return;
}
let remaining = RECOVERY_REASONING_PREVIEW_LIMIT.saturating_sub(preview.chars().count());
preview.extend(delta.chars().take(remaining));
}
fn sanitize_recovery_preview(text: &str) -> String {
const RECOVERY_REASONING_PREVIEW_LIMIT: usize = 4_000;
let redacted = redact_sensitive_text(text.trim());
redacted
.chars()
.take(RECOVERY_REASONING_PREVIEW_LIMIT)
.collect()
}
fn handle_reasoning_summary_completion(
collector: &mut ProviderEventCollector<'_, '_, '_>,
reasoning_summaries: &mut Vec<String>,
observed_reasoning_delta: &mut bool,
reasoning_preview: &mut String,
text: String,
identity: Option<(Option<String>, Option<String>)>,
) -> anyhow::Result<()> {
if let Err(error) = collector.assistant_chunk_batch.flush() {
collector
.session_persistence
.warn_once(collector.output_sink, &error)?;
}
let (item_id, turn_id) = identity.unwrap_or((None, None));
let event = if item_id.is_some() || turn_id.is_some() {
OutputEvent::ThinkingSummaryCompleteIdentified {
text: text.clone(),
item_id: item_id.clone(),
turn_id: turn_id.clone(),
}
} else {
OutputEvent::ThinkingSummaryComplete { text: text.clone() }
};
if let Some(sink) = collector.output_sink.as_deref_mut() {
sink.output_event(event)?;
}
let mut payload = json!({"text": text});
if let Some(item_id) = item_id {
payload["item_id"] = json!(item_id);
}
if let Some(turn_id) = turn_id {
payload["turn_id"] = json!(turn_id);
}
collector.session_persistence.try_record(
SessionEventKind::ReasoningSummary,
payload,
collector.output_sink,
)?;
reasoning_summaries.push(text);
*observed_reasoning_delta = false;
reasoning_preview.clear();
Ok(())
}
fn unsafe_tool_call_progress_from_error(error: &anyhow::Error) -> bool {
error.chain().any(|cause| {
cause
.to_string()
.contains("unsafe tool-call progress observed")
})
}
fn ttsr_match_from_error(error: &anyhow::Error) -> Option<Box<TtsrMatch>> {
error
.chain()
.find_map(|cause| cause.downcast_ref::<TtsrInterrupted>())
.map(|interrupted| Box::new(interrupted.ttsr_match.clone()))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::providers::error::ProviderError;
#[test]
fn anthropic_family_usage_accounting_includes_cached_tokens() {
let usage = Usage {
input: 10,
cache_read: 20,
cache_write: 30,
..Usage::default()
};
assert_eq!(provider_usage_input_tokens(ANTHROPIC_PROVIDER, &usage), 60);
assert_eq!(
provider_usage_input_tokens(crate::providers::CLAUDE_CODE_PROVIDER, &usage),
60
);
assert_eq!(provider_usage_input_tokens("custom", &usage), 10);
}
#[test]
fn auto_continue_block_reason_explains_partial_tool_call_progress() {
let failure = ProviderStreamFailure {
error: ProviderError::stream_failed_incomplete(
"provider stream ended prematurely after partial response; response is incomplete: provider stream no semantic progress before timeout",
)
.into(),
partial_assistant_text: true,
observed_tool_call: false,
observed_function_response_item: false,
observed_reasoning_delta: false,
unsafe_tool_call_progress: true,
reasoning_preview: None,
partial_tool_call_summary: Some("partial provider tool-call progress observed; raw arguments omitted".to_string()),
provider_stream_trace: None,
cancelled: false,
ttsr_match: None,
};
assert!(!failure.eligible_for_auto_continue());
assert_eq!(
failure.auto_continue_block_reason(),
Some(
"partial provider tool-call progress was observed; automatic recovery cannot safely reconstruct arguments"
)
);
}
}