use std::collections::{HashMap, HashSet};
use async_trait::async_trait;
use futures::SinkExt as _;
use polyc_crypto::canon::canon_args;
use polyc_llm::{
CompletionRequest, Content as LlmContent, LlmProvider, Message as LlmMessage, Role, StopReason,
ToolSpec, Usage,
request::ToolCall,
turn::{collect_turn, collect_turn_observed},
};
use polyc_proto::proto::polychrome::agent::v1::Message;
use crate::{
ApprovalOverride, CallContext, CallDisposition, HandoffRequest, PendingApproval, ResolvedCall,
RunTurnOptions, ToolExecutor, append_injected_notes, cap_tool_result, forced_result,
gate_decision, gate_missing, push_internal_note, push_reasoning, resolve_approved_call,
run_and_redact, session_approves, splice_results_after, text_message, tool_result_message,
untrusted_content_in_context,
};
pub(crate) const RESUME_EXECUTED_GROUND_TRUTH_NOTE: &str = "A person approved this request and the tool has \
already run — the results above are final. Tell the user what happened; do not describe \
approval status, the system already showed it.";
#[allow(clippy::struct_excessive_bools)]
pub struct TurnCtx<'a, P, T>
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
pub provider: &'a P,
pub tools: &'a T,
pub model: &'a str,
pub options: &'a RunTurnOptions,
pub messages: Vec<LlmMessage>,
pub outputs: Vec<Message>,
pub total_usage: Usage,
pub last_stop: Option<StopReason>,
pub executed_tools: bool,
pub produced_text: bool,
pub pending_handoff: Option<HandoffRequest>,
pub denied_sigs: HashSet<(String, String)>,
pub approved_remaining: HashSet<(String, String, String)>,
pub denial_reprompts: usize,
pub saw_sig_match_denial: bool,
pub grant_replays: Vec<crate::GrantReplayClear>,
pub unattended_denials: Vec<crate::UnattendedDenial>,
pub escape_hatch_fired: bool,
pub delegate_records: Vec<crate::DelegateRecord>,
}
impl<P, T> TurnCtx<'_, P, T>
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
#[must_use]
pub fn finish(
self,
pending_approvals: Vec<PendingApproval>,
handoff: Option<HandoffRequest>,
) -> crate::TurnResult {
crate::TurnResult {
messages: self.outputs,
usage: self.total_usage,
stop: self.last_stop,
pending_approvals,
handoff,
grant_replays: self.grant_replays,
unattended_denials: self.unattended_denials,
mid_stream_failure: None,
delegate_records: self.delegate_records,
}
}
#[must_use]
pub fn finish_failed(mut self, failure: crate::MidStreamFailure) -> crate::TurnResult {
let handoff = self.pending_handoff.take();
let mut result = self.finish(Vec::new(), handoff);
result.mid_stream_failure = Some(failure);
result
}
}
pub enum StepOutcome {
Continue,
Pause(Vec<PendingApproval>),
Done,
}
#[async_trait]
pub trait TurnStep<P, T>: Send + Sync
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
async fn run(&self, ctx: &mut TurnCtx<'_, P, T>) -> Result<StepOutcome, P::Error>;
}
pub struct ForcedCompletion;
#[async_trait]
impl<P, T> TurnStep<P, T> for ForcedCompletion
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
async fn run(&self, ctx: &mut TurnCtx<'_, P, T>) -> Result<StepOutcome, P::Error> {
if !(ctx.executed_tools && !ctx.produced_text && ctx.pending_handoff.is_none()) {
return Ok(StepOutcome::Continue);
}
let mut req = CompletionRequest::new(ctx.model);
req.messages.clone_from(&ctx.messages);
req.messages.push(LlmMessage {
role: Role::System,
content: vec![LlmContent::Text(
"No tools are available for the remainder of this turn. Give the \
user a direct, plain-text answer using the information already \
gathered."
.to_owned(),
)],
});
req.tools = Vec::new();
req.web_search = false;
if let Ok(stream) = ctx.provider.complete(req).await {
let turn = if let Some(tx) = ctx.options.stream_tx.clone() {
let mut tx = tx;
collect_turn_observed(stream, async move |ev| {
let _ = tx.send(ev).await;
})
.await
} else {
collect_turn(stream).await
};
if let Ok(turn) = turn {
ctx.total_usage.input_tokens += turn.usage.input_tokens;
ctx.total_usage.output_tokens += turn.usage.output_tokens;
push_reasoning(&mut ctx.outputs, &turn.reasoning);
if !turn.text.is_empty() {
ctx.outputs.push(text_message("model", &turn.text));
}
ctx.last_stop = turn.stop;
tracing::info!(
"forced closing completion (tool loop produced no text); turn now yields a reply"
);
}
}
Ok(StepOutcome::Continue)
}
}
pub struct ResumePrePass<'a> {
pub tool_specs: &'a [ToolSpec],
pub approved_overrides: &'a HashMap<(String, String, String), ApprovalOverride>,
pub denied_call_ids: &'a HashSet<(String, String, String)>,
}
#[async_trait]
impl<P, T> TurnStep<P, T> for ResumePrePass<'_>
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
#[allow(clippy::too_many_lines)] async fn run(&self, ctx: &mut TurnCtx<'_, P, T>) -> Result<StepOutcome, P::Error> {
let answered: HashSet<&str> = ctx
.messages
.iter()
.flat_map(|m| m.content.iter())
.filter_map(|c| match c {
LlmContent::ToolResult(tr) => Some(tr.tool_call_id.as_str()),
_ => None,
})
.collect();
let mut unanswered: Vec<(usize, ToolCall)> = Vec::new();
for (idx, m) in ctx.messages.iter().enumerate() {
for c in &m.content {
if let LlmContent::ToolUse(tc) = c
&& !answered.contains(tc.id.as_str())
{
unanswered.push((idx, tc.clone()));
}
}
}
if !ctx.options.approved_call_ids.is_empty() {
let unanswered_ids: HashSet<&str> =
unanswered.iter().map(|(_, tc)| tc.id.as_str()).collect();
for (id, name, _args) in &ctx.options.approved_call_ids {
if unanswered_ids.contains(id.as_str()) {
continue;
}
if answered.contains(id.as_str()) {
tracing::info!(
request_id = %id,
tool = %name,
"approved call already answered on this resume; decision is a no-op re-forward"
);
} else {
tracing::warn!(
request_id = %id,
tool = %name,
"approved call id matches no tool call in the resumed \
transcript; the signed decision cannot resolve to anything"
);
}
}
}
tracing::info!(
approved = ctx.options.approved_call_ids.len(),
denied = ctx.options.denied_call_ids.len(),
unanswered = unanswered.len(),
"resume pre-pass: resolving forwarded decisions against the resumed transcript"
);
if unanswered.is_empty() {
return Ok(StepOutcome::Continue);
}
let untrusted_in_context =
untrusted_content_in_context(&ctx.messages) || ctx.options.untrusted_context_seed;
let dispositions: Vec<CallDisposition> = unanswered
.iter()
.map(|(_, tc)| {
let gate = gate_decision(
ctx.tools,
ctx.options,
untrusted_in_context,
&tc.name,
&tc.args_json,
);
let key = (tc.id.clone(), tc.name.clone(), canon_args(&tc.args_json));
let is_denied = self.denied_call_ids.contains(&key);
let is_approved = ctx.approved_remaining.contains(&key)
|| session_approves(ctx.options, ctx.tools, &tc.name, gate_missing(&gate));
CallDisposition::classify(
gate,
CallContext {
approved: is_approved,
denied: is_denied,
..CallContext::default()
},
)
})
.collect();
let mut execute = 0usize;
let mut denied = 0usize;
let mut policy_denied = 0usize;
let mut pending = 0usize;
for disposition in &dispositions {
match disposition {
CallDisposition::Execute => execute += 1,
CallDisposition::Denied { .. } => denied += 1,
CallDisposition::PolicyDenied { .. } | CallDisposition::UnattendedDenied { .. } => {
policy_denied += 1;
}
CallDisposition::Recovered { .. } => {}
CallDisposition::Pending { .. } => pending += 1,
}
}
tracing::info!(
execute,
denied,
policy_denied,
pending,
"resume pre-pass: classified every unanswered dangling call"
);
if dispositions
.iter()
.any(|d| matches!(d, CallDisposition::Pending { .. }))
{
let pending = unanswered
.iter()
.zip(&dispositions)
.filter_map(|((_, tc), d)| {
let CallDisposition::Pending { reason, missing } = d else {
return None;
};
let title = self
.tool_specs
.iter()
.find(|s| s.name == tc.name)
.and_then(|s| s.title.clone())
.unwrap_or_default();
Some(PendingApproval {
id: tc.id.clone(),
name: tc.name.clone(),
args_json: tc.args_json.clone(),
title,
sandbox_mode: String::new(),
reason: reason.clone(),
missing_capabilities: missing
.names()
.iter()
.map(|n| (*n).to_owned())
.collect(),
})
})
.collect::<Vec<_>>();
return Ok(StepOutcome::Pause(pending));
}
let pre_resolutions: Vec<ResolvedCall> = unanswered
.iter()
.map(|(_, tc)| {
let key = (tc.id.clone(), tc.name.clone(), canon_args(&tc.args_json));
resolve_approved_call(&tc.args_json, self.approved_overrides.get(&key))
})
.collect();
let tools = ctx.tools;
let recorder = ctx.options.dispatch_recorder.clone();
let futures = unanswered
.iter()
.zip(&dispositions)
.zip(&pre_resolutions)
.map(|(((_, tc), disposition), resolved)| {
if matches!(disposition, CallDisposition::Denied { .. }) {
ctx.denied_sigs
.insert((tc.name.clone(), canon_args(&tc.args_json)));
}
let forced = forced_result(disposition);
let name = tc.name.clone();
let args = resolved.args_json.clone();
let call_id = tc.id.clone();
let recorder = recorder.clone();
async move {
if let Some(result) = forced {
result
} else {
run_and_redact(tools, recorder.as_ref(), call_id, name, args).await
}
}
})
.collect::<Vec<_>>();
let results = futures::future::join_all(futures).await;
ctx.executed_tools = true;
for ((_, tc), disposition) in unanswered.iter().zip(&dispositions) {
if matches!(disposition, CallDisposition::Execute) {
ctx.approved_remaining.remove(&(
tc.id.clone(),
tc.name.clone(),
canon_args(&tc.args_json),
));
}
}
let mut result_msgs = Vec::with_capacity(unanswered.len());
for ((_, tc), result) in unanswered.iter().zip(results) {
let result = cap_tool_result(&result);
let first_party = !ctx.tools.ingests_untrusted_content(&tc.name);
ctx.outputs
.push(tool_result_message(&tc.id, &result, first_party));
result_msgs.push(LlmMessage {
role: Role::Tool,
content: vec![LlmContent::tool_result(
tc.id.clone(),
result,
false,
first_party,
)],
});
}
let after = unanswered
.iter()
.map(|(idx, _)| *idx)
.max()
.unwrap_or(ctx.messages.len());
ctx.messages = splice_results_after(std::mem::take(&mut ctx.messages), after, result_msgs);
append_injected_notes(&mut ctx.outputs, &mut ctx.messages, &pre_resolutions);
if dispositions
.iter()
.any(|d| matches!(d, CallDisposition::Execute))
{
push_internal_note(
&mut ctx.outputs,
&mut ctx.messages,
RESUME_EXECUTED_GROUND_TRUTH_NOTE,
);
}
Ok(StepOutcome::Continue)
}
}
pub struct CircuitBreaker;
#[async_trait]
impl<P, T> TurnStep<P, T> for CircuitBreaker
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
async fn run(&self, ctx: &mut TurnCtx<'_, P, T>) -> Result<StepOutcome, P::Error> {
if ctx.saw_sig_match_denial {
ctx.denial_reprompts += 1;
if ctx.denial_reprompts >= crate::MAX_DENIAL_REPROMPTS {
tracing::warn!(
denial_reprompts = ctx.denial_reprompts,
max = crate::MAX_DENIAL_REPROMPTS,
"HITL circuit breaker: model re-emitted a denied action repeatedly; \
ending turn instead of re-prompting"
);
return Ok(StepOutcome::Done);
}
}
Ok(StepOutcome::Continue)
}
}