use std::collections::HashSet;
use serde_json::Value;
use crate::core::tool_calls::{
normalize_tool_calls, tool_call_arguments_string, tool_call_id, tool_call_name,
};
struct ToolResolution {
next_index: usize,
resolved_ids: HashSet<String>,
tool_result_blocks: Vec<Value>,
trailing_user_blocks: Vec<Value>,
}
pub fn split_system_and_conversation(messages_payload: &[Value]) -> (Vec<String>, Vec<Value>) {
let mut system_parts = Vec::new();
let mut conversation = Vec::new();
let mut index = 0;
while index < messages_payload.len() {
index = split_message(
messages_payload,
index,
&mut system_parts,
&mut conversation,
);
}
(system_parts, normalize_messages(conversation))
}
fn split_message(
messages_payload: &[Value],
index: usize,
system_parts: &mut Vec<String>,
conversation: &mut Vec<Value>,
) -> usize {
let msg = &messages_payload[index];
if collect_system_block(msg, system_parts) {
return index + 1;
}
if is_tool_assistant(msg) {
return process_tool_assistant(messages_payload, index, msg, conversation);
}
if message_role(msg) != "tool" {
push_conversation_entry(msg, conversation);
}
index + 1
}
fn collect_system_block(msg: &Value, system_parts: &mut Vec<String>) -> bool {
if !matches!(message_role(msg), "system" | "developer") {
return false;
}
if let Some(text) = extract_system_text(msg.get("content"))
&& !text.is_empty()
{
system_parts.push(text);
}
true
}
fn is_tool_assistant(msg: &Value) -> bool {
message_role(msg) == "assistant" && has_tool_calls(msg)
}
fn process_tool_assistant(
messages_payload: &[Value],
index: usize,
msg: &Value,
conversation: &mut Vec<Value>,
) -> usize {
let resolution = resolve_tool_results(messages_payload, index + 1, &tool_call_ids(msg));
if resolution.resolved_ids.is_empty() {
push_assistant_entry(msg, conversation, &HashSet::new());
return index + 1;
}
push_assistant_entry(msg, conversation, &resolution.resolved_ids);
push_tool_result_entry(
resolution.tool_result_blocks,
resolution.trailing_user_blocks,
conversation,
);
resolution.next_index
}
fn resolve_tool_results(
messages_payload: &[Value],
mut lookahead: usize,
call_ids: &HashSet<String>,
) -> ToolResolution {
let mut resolved_ids = HashSet::new();
let mut tool_result_blocks = Vec::new();
let mut trailing_user_blocks = Vec::new();
while lookahead < messages_payload.len() {
let next = &messages_payload[lookahead];
if message_role(next) == "assistant" {
break;
}
process_followup_message(
next,
call_ids,
&mut resolved_ids,
&mut tool_result_blocks,
&mut trailing_user_blocks,
);
lookahead += 1;
}
ToolResolution {
next_index: lookahead,
resolved_ids,
tool_result_blocks,
trailing_user_blocks,
}
}
fn process_followup_message(
msg: &Value,
call_ids: &HashSet<String>,
resolved_ids: &mut HashSet<String>,
tool_result_blocks: &mut Vec<Value>,
trailing_user_blocks: &mut Vec<Value>,
) {
match message_role(msg) {
"tool" => resolve_tool_result(msg, call_ids, resolved_ids, tool_result_blocks),
"system" | "developer" => {}
_ => trailing_user_blocks.extend(message_blocks(msg)),
}
}
fn resolve_tool_result(
msg: &Value,
call_ids: &HashSet<String>,
resolved_ids: &mut HashSet<String>,
tool_result_blocks: &mut Vec<Value>,
) {
if let Some((tool_use_id, block)) = tool_result_block(msg)
&& call_ids.contains(&tool_use_id)
{
resolved_ids.insert(tool_use_id);
tool_result_blocks.push(block);
}
}
fn push_assistant_entry(
msg: &Value,
conversation: &mut Vec<Value>,
allowed_tool_ids: &HashSet<String>,
) {
let blocks = assistant_content_blocks(msg, Some(allowed_tool_ids));
push_blocks("assistant", blocks, conversation);
}
fn push_tool_result_entry(
mut tool_result_blocks: Vec<Value>,
trailing_user_blocks: Vec<Value>,
conversation: &mut Vec<Value>,
) {
tool_result_blocks.extend(trailing_user_blocks);
push_blocks("user", tool_result_blocks, conversation);
}
fn push_conversation_entry(msg: &Value, conversation: &mut Vec<Value>) {
push_blocks(message_role(msg), message_blocks(msg), conversation);
}
fn push_blocks(role: &str, blocks: Vec<Value>, conversation: &mut Vec<Value>) {
if !blocks.is_empty() {
conversation.push(serde_json::json!({
"role": role,
"content": blocks,
}));
}
}
pub fn normalize_messages(messages: Vec<Value>) -> Vec<Value> {
if messages.is_empty() {
return vec![continue_message()];
}
let mut merged = merge_consecutive_roles(messages);
ensure_user_bookends(&mut merged);
merged
}
fn continue_message() -> Value {
serde_json::json!({"role": "user", "content": [{"type": "text", "text": "Continue."}]})
}
fn merge_consecutive_roles(messages: Vec<Value>) -> Vec<Value> {
let mut merged: Vec<Value> = Vec::with_capacity(messages.len());
for msg in messages {
let role = message_role(&msg);
let same_as_last = merged.last().is_some_and(|last| message_role(last) == role);
if same_as_last {
let last = merged
.last_mut()
.expect("SAFETY: same_as_last implies non-empty");
let existing = last.get("content").cloned().unwrap_or(Value::Null);
let new_content = msg.get("content").cloned().unwrap_or(Value::Null);
last["content"] = Value::Array(merge_blocks(existing, new_content));
} else {
merged.push(msg);
}
}
merged
}
fn ensure_user_bookends(merged: &mut Vec<Value>) {
if merged.first().map(message_role) != Some("user") {
merged.insert(0, continue_message());
}
if merged.last().map(message_role) != Some("user") {
merged.push(continue_message());
}
}
fn has_tool_calls(msg: &Value) -> bool {
msg.get("tool_calls")
.and_then(|value| value.as_array())
.map(|calls| !calls.is_empty())
.unwrap_or(false)
}
fn tool_call_ids(msg: &Value) -> HashSet<String> {
msg.get("tool_calls")
.and_then(|value| value.as_array())
.map(|calls| normalize_tool_calls(calls))
.unwrap_or_default()
.into_iter()
.filter_map(|call| tool_call_id(&call).map(ToOwned::to_owned))
.collect()
}
fn assistant_content_blocks(msg: &Value, allowed_tool_ids: Option<&HashSet<String>>) -> Vec<Value> {
let mut blocks = content_blocks(msg.get("content").cloned().unwrap_or(Value::Null));
if let Some(tool_calls) = msg.get("tool_calls").and_then(|value| value.as_array()) {
let tool_use_blocks = normalize_tool_calls(tool_calls)
.into_iter()
.filter_map(|tc| tool_call_to_tool_use_block(&tc, allowed_tool_ids));
blocks.extend(tool_use_blocks);
}
blocks
}
fn tool_call_to_tool_use_block(
tool_call: &Value,
allowed_tool_ids: Option<&HashSet<String>>,
) -> Option<Value> {
let id = tool_call_id(tool_call).unwrap_or_default();
if let Some(allowed) = allowed_tool_ids
&& !allowed.contains(id)
{
return None;
}
let name = tool_call_name(tool_call).unwrap_or_default();
let arguments = tool_call_arguments_string(tool_call);
let input = serde_json::from_str(&arguments).unwrap_or_else(|_| serde_json::json!({}));
Some(serde_json::json!({
"type": "tool_use",
"id": id,
"name": name,
"input": input,
}))
}
fn tool_result_block(msg: &Value) -> Option<(String, Value)> {
let tool_use_id = msg
.get("tool_call_id")
.and_then(|value| value.as_str())
.filter(|s| !s.is_empty())?;
let content = msg
.get("content")
.and_then(|value| value.as_str())
.unwrap_or_default();
let block = serde_json::json!({
"type": "tool_result",
"tool_use_id": tool_use_id,
"content": content,
});
Some((tool_use_id.to_owned(), block))
}
fn extract_system_text(content: Option<&Value>) -> Option<String> {
match content {
Some(Value::String(text)) => Some(text.clone()),
Some(Value::Array(items)) => {
let parts: Vec<&str> = items
.iter()
.filter(|item| item.get("type").and_then(|v| v.as_str()) == Some("text"))
.filter_map(|item| item.get("text").and_then(|v| v.as_str()))
.collect();
let joined = parts.join("\n\n");
(!joined.is_empty()).then_some(joined)
}
_ => None,
}
}
fn merge_blocks(existing: Value, new_content: Value) -> Vec<Value> {
let mut blocks = content_blocks(existing);
blocks.extend(content_blocks(new_content));
blocks
}
fn message_blocks(msg: &Value) -> Vec<Value> {
content_blocks(msg.get("content").cloned().unwrap_or(Value::Null))
}
fn content_blocks(content: Value) -> Vec<Value> {
match content {
Value::Null => Vec::new(),
Value::String(text) => {
if text.is_empty() {
Vec::new()
} else {
vec![serde_json::json!({"type": "text", "text": text})]
}
}
Value::Array(items) => items,
Value::Object(_) => vec![content],
other => vec![serde_json::json!({"type": "text", "text": other.to_string()})],
}
}
fn message_role(msg: &Value) -> &str {
msg.get("role")
.and_then(|value| value.as_str())
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_split_system_and_conversation_keeps_tool_results_adjacent() {
let messages = vec![
json!({"role": "system", "content": "system rules"}),
json!({"role": "user", "content": "find papers"}),
json!({
"role": "assistant",
"tool_calls": [
{"id": "toolu_1", "type": "function", "function": {"name": "bash", "arguments": "{\"cmd\":\"pwd\"}"}},
{"id": "toolu_2", "type": "function", "function": {"name": "skill", "arguments": "{\"name\":\"active-research\"}"}}
]
}),
json!({"role": "tool", "tool_call_id": "toolu_1", "content": "ok-1"}),
json!({"role": "tool", "tool_call_id": "toolu_2", "content": "ok-2"}),
json!({"role": "user", "content": "继续"}),
];
let (system_parts, conversation) = split_system_and_conversation(&messages);
assert_eq!(system_parts, vec!["system rules".to_owned()]);
assert_eq!(conversation.len(), 3);
assert_eq!(conversation[0]["role"], "user");
assert_eq!(conversation[1]["role"], "assistant");
assert_eq!(conversation[2]["role"], "user");
let assistant_blocks = conversation[1]["content"].as_array().unwrap();
assert_eq!(assistant_blocks.len(), 2);
assert_eq!(assistant_blocks[0]["type"], "tool_use");
assert_eq!(assistant_blocks[1]["type"], "tool_use");
let user_blocks = conversation[2]["content"].as_array().unwrap();
assert_eq!(user_blocks[0]["type"], "tool_result");
assert_eq!(user_blocks[1]["type"], "tool_result");
assert_eq!(user_blocks[2]["type"], "text");
assert_eq!(user_blocks[2]["text"], "继续");
}
#[test]
fn test_split_system_and_conversation_drops_unresolved_tool_use_blocks() {
let messages = vec![
json!({"role": "user", "content": "hello"}),
json!({
"role": "assistant",
"content": "I will inspect that.",
"tool_calls": [
{"id": "toolu_1", "type": "function", "function": {"name": "bash", "arguments": "{\"cmd\":\"pwd\"}"}}
]
}),
json!({"role": "user", "content": "继续"}),
];
let (_system_parts, conversation) = split_system_and_conversation(&messages);
assert_eq!(conversation.len(), 3);
assert_eq!(conversation[0]["role"], "user");
assert_eq!(conversation[1]["role"], "assistant");
assert_eq!(conversation[2]["role"], "user");
let assistant_blocks = conversation[1]["content"].as_array().unwrap();
assert_eq!(assistant_blocks.len(), 1);
assert_eq!(assistant_blocks[0]["type"], "text");
assert_eq!(conversation[2]["content"][0]["text"], "继续");
}
#[test]
fn test_normalize_messages_merges_tool_results_and_user_text() {
let messages = vec![
json!({
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "ok"}]
}),
json!({
"role": "user",
"content": [{"type": "text", "text": "continue"}]
}),
];
let normalized = normalize_messages(messages);
assert_eq!(normalized.len(), 1);
let blocks = normalized[0]["content"].as_array().unwrap();
assert_eq!(blocks[0]["type"], "tool_result");
assert_eq!(blocks[1]["type"], "text");
}
}