use std::collections::HashSet;
use crate::product::protocol::models::TranscriptItem;
use lha_llm::ToolCallPayload;
use lha_llm::ToolResultPayload;
use crate::product::agent::util::error_or_panic;
use tracing::info;
pub(crate) fn ensure_call_outputs_present(items: &mut Vec<TranscriptItem>) {
let mut missing_outputs_to_insert: Vec<(usize, TranscriptItem)> = Vec::new();
for (idx, item) in items.iter().enumerate() {
if let TranscriptItem::ToolCall {
call_id,
tool_name,
payload,
..
} = item
{
let has_output = items.iter().any(
|i| matches!(i, TranscriptItem::ToolResult { call_id: existing, .. } if existing == call_id),
);
if !has_output {
let result = match payload {
ToolCallPayload::JsonArguments { .. } => {
info!("Tool result is missing for call id: {call_id}");
TranscriptItem::ToolResult {
call_id: call_id.clone(),
tool_name: tool_name.clone(),
payload: ToolResultPayload::Structured {
content: "aborted".to_string(),
content_items: None,
success: None,
},
}
}
ToolCallPayload::TextInput { .. } => {
error_or_panic(format!(
"Custom tool call output is missing for call id: {call_id}"
));
TranscriptItem::ToolResult {
call_id: call_id.clone(),
tool_name: tool_name.clone(),
payload: ToolResultPayload::Text {
output: "aborted".to_string(),
},
}
}
};
missing_outputs_to_insert.push((idx, result));
}
}
}
for (idx, output_item) in missing_outputs_to_insert.into_iter().rev() {
items.insert(idx + 1, output_item);
}
}
pub(crate) fn remove_orphan_outputs(items: &mut Vec<TranscriptItem>) {
let tool_call_ids: HashSet<String> = items
.iter()
.filter_map(|i| match i {
TranscriptItem::ToolCall { call_id, .. } => Some(call_id.clone()),
_ => None,
})
.collect();
items.retain(|item| match item {
TranscriptItem::ToolResult { call_id, .. } => {
if call_id.is_empty() {
return false;
}
let has_match = tool_call_ids.contains(call_id);
if !has_match {
error_or_panic(format!("Orphan tool result for call id: {call_id}"));
}
has_match
}
_ => true,
});
}
pub(crate) fn remove_corresponding_for(items: &mut Vec<TranscriptItem>, item: &TranscriptItem) {
match item {
TranscriptItem::ToolCall { call_id, .. } => {
remove_first_matching(items, |i| {
matches!(
i,
TranscriptItem::ToolResult {
call_id: existing, ..
} if existing == call_id
)
});
}
TranscriptItem::ToolResult { call_id, .. } => {
remove_first_matching(items, |i| {
matches!(
i,
TranscriptItem::ToolCall {
call_id: existing, ..
} if existing == call_id
)
});
}
_ => {}
}
}
fn remove_first_matching<F>(items: &mut Vec<TranscriptItem>, predicate: F)
where
F: Fn(&TranscriptItem) -> bool,
{
if let Some(pos) = items.iter().position(predicate) {
items.remove(pos);
}
}