use super::{
Agent, AgentRunTime, FailureKind, PendingToolCall, ToolRunState, prompt::PromptComposeInput,
storage_failure,
};
use crate::{
error::{AgentError, StorageError},
middleware::MiddlewareCtx,
providers::{FinishReason, ProviderError, ProviderErrorKind, ProviderStreamChunk},
shared::{ContentBlock, Message, Role, ToolCall},
state::AgentState,
trajectory::TrajectoryEventKind,
transcript::Transcript,
};
use futures_util::StreamExt;
use serde_json::Value;
use std::collections::{BTreeMap, HashSet};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
impl AgentRunTime {
pub(crate) async fn handle_reasoning(
&self,
ctx: AgentState,
cancellation_token: CancellationToken,
) -> Result<Agent, AgentError> {
let history = match self.load_repaired_history(&ctx.session_id).await {
Ok(history) => history,
Err(e) => return Ok(storage_failure(ctx, e)),
};
let mut req = self.prompt_composer.compose(
PromptComposeInput {
tools: self.tools.definitions(),
model: self.model.to_string(),
system_prompt: self.system_prompt.clone(),
history,
skills: self
.skills
.as_ref()
.map(|skills| skills.list())
.unwrap_or_default(),
temperature: self.temperature,
max_tokens: self.max_tokens,
},
&ctx,
);
{
let mw_ctx = MiddlewareCtx {
session_id: &ctx.session_id,
state: &ctx,
};
if let Err(e) = self.middleware.before_model(&mw_ctx, &mut req).await {
return Ok(self.fail_with(ctx, FailureKind::Middleware, e));
}
}
self.emit(
&ctx.session_id,
TrajectoryEventKind::ReasoningStarted {
request: self.provider.request_snapshot(&req),
},
)
.await;
let mut stream = match self.provider.as_ref().chat_completion(req).await {
Ok(stream) => stream,
Err(e) => {
let message = format!("[{}] {e}", self.provider.name());
return Ok(self.fail_with(ctx, provider_failure_kind(&e, false), message));
}
};
let mut content = String::new();
let mut thinking = String::new();
let mut tool_calls = BTreeMap::<usize, PendingToolCall>::new();
let mut stream_error = None;
let mut finish_reason = None;
loop {
tokio::select! {
_ = cancellation_token.cancelled() => {
if let Err(e) = self.save_interruption(&Agent::Reasoning(ctx.clone())).await {
return match e {
AgentError::Storage(storage) => Ok(storage_failure(ctx.clone(), storage)),
other => Err(other),
};
}
return Ok(Agent::Interrupted(ctx));
}
maybe_chunk = stream.next() => {
let Some(chunk) = maybe_chunk else {
break;
};
match chunk {
Ok(ProviderStreamChunk::ContentDelta(text)) => {
content.push_str(&text)
}
Ok(ProviderStreamChunk::ThinkingDelta(text)) => {
thinking.push_str(&text)
}
Ok(ProviderStreamChunk::ToolCallDelta {
index,
id,
name,
arguments_delta,
}) => {
let call = tool_calls.entry(index).or_default();
if id.is_some() {
call.id = id;
}
if name.is_some() {
call.name = name;
}
call.arguments.push_str(&arguments_delta);
}
Ok(ProviderStreamChunk::Done { finish_reason: reason, .. }) => {
finish_reason = Some(reason);
break;
}
Err(e) => {
stream_error = Some(e);
break;
}
}
}
}
}
if let Some(error) = stream_error {
let message = error.to_string();
return Ok(self.fail_with(ctx, provider_failure_kind(&error, true), message));
}
let Some(finish_reason) = finish_reason else {
return Ok(self.fail_with(
ctx,
FailureKind::ProviderStream,
"stream ended before completion (no finish reason)",
));
};
match finish_reason {
FinishReason::Stop | FinishReason::ToolCalls => {}
FinishReason::Length => {
return Ok(self.fail_with(
ctx,
FailureKind::Truncated,
"model output truncated by length limit",
));
}
FinishReason::ContentFilter => {
return Ok(self.fail_with(
ctx,
FailureKind::ContentFilter,
"model output blocked by content filter",
));
}
FinishReason::Error => {
return Ok(self.fail_with(
ctx,
FailureKind::ProviderStream,
"provider reported an error finish reason",
));
}
FinishReason::Unknown(other) => {
return Ok(self.fail_with(
ctx,
FailureKind::ProviderStream,
format!("unknown finish reason: {}", other),
));
}
}
let assistant_text = if content.is_empty() {
None
} else {
Some(content.clone())
};
let mut blocks = Vec::new();
if !thinking.is_empty() {
blocks.push(ContentBlock::Thinking { text: thinking });
}
if !content.is_empty() {
blocks.push(ContentBlock::Text { text: content });
}
let mut pending_tool_calls = Vec::new();
for (_, call) in tool_calls {
let Some(name) = call.name else {
continue;
};
let id = call.id.unwrap_or_else(|| Uuid::new_v4().to_string());
let input = parse_tool_arguments(&call.arguments);
blocks.push(ContentBlock::ToolUse {
id: id.clone(),
name: name.clone(),
input: input.clone(),
});
pending_tool_calls.push(ToolCall {
call_id: id,
tool_name: name,
arguments: input,
});
}
let mut assistant_msg = Message {
role: Role::Assistant,
blocks,
};
{
let mw_ctx = MiddlewareCtx {
session_id: &ctx.session_id,
state: &ctx,
};
if let Err(e) = self
.middleware
.after_model(&mw_ctx, &mut assistant_msg)
.await
{
return Ok(self.fail_with(ctx, FailureKind::Middleware, e));
}
}
self.emit(
&ctx.session_id,
TrajectoryEventKind::AssistantMessage {
blocks: assistant_msg.blocks.clone(),
},
)
.await;
if let Err(e) = self
.transcript_storage
.append_message(&ctx.session_id, assistant_msg)
.await
{
return Ok(storage_failure(ctx, e));
}
let mut next_ctx = ctx.clone();
next_ctx.consecutive_fail_count = 0;
if pending_tool_calls.is_empty() {
next_ctx.final_resp = assistant_text.clone();
self.emit(
&ctx.session_id,
TrajectoryEventKind::Succeeded {
final_resp: assistant_text,
},
)
.await;
Ok(Agent::Success(next_ctx))
} else {
Ok(Agent::ToolExecuting(
next_ctx,
ToolRunState::new(pending_tool_calls),
))
}
}
async fn load_repaired_history(&self, session_id: &str) -> Result<Vec<Message>, StorageError> {
let mut transcript = self
.transcript_storage
.get_transcript(session_id)
.await?
.unwrap_or_else(|| Transcript::new(session_id));
let repairs = orphan_tool_use_repairs(&transcript.messages);
if !repairs.is_empty() {
transcript.messages.extend(repairs.clone());
self.transcript_storage
.append_messages(session_id, repairs)
.await?;
}
Ok(transcript.messages)
}
}
fn provider_failure_kind(error: &ProviderError, stream: bool) -> FailureKind {
match error.kind {
ProviderErrorKind::Permanent => FailureKind::ProviderPermanent,
ProviderErrorKind::RateLimited { .. } => FailureKind::ProviderRateLimited,
ProviderErrorKind::Transient if stream => FailureKind::ProviderStream,
ProviderErrorKind::Transient => FailureKind::Provider,
}
}
fn orphan_tool_use_repairs(messages: &[Message]) -> Vec<Message> {
let answered: HashSet<&str> = messages
.iter()
.flat_map(|message| message.blocks.iter())
.filter_map(|block| match block {
ContentBlock::ToolResult { tool_use_id, .. } => Some(tool_use_id.as_str()),
_ => None,
})
.collect();
let mut repairs = Vec::new();
for block in messages.iter().flat_map(|message| message.blocks.iter()) {
let ContentBlock::ToolUse { id, .. } = block else {
continue;
};
if answered.contains(id.as_str()) {
continue;
}
repairs.push(Message {
role: Role::Tool,
blocks: vec![ContentBlock::ToolResult {
tool_use_id: id.clone(),
content: vec![ContentBlock::Text {
text: "tool result missing after previous crash; treating as failed"
.to_string(),
}],
is_error: true,
}],
});
}
repairs
}
fn parse_tool_arguments(arguments: &str) -> Value {
if arguments.trim().is_empty() {
Value::Object(Default::default())
} else {
serde_json::from_str(arguments).unwrap_or_else(|_| Value::String(arguments.to_string()))
}
}