use async_trait::async_trait;
use buffa_types::google::protobuf::Struct;
use polyc_llm::request::ToolCall;
use polyc_llm::{
CompletionRequest, Content as LlmContent, LlmProvider, Message as LlmMessage, Role, StopReason,
ToolSpec, Usage,
turn::{collect_turn, collect_turn_observed},
};
use polyc_proto::proto::polychrome::agent::v1::{
Content, FunctionCallContent, FunctionResultContent, Message, TextContent, ToolCallContent,
ToolResultContent, content, function_result_content, thought_summary_content,
tool_call_content, tool_result_content,
};
pub mod handoff;
pub mod llm_summarizer;
pub mod participation;
pub use handoff::{
DEFAULT_MAX_CARRY, HANDOFF_TOOL_NAME, HandoffRequest, handoff_tool_spec, parse_handoff_args,
};
pub use llm_summarizer::LlmSummarizer;
pub use polyc_llm::turn::TurnStreamEvent;
#[must_use]
pub const fn llm_stop_to_wire_i32(stop: StopReason) -> i32 {
use polyc_proto::proto::polychrome::harness::v1::StopReason as Wire;
match stop {
StopReason::EndTurn => Wire::STOP_REASON_END_TURN as i32,
StopReason::ToolUse => Wire::STOP_REASON_TOOL_USE as i32,
StopReason::MaxTokens => Wire::STOP_REASON_MAX_TOKENS as i32,
StopReason::Refusal => Wire::STOP_REASON_REFUSAL as i32,
StopReason::StopSequence => Wire::STOP_REASON_STOP_SEQUENCE as i32,
_ => Wire::STOP_REASON_UNSPECIFIED as i32,
}
}
#[must_use]
pub const fn wire_to_llm_stop(wire: i32) -> Option<StopReason> {
use polyc_proto::proto::polychrome::harness::v1::StopReason as Wire;
match wire {
x if x == Wire::STOP_REASON_END_TURN as i32 => Some(StopReason::EndTurn),
x if x == Wire::STOP_REASON_TOOL_USE as i32 => Some(StopReason::ToolUse),
x if x == Wire::STOP_REASON_MAX_TOKENS as i32 => Some(StopReason::MaxTokens),
x if x == Wire::STOP_REASON_REFUSAL as i32 => Some(StopReason::Refusal),
x if x == Wire::STOP_REASON_STOP_SEQUENCE as i32 => Some(StopReason::StopSequence),
_ => None,
}
}
#[async_trait]
pub trait Summarizer: Send + Sync {
async fn summarize(&self, prior_summary: &str, transcript: &[LlmMessage]) -> String;
}
pub const SUMMARY_TRIGGER: usize = 40;
#[derive(Clone, Copy, Default)]
pub struct StubSummarizer;
#[async_trait]
impl Summarizer for StubSummarizer {
async fn summarize(&self, prior_summary: &str, transcript: &[LlmMessage]) -> String {
let head = transcript
.iter()
.take(2)
.filter_map(|m| match m.content.first() {
Some(LlmContent::Text(t)) => Some(format!("{:?}: {}", m.role, snippet(t, 80))),
_ => None,
})
.collect::<Vec<_>>()
.join("; ");
let tail = transcript
.iter()
.rev()
.take(2)
.rev()
.filter_map(|m| match m.content.first() {
Some(LlmContent::Text(t)) => Some(format!("{:?}: {}", m.role, snippet(t, 80))),
_ => None,
})
.collect::<Vec<_>>()
.join("; ");
let count = transcript.len();
if prior_summary.is_empty() {
format!("[summary: {count} prior messages — head: {head}; tail: {tail}]")
} else {
format!("{prior_summary}\n[+{count} messages: head: {head}; tail: {tail}]")
}
}
}
fn snippet(s: &str, max: usize) -> String {
if s.len() <= max {
return s.to_owned();
}
let mut end = max;
while !s.is_char_boundary(end) && end > 0 {
end -= 1;
}
format!("{}…", &s[..end])
}
#[async_trait]
pub trait ToolExecutor: Send + Sync {
fn specs(&self) -> Vec<ToolSpec> {
Vec::new()
}
fn owns(&self, name: &str) -> bool {
self.specs().iter().any(|s| s.name == name)
}
fn needs_approval(&self, _name: &str) -> bool {
false
}
async fn execute(&self, name: &str, args_json: &str) -> String;
}
#[derive(Clone, Copy, Default)]
pub struct StubTools;
#[async_trait]
impl ToolExecutor for StubTools {
async fn execute(&self, name: &str, args_json: &str) -> String {
format!(r#"{{"unhandled_tool":"{name}","args":{args_json}}}"#)
}
}
const MAX_STEPS: usize = 8;
const MAX_DENIAL_REPROMPTS: usize = 2;
const DENIAL_RESULT_JSON: &str = r#"{"approved":false,"error":"denied by human approver"}"#;
#[derive(Debug, Clone)]
pub struct PendingApproval {
pub id: String,
pub name: String,
pub args_json: String,
pub title: String,
}
#[derive(Debug, Default, Clone)]
pub struct TurnResult {
pub messages: Vec<Message>,
pub usage: Usage,
pub stop: Option<StopReason>,
pub pending_approvals: Vec<PendingApproval>,
pub handoff: Option<HandoffRequest>,
}
#[derive(Debug, Default, Clone)]
pub struct RunTurnOptions {
pub approved_call_ids: std::collections::HashSet<(String, String, String)>,
pub denied_call_ids: std::collections::HashSet<(String, String, String)>,
pub stream_tx: Option<futures::channel::mpsc::UnboundedSender<TurnStreamEvent>>,
pub web_search: bool,
}
tokio::task_local! {
static CURRENT_TOOL_CALL_ID: String;
}
#[must_use]
pub fn current_tool_call_id() -> Option<String> {
CURRENT_TOOL_CALL_ID.try_with(String::clone).ok()
}
pub async fn with_tool_call_id<F>(id: String, fut: F) -> F::Output
where
F: std::future::Future,
{
CURRENT_TOOL_CALL_ID.scope(id, fut).await
}
pub async fn run_turn<P, T>(
provider: &P,
tools: &T,
model: &str,
messages: Vec<LlmMessage>,
) -> Result<TurnResult, P::Error>
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
run_turn_with(provider, tools, model, messages, RunTurnOptions::default()).await
}
enum CallDisposition {
Pending,
Denied { sig_match: bool },
Execute,
}
#[allow(clippy::too_many_lines)] #[tracing::instrument(skip_all, fields(model = model, input_messages = messages.len(), approved = options.approved_call_ids.len(), denied = options.denied_call_ids.len()))]
pub async fn run_turn_with<P, T>(
provider: &P,
tools: &T,
model: &str,
mut messages: Vec<LlmMessage>,
options: RunTurnOptions,
) -> Result<TurnResult, P::Error>
where
P: LlmProvider + ?Sized,
T: ToolExecutor + ?Sized,
{
let mut outputs = Vec::new();
let mut total_usage = Usage::default();
let mut last_stop: Option<StopReason> = None;
let mut pending_handoff: Option<HandoffRequest> = None;
let mut denied_sigs: std::collections::HashSet<(String, String)> =
std::collections::HashSet::new();
let mut denial_reprompts: usize = 0;
for _ in 0..MAX_STEPS {
let mut tool_specs = tools.specs();
if !tool_specs.iter().any(|s| s.name == HANDOFF_TOOL_NAME) {
tool_specs.push(handoff_tool_spec());
}
let mut req = CompletionRequest::new(model);
req.messages.clone_from(&messages);
req.tools = tool_specs.clone();
req.web_search = options.web_search;
let stream = provider.complete(req).await?;
let turn = if let Some(tx) = options.stream_tx.clone() {
collect_turn_observed(stream, move |ev| {
let _ = tx.unbounded_send(ev);
})
.await?
} else {
collect_turn(stream).await?
};
total_usage.input_tokens += turn.usage.input_tokens;
total_usage.output_tokens += turn.usage.output_tokens;
last_stop = turn.stop;
if !turn.text.is_empty() {
outputs.push(text_message("model", &turn.text));
}
for tc in &turn.tool_calls {
outputs.push(tool_call_message(tc));
}
let mut assistant = LlmMessage::assistant(turn.text.clone());
for tc in &turn.tool_calls {
assistant.content.push(LlmContent::tool_use_signed(
tc.id.clone(),
tc.name.clone(),
tc.args_json.clone(),
tc.signature.clone(),
));
}
messages.push(assistant);
let wants_tools = !turn.tool_calls.is_empty()
&& !matches!(turn.stop, Some(StopReason::MaxTokens | StopReason::Refusal));
if !wants_tools {
break;
}
if let Some(tc) = turn.tool_calls.iter().find(|t| t.name == HANDOFF_TOOL_NAME)
&& let Some(req) = handoff::parse_handoff_args(&tc.id, &tc.args_json, &messages)
{
pending_handoff = Some(req);
break;
}
let dispositions = turn
.tool_calls
.iter()
.map(|tc| {
let needs_approval = tools.needs_approval(&tc.name);
let sig = (tc.name.clone(), tc.args_json.clone());
let sig_denied = denied_sigs.contains(&sig);
let approval_key = (tc.id.clone(), tc.name.clone(), tc.args_json.clone());
let is_denied = options.denied_call_ids.contains(&approval_key) || sig_denied;
let is_approved = options.approved_call_ids.contains(&approval_key);
if needs_approval && is_denied {
CallDisposition::Denied {
sig_match: sig_denied,
}
} else if needs_approval && !is_approved {
CallDisposition::Pending
} else {
CallDisposition::Execute
}
})
.collect::<Vec<_>>();
let batch_needs_approval = dispositions
.iter()
.any(|d| matches!(d, CallDisposition::Pending));
if batch_needs_approval {
let pending = turn
.tool_calls
.iter()
.zip(&dispositions)
.filter(|(_, d)| matches!(d, CallDisposition::Pending))
.map(|(tc, _)| {
let title = tool_specs
.iter()
.find(|s| s.name == tc.name)
.and_then(|s| s.title.clone())
.unwrap_or_default();
PendingApproval {
id: tc.id.clone(),
name: tc.name.clone(),
args_json: tc.args_json.clone(),
title,
}
})
.collect::<Vec<_>>();
return Ok(TurnResult {
messages: outputs,
usage: total_usage,
stop: last_stop,
pending_approvals: pending,
handoff: None,
});
}
let mut saw_sig_match_denial = false;
let tool_futures = turn
.tool_calls
.iter()
.zip(&dispositions)
.map(|(tc, disposition)| {
let denied = matches!(disposition, CallDisposition::Denied { .. });
if let CallDisposition::Denied { sig_match } = disposition {
denied_sigs.insert((tc.name.clone(), tc.args_json.clone()));
if *sig_match {
saw_sig_match_denial = true;
}
}
let name = tc.name.clone();
let args = tc.args_json.clone();
let call_id = tc.id.clone();
async move {
if denied {
DENIAL_RESULT_JSON.to_owned()
} else {
CURRENT_TOOL_CALL_ID
.scope(call_id, tools.execute(&name, &args))
.await
}
}
})
.collect::<Vec<_>>();
let results = futures::future::join_all(tool_futures).await;
for (tc, result) in turn.tool_calls.iter().zip(results) {
outputs.push(tool_result_message(&tc.id, &result));
messages.push(LlmMessage {
role: Role::Tool,
content: vec![LlmContent::tool_result(tc.id.clone(), result, false)],
});
}
if saw_sig_match_denial {
denial_reprompts += 1;
if denial_reprompts >= MAX_DENIAL_REPROMPTS {
tracing::warn!(
denial_reprompts,
max = MAX_DENIAL_REPROMPTS,
"HITL circuit breaker: model re-emitted a denied action repeatedly; \
ending turn instead of re-prompting"
);
break;
}
}
}
Ok(TurnResult {
messages: outputs,
usage: total_usage,
stop: last_stop,
pending_approvals: Vec::new(),
handoff: pending_handoff,
})
}
#[must_use]
pub fn llm_to_wire(msg: &LlmMessage) -> Message {
let role = match msg.role {
Role::Assistant => "model",
Role::Tool => "tool",
Role::System => "system",
_ => "user",
};
let text = msg
.content
.iter()
.filter_map(|c| match c {
LlmContent::Text(s) => Some(s.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("");
text_message(role, &text)
}
#[must_use]
pub fn wire_to_llm(msg: &Message) -> LlmMessage {
let role = match msg.role.as_str() {
"model" | "assistant" => Role::Assistant,
"tool" | "function" => Role::Tool,
"system" => Role::System,
_ => Role::User,
};
let content = match msg.content.as_option().and_then(|c| c.r#type.as_ref()) {
Some(content::Type::Text(t)) => vec![LlmContent::Text(t.text.clone())],
Some(content::Type::ToolCall(tc)) => {
let (name, args_json) = match tc.r#type.as_ref() {
Some(tool_call_content::Type::FunctionCall(fc)) => {
let args_json = fc
.arguments
.as_option()
.and_then(|s| serde_json::to_string(s).ok())
.unwrap_or_else(|| "{}".to_owned());
(fc.name.clone(), args_json)
}
None => (String::new(), "{}".to_owned()),
};
let signature = (!tc.signature.is_empty())
.then(|| String::from_utf8_lossy(&tc.signature).into_owned());
vec![LlmContent::tool_use_signed(
tc.id.clone(),
name,
args_json,
signature,
)]
}
Some(content::Type::ToolResult(tr)) => {
let result_json = match tr.r#type.as_ref() {
Some(tool_result_content::Type::FunctionResult(fr)) => match fr.result.as_ref() {
Some(function_result_content::Result::Response(resp)) => {
serde_json::to_string(resp).unwrap_or_else(|_| "{}".to_owned())
}
None => "{}".to_owned(),
},
None => "{}".to_owned(),
};
vec![LlmContent::tool_result(
tr.call_id.clone(),
result_json,
false,
)]
}
Some(content::Type::Thought(t)) => {
let mut buf = String::new();
for s in &t.summary {
if let Some(thought_summary_content::Type::Text(text)) = s.r#type.as_ref()
&& !text.text.is_empty()
{
if !buf.is_empty() {
buf.push(' ');
}
buf.push_str(&text.text);
}
}
if buf.is_empty() {
Vec::new()
} else {
vec![LlmContent::Text(buf)]
}
}
_ => Vec::new(),
};
LlmMessage { role, content }
}
#[must_use]
pub fn tool_call_message(tc: &ToolCall) -> Message {
let arguments = serde_json::from_str::<Struct>(&tc.args_json)
.map(buffa::MessageField::some)
.unwrap_or_default();
Message {
role: "model".to_owned(),
content: buffa::MessageField::some(Content {
r#type: Some(content::Type::ToolCall(Box::new(ToolCallContent {
id: tc.id.clone(),
signature: tc
.signature
.clone()
.map(String::into_bytes)
.unwrap_or_default(),
r#type: Some(tool_call_content::Type::FunctionCall(Box::new(
FunctionCallContent {
name: tc.name.clone(),
arguments,
..Default::default()
},
))),
..Default::default()
}))),
..Default::default()
}),
internal_only: false,
..Default::default()
}
}
#[must_use]
pub fn tool_result_message(call_id: &str, result_json: &str) -> Message {
let response = serde_json::from_str::<Struct>(result_json)
.ok()
.map(|s| function_result_content::Result::Response(Box::new(s)));
Message {
role: "tool".to_owned(),
content: buffa::MessageField::some(Content {
r#type: Some(content::Type::ToolResult(Box::new(ToolResultContent {
call_id: call_id.to_owned(),
r#type: Some(tool_result_content::Type::FunctionResult(Box::new(
FunctionResultContent {
result: response,
..Default::default()
},
))),
..Default::default()
}))),
..Default::default()
}),
internal_only: false,
..Default::default()
}
}
#[must_use]
pub fn text_message(role: &str, text: &str) -> Message {
Message {
role: role.to_owned(),
content: buffa::MessageField::some(Content {
r#type: Some(content::Type::Text(Box::new(TextContent {
text: text.to_owned(),
..Default::default()
}))),
..Default::default()
}),
internal_only: false,
..Default::default()
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use futures::{StreamExt, stream};
use polyc_llm::{Chunk, CompletionRequest, LlmProvider, error::DummyError, turn::StubProvider};
use std::sync::atomic::{AtomicUsize, Ordering};
use super::*;
#[tokio::test]
async fn stub_turn_yields_one_assistant_message() {
let out = run_turn(
&StubProvider,
&StubTools,
"stub",
vec![LlmMessage::user("hi")],
)
.await
.expect("turn");
assert_eq!(out.messages.len(), 1);
assert_eq!(out.messages[0].role, "model");
assert!(out.pending_approvals.is_empty());
}
struct ScriptedToolCallProvider {
calls: AtomicUsize,
}
#[async_trait]
impl LlmProvider for ScriptedToolCallProvider {
type Error = DummyError;
async fn complete(
&self,
_req: CompletionRequest,
) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
{
let n = self.calls.fetch_add(1, Ordering::SeqCst);
let chunks = if n == 0 {
vec![
Ok(Chunk::tool_call_start("call-1", "dangerous_tool")),
Ok(Chunk::tool_call_args_delta("call-1", r#"{"rm":"-rf"}"#)),
Ok(Chunk::tool_call_end("call-1")),
Ok(Chunk::Stop(polyc_llm::StopReason::ToolUse)),
]
} else {
vec![
Ok(Chunk::text_delta("done")),
Ok(Chunk::Stop(polyc_llm::StopReason::EndTurn)),
]
};
Ok(stream::iter(chunks).boxed())
}
}
#[derive(Default)]
struct ApprovalGatedTools {
executed: std::sync::Mutex<Vec<String>>,
}
#[async_trait]
impl ToolExecutor for ApprovalGatedTools {
fn needs_approval(&self, name: &str) -> bool {
name == "dangerous_tool"
}
async fn execute(&self, name: &str, args_json: &str) -> String {
self.executed.lock().unwrap().push(name.to_owned());
format!(r#"{{"ran":"{name}","args":{args_json}}}"#)
}
}
#[tokio::test]
async fn needs_approval_tool_pauses_with_pending_approval() {
let provider = ScriptedToolCallProvider {
calls: AtomicUsize::new(0),
};
let tools = ApprovalGatedTools::default();
let out = run_turn(&provider, &tools, "scripted", vec![LlmMessage::user("hi")])
.await
.expect("turn");
assert_eq!(
out.pending_approvals.len(),
1,
"needs_approval tool short-circuits the loop"
);
let pa = &out.pending_approvals[0];
assert_eq!(pa.id, "call-1");
assert_eq!(pa.name, "dangerous_tool");
assert_eq!(pa.args_json, r#"{"rm":"-rf"}"#);
assert!(
tools.executed.lock().unwrap().is_empty(),
"execute() must not be called when needs_approval=true"
);
}
#[tokio::test]
async fn pending_approval_default_is_empty() {
let out = run_turn(
&StubProvider,
&StubTools,
"stub",
vec![LlmMessage::user("hi")],
)
.await
.expect("turn");
assert!(out.pending_approvals.is_empty());
}
#[derive(Default)]
struct ReadOnlyTools;
#[async_trait]
impl ToolExecutor for ReadOnlyTools {
async fn execute(&self, _name: &str, _args_json: &str) -> String {
r#"{"result":"ok"}"#.to_owned()
}
}
struct ScriptedBenignProvider {
calls: AtomicUsize,
}
#[async_trait]
impl LlmProvider for ScriptedBenignProvider {
type Error = DummyError;
async fn complete(
&self,
_req: CompletionRequest,
) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
{
let n = self.calls.fetch_add(1, Ordering::SeqCst);
let chunks = if n == 0 {
vec![
Ok(Chunk::tool_call_start("call-1", "read_only")),
Ok(Chunk::tool_call_args_delta("call-1", "{}")),
Ok(Chunk::tool_call_end("call-1")),
Ok(Chunk::Stop(polyc_llm::StopReason::ToolUse)),
]
} else {
vec![
Ok(Chunk::text_delta("done")),
Ok(Chunk::Stop(polyc_llm::StopReason::EndTurn)),
]
};
Ok(stream::iter(chunks).boxed())
}
}
#[tokio::test]
async fn previously_approved_tool_executes_on_resume() {
let provider = ScriptedToolCallProvider {
calls: AtomicUsize::new(0),
};
let tools = ApprovalGatedTools::default();
let mut approved = std::collections::HashSet::new();
approved.insert((
"call-1".to_owned(),
"dangerous_tool".to_owned(),
r#"{"rm":"-rf"}"#.to_owned(),
));
let out = run_turn_with(
&provider,
&tools,
"scripted",
vec![LlmMessage::user("hi")],
RunTurnOptions {
approved_call_ids: approved,
..Default::default()
},
)
.await
.expect("turn");
assert!(
out.pending_approvals.is_empty(),
"approved call must NOT re-pause the loop"
);
let executed = tools.executed.lock().unwrap().clone();
assert_eq!(
executed,
vec!["dangerous_tool".to_owned()],
"tool executes after approval lands"
);
}
#[tokio::test]
async fn approval_does_not_inherit_across_changed_args() {
let provider = ScriptedToolCallProvider {
calls: AtomicUsize::new(0),
};
let tools = ApprovalGatedTools::default();
let mut approved = std::collections::HashSet::new();
approved.insert((
"call-1".to_owned(),
"dangerous_tool".to_owned(),
r#"{"rm":"/tmp/safe"}"#.to_owned(),
));
let out = run_turn_with(
&provider,
&tools,
"scripted",
vec![LlmMessage::user("hi")],
RunTurnOptions {
approved_call_ids: approved,
..Default::default()
},
)
.await
.expect("turn");
assert_eq!(
out.pending_approvals.len(),
1,
"an approval for different args must NOT authorize this call — it re-pauses"
);
assert!(
tools.executed.lock().unwrap().is_empty(),
"the tool must NOT execute under a mismatched-args approval"
);
}
#[tokio::test]
async fn denied_tool_resolves_without_executing_or_repausing() {
let provider = ScriptedToolCallProvider {
calls: AtomicUsize::new(0),
};
let tools = ApprovalGatedTools::default();
let mut denied = std::collections::HashSet::new();
denied.insert((
"call-1".to_owned(),
"dangerous_tool".to_owned(),
r#"{"rm":"-rf"}"#.to_owned(),
));
let out = run_turn_with(
&provider,
&tools,
"scripted",
vec![LlmMessage::user("hi")],
RunTurnOptions {
denied_call_ids: denied,
..Default::default()
},
)
.await
.expect("turn");
assert!(
out.pending_approvals.is_empty(),
"denied call must NOT re-pause the loop"
);
assert_eq!(
provider.calls.load(Ordering::SeqCst),
2,
"first signed denial must not trip the breaker; model ends the turn itself"
);
assert!(
tools.executed.lock().unwrap().is_empty(),
"execute() must not be called for a denied call"
);
let denial = out
.messages
.iter()
.find(|m| {
m.role == "tool"
&& matches!(
m.content.as_option().and_then(|c| c.r#type.as_ref()),
Some(content::Type::ToolResult(tr)) if tr.call_id == "call-1"
)
})
.expect("denied call must produce a tool_result message");
let llm = wire_to_llm(denial);
match &llm.content[0] {
LlmContent::ToolResult(tr) => {
let parsed: serde_json::Value =
serde_json::from_str(&tr.result_json).expect("denial result is valid json");
assert_eq!(
parsed.get("approved"),
Some(&serde_json::Value::Bool(false)),
"denial result must carry approved=false"
);
assert!(
parsed.get("error").is_some(),
"denial result must carry an error explanation"
);
}
other => panic!("expected ToolResult, got {other:?}"),
}
}
struct ReEmittingDeniedProvider {
calls: AtomicUsize,
}
#[async_trait]
impl LlmProvider for ReEmittingDeniedProvider {
type Error = DummyError;
async fn complete(
&self,
_req: CompletionRequest,
) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
{
let n = self.calls.fetch_add(1, Ordering::SeqCst);
let id = format!("call-{}", n + 1);
let chunks = vec![
Ok(Chunk::tool_call_start(&id, "dangerous_tool")),
Ok(Chunk::tool_call_args_delta(&id, r#"{"rm":"-rf"}"#)),
Ok(Chunk::tool_call_end(&id)),
Ok(Chunk::Stop(polyc_llm::StopReason::ToolUse)),
];
Ok(stream::iter(chunks).boxed())
}
}
#[tokio::test]
async fn reemitted_denied_signature_is_auto_denied_and_circuit_breaker_bounds_loop() {
let provider = ReEmittingDeniedProvider {
calls: AtomicUsize::new(0),
};
let tools = ApprovalGatedTools::default();
let mut denied = std::collections::HashSet::new();
denied.insert((
"call-1".to_owned(),
"dangerous_tool".to_owned(),
r#"{"rm":"-rf"}"#.to_owned(),
));
let out = run_turn_with(
&provider,
&tools,
"scripted",
vec![LlmMessage::user("hi")],
RunTurnOptions {
denied_call_ids: denied,
..Default::default()
},
)
.await
.expect("turn");
assert!(
out.pending_approvals.is_empty(),
"re-emitted denied signature must auto-deny, not re-prompt"
);
assert!(
tools.executed.lock().unwrap().is_empty(),
"auto-denied calls must never execute"
);
let denial_results = out
.messages
.iter()
.filter(|m| {
m.role == "tool"
&& matches!(
m.content.as_option().and_then(|c| c.r#type.as_ref()),
Some(content::Type::ToolResult(_))
)
})
.count();
assert!(
denial_results >= 1,
"each auto-denied call must still produce a tool_result"
);
let driven = provider.calls.load(Ordering::SeqCst);
assert!(
driven <= MAX_DENIAL_REPROMPTS + 1,
"circuit breaker must bound re-prompts: driven={driven} > {}",
MAX_DENIAL_REPROMPTS + 1
);
assert!(
driven < MAX_STEPS,
"circuit breaker must end the turn before burning MAX_STEPS"
);
}
#[tokio::test]
async fn read_only_batch_runs_through_without_approval_pause() {
let provider = ScriptedBenignProvider {
calls: AtomicUsize::new(0),
};
let tools = ReadOnlyTools;
let out = run_turn(&provider, &tools, "scripted", vec![LlmMessage::user("hi")])
.await
.expect("turn");
assert!(
out.pending_approvals.is_empty(),
"no approval needed for read-only tools"
);
assert!(out.messages.iter().any(|m| m.role == "tool"));
}
#[test]
fn wire_to_llm_preserves_tool_call_and_result() {
use buffa::MessageField;
use buffa_types::google::protobuf::Struct;
use polyc_proto::proto::polychrome::agent::v1::{
FunctionCallContent, FunctionResultContent, ToolCallContent, ToolResultContent,
};
fn wire(role: &str, ty: content::Type) -> Message {
Message {
role: role.to_owned(),
content: MessageField::some(Content {
r#type: Some(ty),
..Default::default()
}),
internal_only: false,
..Default::default()
}
}
let args: Struct = serde_json::from_str(r#"{"query":"rust"}"#).expect("args struct");
let call = wire(
"model",
content::Type::ToolCall(Box::new(ToolCallContent {
id: "call_1".to_owned(),
r#type: Some(tool_call_content::Type::FunctionCall(Box::new(
FunctionCallContent {
name: "search".to_owned(),
arguments: MessageField::some(args),
..Default::default()
},
))),
..Default::default()
})),
);
let llm_call = wire_to_llm(&call);
assert_eq!(llm_call.role, Role::Assistant);
assert_eq!(llm_call.content.len(), 1);
match &llm_call.content[0] {
LlmContent::ToolUse(tc) => {
assert_eq!(tc.id, "call_1");
assert_eq!(tc.name, "search", "function name must survive");
let parsed: serde_json::Value =
serde_json::from_str(&tc.args_json).expect("args_json is valid json");
assert_eq!(
parsed,
serde_json::json!({ "query": "rust" }),
"args must survive, not a placeholder"
);
}
other => panic!("expected ToolUse, got {other:?}"),
}
let resp: Struct = serde_json::from_str(r#"{"answer":42}"#).expect("resp struct");
let result = wire(
"tool",
content::Type::ToolResult(Box::new(ToolResultContent {
call_id: "call_1".to_owned(),
r#type: Some(tool_result_content::Type::FunctionResult(Box::new(
FunctionResultContent {
name: "search".to_owned(),
result: Some(function_result_content::Result::Response(Box::new(resp))),
..Default::default()
},
))),
..Default::default()
})),
);
let llm_result = wire_to_llm(&result);
assert_eq!(llm_result.role, Role::Tool);
assert_eq!(llm_result.content.len(), 1);
match &llm_result.content[0] {
LlmContent::ToolResult(tr) => {
assert_eq!(tr.tool_call_id, "call_1", "call id must correlate");
assert!(!tr.is_error);
let parsed: serde_json::Value =
serde_json::from_str(&tr.result_json).expect("result_json is valid json");
assert_eq!(
parsed,
serde_json::json!({ "answer": 42.0 }),
"result payload must survive, not a placeholder"
);
}
other => panic!("expected ToolResult, got {other:?}"),
}
}
#[test]
fn tool_call_message_round_trips_through_wire_to_llm_with_signature() {
let tc = ToolCall {
id: "call-7".to_owned(),
name: "search".to_owned(),
args_json: r#"{"query":"rust"}"#.to_owned(),
signature: Some("sig-abc123".to_owned()),
};
let wire = tool_call_message(&tc);
assert_eq!(wire.role, "model");
let back = wire_to_llm(&wire);
match &back.content[0] {
LlmContent::ToolUse(rt) => {
assert_eq!(rt.id, "call-7");
assert_eq!(rt.name, "search");
let parsed: serde_json::Value = serde_json::from_str(&rt.args_json).unwrap();
assert_eq!(parsed, serde_json::json!({ "query": "rust" }));
assert_eq!(
rt.signature.as_deref(),
Some("sig-abc123"),
"thought signature must survive the wire round-trip"
);
}
other => panic!("expected ToolUse, got {other:?}"),
}
}
#[test]
fn tool_result_message_round_trips_through_wire_to_llm() {
let wire = tool_result_message("call-7", r#"{"answer":42}"#);
assert_eq!(wire.role, "tool");
let back = wire_to_llm(&wire);
match &back.content[0] {
LlmContent::ToolResult(tr) => {
assert_eq!(tr.tool_call_id, "call-7");
let parsed: serde_json::Value = serde_json::from_str(&tr.result_json).unwrap();
assert_eq!(parsed, serde_json::json!({ "answer": 42.0 }));
}
other => panic!("expected ToolResult, got {other:?}"),
}
}
}