use af_agent::HookDecision;
use af_agent_session::{
Event, InteractionResolution, RunStatus, SessionProjection, ToolAuthorizationStatus,
};
use af_llm::ChatMessage;
use serde_json::json;
use crate::replay::{interaction_resolution_for_call, tool_authorization_for_call, tool_message};
use crate::tools::PlannedCall;
use crate::{
failure_events, recovery_events, AgentRuntime, CancellationToken, EventWriter, RuntimeError,
RuntimeOutcome, TurnRequest,
};
pub(super) struct RecoveryPhase {
pub(super) first_step: u32,
pub(super) terminal: Option<RuntimeOutcome>,
}
impl AgentRuntime {
pub(super) async fn recover_open_surface(
&self,
request: &TurnRequest,
projection: &SessionProjection,
writer: &dyn EventWriter,
cancellation: CancellationToken,
transcript: &mut Vec<ChatMessage>,
usage: (u64, u64),
) -> Result<RecoveryPhase, RuntimeError> {
if projection.open_turn.is_none() {
return Ok(RecoveryPhase {
first_step: 1,
terminal: None,
});
}
if projection.open_compaction.is_some() {
let mut events = Vec::new();
let mut terminal_usage = usage;
if let Some((compaction_id, _)) = &projection.open_compaction {
let operation_prefix = format!("{compaction_id}:attempt:");
if let Some(Event::ModelRequestPrepared {
operation_id,
reserved_prompt_tokens,
reserved_completion_tokens,
..
}) = request
.history
.iter()
.rev()
.find_map(|event| match &event.event {
prepared @ Event::ModelRequestPrepared {
run_id,
operation_id,
..
} if run_id == &request.run_id
&& operation_id.starts_with(&operation_prefix) =>
{
Some(prepared)
}
_ => None,
})
{
let already_recorded = request.history.iter().any(|event| {
matches!(
&event.event,
Event::UsageRecorded { run_id, operation_id: recorded, .. }
if run_id == &request.run_id && recorded == operation_id
)
});
if !already_recorded {
terminal_usage.0 += reserved_prompt_tokens;
terminal_usage.1 += reserved_completion_tokens;
events.push(Event::UsageRecorded {
run_id: request.run_id.clone(),
operation_id: operation_id.clone(),
prompt_tokens: *reserved_prompt_tokens,
completion_tokens: *reserved_completion_tokens,
cost_units: 0,
});
}
}
}
events.extend(recovery_events(projection));
writer.append(events).await?;
return Ok(RecoveryPhase {
first_step: 1,
terminal: Some(failed(terminal_usage)),
});
}
if !projection.started_tool_calls.is_empty() {
writer
.append(failure_events(projection, "tool_outcome_unknown"))
.await?;
return Ok(RecoveryPhase {
first_step: 1,
terminal: Some(failed(usage)),
});
}
let mut pending = projection
.open_tool_calls
.iter()
.map(|(id, call)| PlannedCall {
transcript_id: id.clone(),
id: id.clone(),
name: call.tool.clone(),
arguments: call.arguments.clone(),
step: call.step,
source_event_seq: call.source_event_seq,
preflight_error: (!self.tools.contains(&call.tool))
.then(|| format!("tool '{}' is unavailable", call.tool)),
})
.collect::<Vec<_>>();
pending.sort_by_key(|call| call.source_event_seq);
let mut executable = Vec::new();
for call in pending {
if let Some(reason) = &call.preflight_error {
deny_recovered(writer, request, &call, reason, json!({"error":reason})).await?;
transcript.push(tool_message(&call, json!({"error":reason})));
continue;
}
let prior_authorization =
tool_authorization_for_call(&request.history, &request.run_id, &call.id);
let resolution = interaction_resolution_for_call(
&request.history,
&request.run_id,
&call.id,
call.source_event_seq,
);
if prior_authorization == Some(ToolAuthorizationStatus::Denied) {
let value = json!({"error":"tool_denied"});
writer
.append(vec![Event::ToolResult {
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
result: value.clone(),
is_error: true,
}])
.await?;
transcript.push(tool_message(&call, value));
continue;
}
if prior_authorization == Some(ToolAuthorizationStatus::Waiting) && resolution.is_none()
{
writer
.append(failure_events(projection, "tool_outcome_unknown"))
.await?;
return Ok(RecoveryPhase {
first_step: 1,
terminal: Some(failed(usage)),
});
}
if resolution == Some(InteractionResolution::Rejected) {
let value = json!({"error":"user_rejected"});
deny_recovered(writer, request, &call, "user_rejected", value.clone()).await?;
transcript.push(tool_message(&call, value));
continue;
}
if resolution == Some(InteractionResolution::Answered) {
let value = json!({"error":"user_answer_requires_replan"});
deny_recovered(
writer,
request,
&call,
"user_answer_requires_replan",
value.clone(),
)
.await?;
transcript.push(tool_message(&call, value));
continue;
}
match self
.authorize_call(request, &call, resolution, cancellation.clone())
.await
{
HookDecision::Continue => {
writer
.append(vec![Event::ToolAuthorization {
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
status: ToolAuthorizationStatus::Allowed,
reason: None,
}])
.await?;
executable.push((call, resolution));
}
HookDecision::Deny { reason } | HookDecision::WaitForInput { kind: reason, .. } => {
let value = json!({"error":reason});
deny_recovered(writer, request, &call, &reason, value.clone()).await?;
transcript.push(tool_message(&call, value));
}
}
}
for (call, resolution) in executable {
let result = self
.execute_tools(
request,
std::slice::from_ref(&call),
writer,
cancellation.clone(),
resolution,
)
.await?
.pop()
.expect("one recovered tool produces one result");
let outcome_unknown = matches!(
&result,
Err(error) if error.contains("tool_outcome_unknown")
);
let value = match result {
Ok(value) => value,
Err(error) => json!({"error": error}),
};
if outcome_unknown {
writer
.append(failure_events(projection, "tool_outcome_unknown"))
.await?;
return Ok(RecoveryPhase {
first_step: 1,
terminal: Some(failed(usage)),
});
}
if let Err(error) = self
.run_after_hooks(request, &call, &value, resolution, cancellation.clone())
.await
{
writer
.append(vec![Event::Extension {
run_id: request.run_id.clone(),
plugin_id: "agentfactory.runtime".into(),
event_type: "after_tool_failed".into(),
payload: json!({"call_id":call.id,"error":error}),
}])
.await?;
}
transcript.push(tool_message(&call, value));
}
for step in &projection.open_steps {
writer
.append(vec![Event::StepFinished {
run_id: request.run_id.clone(),
step: *step,
}])
.await?;
}
Ok(RecoveryPhase {
first_step: projection.open_steps.last().copied().unwrap_or(0) + 1,
terminal: None,
})
}
}
async fn deny_recovered(
writer: &dyn EventWriter,
request: &TurnRequest,
call: &PlannedCall,
reason: &str,
result: serde_json::Value,
) -> Result<(), RuntimeError> {
writer
.append(vec![
Event::ToolAuthorization {
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
status: ToolAuthorizationStatus::Denied,
reason: Some(reason.into()),
},
Event::ToolResult {
run_id: request.run_id.clone(),
step: call.step,
call_id: call.id.clone(),
result,
is_error: true,
},
])
.await?;
Ok(())
}
fn failed((prompt_tokens, completion_tokens): (u64, u64)) -> RuntimeOutcome {
RuntimeOutcome {
status: RunStatus::Failed.as_str().into(),
final_text: None,
prompt_tokens,
completion_tokens,
waiting_interaction_id: None,
}
}