use std::collections::{HashMap, HashSet};
use crate::providers::{ChatMessage, ProviderConversationItem, ProviderToolResult, ToolCall};
use super::tool_continuation::{ToolCallFingerprint, function_call_item};
#[derive(Debug, Default)]
pub(super) struct AgentTurnState {
current_turn_items: Vec<ProviderConversationItem>,
current_assistant_segment: String,
current_assistant_replay_pending: String,
seen_tool_calls: HashSet<ToolCallFingerprint>,
seen_tool_call_ids: HashMap<String, ToolCallFingerprint>,
function_call_item_ids: HashSet<String>,
pending_segment_separator: bool,
iteration: usize,
}
impl AgentTurnState {
pub(super) fn request_items_slice(&self) -> &[ProviderConversationItem] {
&self.current_turn_items
}
pub(super) fn iteration(&self) -> usize {
self.iteration
}
pub(super) fn take_segment_separator(&mut self) -> bool {
std::mem::take(&mut self.pending_segment_separator)
}
pub(super) fn push_assistant_delta(&mut self, delta: &str) {
self.current_assistant_segment.push_str(delta);
self.current_assistant_replay_pending.push_str(delta);
}
pub(super) fn assistant_segment(&self) -> &str {
&self.current_assistant_segment
}
pub(super) fn push_response_item(&mut self, item: serde_json::Value) {
self.flush_pending_assistant_replay();
self.push_conversation_item(ProviderConversationItem::ResponseItem(item));
}
pub(super) fn prepare_tool_turn(&mut self) {
if !self.current_assistant_segment.trim().is_empty() {
self.flush_pending_assistant_replay();
}
}
pub(super) fn register_tool_call(&mut self, call: &ToolCall) -> anyhow::Result<()> {
let fingerprint = ToolCallFingerprint::from(call);
if !call.id.trim().is_empty() {
if let Some(existing) = self.seen_tool_call_ids.get(&call.id) {
if existing != &fingerprint {
anyhow::bail!("conflicting duplicate tool call id: {}", call.id);
}
} else {
self.seen_tool_call_ids
.insert(call.id.clone(), fingerprint.clone());
}
}
if !self.seen_tool_calls.insert(fingerprint) {
anyhow::bail!("duplicate tool call suppressed");
}
Ok(())
}
pub(super) fn append_function_call_if_missing(&mut self, call: &ToolCall) {
if !self.has_function_call_item(&call.id) {
self.push_conversation_item(ProviderConversationItem::ResponseItem(
function_call_item(call),
));
}
}
pub(super) fn append_tool_result(&mut self, provider_result: ProviderToolResult) {
self.push_conversation_item(ProviderConversationItem::ToolResult(provider_result));
}
pub(super) fn append_auto_continue(&mut self) {
self.flush_pending_assistant_replay();
self.push_conversation_item(ProviderConversationItem::Message(ChatMessage::user(
"Continue",
)));
}
pub(super) fn append_ttsr_system_reminder(&mut self, reminder: String) {
self.flush_pending_assistant_replay();
self.push_conversation_item(ProviderConversationItem::Message(ChatMessage::system(
reminder,
)));
}
pub(super) fn append_provider_context_items(
&mut self,
items: impl IntoIterator<Item = ProviderConversationItem>,
) {
self.flush_pending_assistant_replay();
let items = items.into_iter();
self.current_turn_items.reserve(items.size_hint().0);
for item in items {
self.push_conversation_item(item);
}
}
pub(super) fn finish_text_action_for_continuation(&mut self) {
self.flush_pending_assistant_replay();
self.current_assistant_segment.clear();
self.current_assistant_replay_pending.clear();
self.pending_segment_separator = true;
self.iteration += 1;
}
pub(super) fn finish_tool_iteration(&mut self) {
self.current_assistant_segment.clear();
self.current_assistant_replay_pending.clear();
self.pending_segment_separator = true;
self.iteration += 1;
}
fn flush_pending_assistant_replay(&mut self) {
if self.current_assistant_replay_pending.trim().is_empty() {
self.current_assistant_replay_pending.clear();
return;
}
let text = std::mem::take(&mut self.current_assistant_replay_pending);
self.push_conversation_item(ProviderConversationItem::Message(ChatMessage::assistant(
text,
)));
}
fn push_conversation_item(&mut self, item: ProviderConversationItem) {
self.track_function_call_item(&item);
self.current_turn_items.push(item);
}
fn track_function_call_item(&mut self, item: &ProviderConversationItem) {
let ProviderConversationItem::ResponseItem(value) = item else {
return;
};
if value.get("type").and_then(serde_json::Value::as_str) != Some("function_call") {
return;
}
if let Some(call_id) = value.get("call_id").and_then(serde_json::Value::as_str) {
self.function_call_item_ids.insert(call_id.to_string());
}
}
fn has_function_call_item(&self, call_id: &str) -> bool {
self.function_call_item_ids.contains(call_id)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn read_call(id: &str) -> ToolCall {
ToolCall {
id: id.to_string(),
name: "read".to_string(),
arguments: json!({"path":"file.txt"}),
}
}
#[test]
fn auto_continue_flushes_partial_assistant_before_user_continue() {
let mut state = AgentTurnState::default();
state.push_assistant_delta("partial answer");
state.append_auto_continue();
assert_eq!(state.current_assistant_segment, "partial answer");
assert_eq!(state.current_turn_items.len(), 2);
assert!(matches!(
&state.current_turn_items[0],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::Assistant
&& message.content == "partial answer"
));
assert!(matches!(
&state.current_turn_items[1],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User
&& message.content == "Continue"
));
}
#[test]
fn response_item_flushes_pending_assistant_replay_first() {
let mut state = AgentTurnState::default();
let item = json!({"type":"function_call", "call_id":"call_1"});
state.push_assistant_delta("I'll read.");
state.push_response_item(item.clone());
assert!(matches!(
&state.current_turn_items[0],
ProviderConversationItem::Message(message)
if message.content == "I'll read."
));
assert_eq!(
state.current_turn_items[1],
ProviderConversationItem::ResponseItem(item)
);
}
#[test]
fn duplicate_registration_blocks_replayed_tool_call_before_state_append() {
let mut state = AgentTurnState::default();
let call = read_call("call_1");
state.register_tool_call(&call).unwrap();
let error = state.register_tool_call(&call).unwrap_err().to_string();
assert_eq!(error, "duplicate tool call suppressed");
assert!(state.current_turn_items.is_empty());
}
#[test]
fn provider_context_items_append_after_tool_result() {
let mut state = AgentTurnState::default();
let call = read_call("call_1");
state.append_function_call_if_missing(&call);
state.append_tool_result(ProviderToolResult {
call_id: "call_1".to_string(),
tool_name: "read".to_string(),
success: true,
output: "file text".to_string(),
});
state.append_provider_context_items(vec![ProviderConversationItem::Message(
ChatMessage::user("injected"),
)]);
assert!(matches!(
&state.current_turn_items[1],
ProviderConversationItem::ToolResult(result) if result.output == "file text"
));
assert!(matches!(
&state.current_turn_items[2],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User
&& message.content == "injected"
));
}
#[test]
fn provider_context_items_flush_pending_assistant_replay_first() {
let mut state = AgentTurnState::default();
state.push_assistant_delta("first");
state.append_provider_context_items([ProviderConversationItem::Message(
ChatMessage::user("injected"),
)]);
assert!(matches!(
&state.current_turn_items[0],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::Assistant
&& message.content == "first"
));
assert!(matches!(
&state.current_turn_items[1],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User
&& message.content == "injected"
));
}
#[test]
fn finish_text_action_for_continuation_flushes_replay_and_resets_segment() {
let mut state = AgentTurnState::default();
state.push_assistant_delta("first");
state.finish_text_action_for_continuation();
state.append_provider_context_items([ProviderConversationItem::Message(
ChatMessage::user(
"Steering update from user while current run was active:\n\nstay concise"
.to_string(),
),
)]);
assert_eq!(state.current_assistant_segment, "");
assert_eq!(state.current_assistant_replay_pending, "");
assert!(state.pending_segment_separator);
assert_eq!(state.iteration, 1);
assert!(matches!(
&state.current_turn_items[0],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::Assistant
&& message.content == "first"
));
assert!(matches!(
&state.current_turn_items[1],
ProviderConversationItem::Message(message)
if message.role == crate::providers::MessageRole::User
&& message.content.starts_with("Steering update from user")
));
}
#[test]
fn existing_provider_function_call_item_is_not_duplicated() {
let mut state = AgentTurnState::default();
let call = read_call("call_1");
let item = function_call_item(&call);
state.push_response_item(item);
state.append_function_call_if_missing(&call);
let function_call_count = state
.current_turn_items
.iter()
.filter(|item| matches!(item, ProviderConversationItem::ResponseItem(value) if value["call_id"] == "call_1"))
.count();
assert_eq!(function_call_count, 1);
}
}