use crate::core::ToolInvocation;
use serde::Deserialize;
use serde_json::Value;
use std::collections::HashMap;
#[derive(Debug, Deserialize)]
pub(crate) struct UsageRecord {
pub(crate) input_tokens: Option<i64>,
pub(crate) output_tokens: Option<i64>,
pub(crate) cache_creation_input_tokens: Option<i64>,
pub(crate) cache_read_input_tokens: Option<i64>,
}
#[derive(Debug, Deserialize)]
struct Message {
content: Option<Value>,
}
#[derive(Debug, Deserialize)]
pub(crate) struct TranscriptRecord {
#[serde(rename = "type")]
pub(crate) record_type: Option<String>,
message: Option<Message>,
}
fn content_blocks(message: &Option<Message>) -> &[Value] {
match message.as_ref().and_then(|m| m.content.as_ref()) {
Some(Value::Array(arr)) => arr,
_ => &[],
}
}
fn value_to_plain_string(v: &Value) -> String {
match v {
Value::String(s) => s.clone(),
Value::Null => "null".into(),
Value::Bool(b) => b.to_string(),
Value::Number(n) => n.to_string(),
_ => serde_json::to_string(v).unwrap_or_default(),
}
}
fn stringify_result(content: Option<&Value>) -> String {
match content {
Some(Value::String(s)) => s.clone(),
Some(Value::Array(arr)) => arr
.iter()
.map(|c| match c {
Value::String(s) => s.clone(),
Value::Object(_) => match c.get("text") {
Some(t) => value_to_plain_string(t),
None => serde_json::to_string(c).unwrap_or_default(),
},
_ => serde_json::to_string(c).unwrap_or_default(),
})
.collect::<Vec<_>>()
.join("\n"),
Some(other) => serde_json::to_string(other).unwrap_or_default(),
None => String::new(),
}
}
pub(crate) fn extract_invocations(records: &[TranscriptRecord]) -> Vec<ToolInvocation> {
let mut invocations: Vec<ToolInvocation> = Vec::new();
let mut index_by_id: HashMap<String, usize> = HashMap::new();
for record in records {
let rtype = record.record_type.as_deref();
let blocks = content_blocks(&record.message);
if rtype == Some("assistant") {
for block in blocks {
if block.get("type").and_then(Value::as_str) != Some("tool_use") {
continue;
}
let ordinal = invocations.len();
if let Some(id) = block.get("id").and_then(Value::as_str) {
index_by_id.insert(id.to_string(), ordinal);
}
invocations.push(ToolInvocation {
name: block
.get("name")
.and_then(Value::as_str)
.unwrap_or("")
.to_string(),
args: block.get("input").cloned(),
result: None,
ordinal: ordinal as u32,
});
}
} else if rtype == Some("user") {
for block in blocks {
if block.get("type").and_then(Value::as_str) != Some("tool_result") {
continue;
}
let Some(id) = block.get("tool_use_id").and_then(Value::as_str) else {
continue;
};
let Some(&idx) = index_by_id.get(id) else {
continue;
};
invocations[idx].result =
Some(Value::String(stringify_result(block.get("content"))));
}
}
}
invocations
}
pub(crate) fn last_assistant_text(records: &[TranscriptRecord]) -> Option<String> {
let mut final_text: Option<String> = None;
for record in records {
if record.record_type.as_deref() != Some("assistant") {
continue;
}
let texts: Vec<&str> = content_blocks(&record.message)
.iter()
.filter(|b| b.get("type").and_then(Value::as_str) == Some("text"))
.filter_map(|b| b.get("text").and_then(Value::as_str))
.collect();
if !texts.is_empty() {
final_text = Some(texts.join("\n"));
}
}
final_text
}
#[cfg(test)]
mod tests {
use super::*;
use crate::adapters::transcript::read_jsonl;
use serde_json::{Value, json};
use std::fs;
use std::path::Path;
use tempfile::TempDir;
fn write_jsonl(path: &Path, lines: &[Value]) {
let body = lines
.iter()
.map(|l| l.to_string())
.collect::<Vec<_>>()
.join("\n");
fs::write(path, format!("{body}\n")).unwrap();
}
fn invocations(path: &Path) -> Vec<ToolInvocation> {
extract_invocations(&read_jsonl::<TranscriptRecord>(path).unwrap())
}
#[test]
fn extracts_tool_use_blocks_with_ordinal_and_args() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("simple.jsonl");
write_jsonl(
&path,
&[
json!({"type": "user", "message": {"role": "user", "content": "Run the tests"}}),
json!({"type": "assistant", "message": {"role": "assistant", "content": [
{"type": "text", "text": "Running tests now."},
{"type": "tool_use", "id": "toolu_001", "name": "Bash", "input": {"command": "bun test"}}
]}}),
json!({"type": "user", "message": {"role": "user", "content": [
{"type": "tool_result", "tool_use_id": "toolu_001", "content": "2 pass\n0 fail"}
]}}),
json!({"type": "assistant", "message": {"role": "assistant", "content": [
{"type": "tool_use", "id": "toolu_002", "name": "Read", "input": {"file_path": "/tmp/x.txt"}}
]}}),
],
);
let result = invocations(&path);
assert_eq!(result.len(), 2);
assert_eq!(result[0].name, "Bash");
assert_eq!(result[0].ordinal, 0);
assert_eq!(result[0].args, Some(json!({"command": "bun test"})));
assert_eq!(
result[0].result,
Some(Value::String("2 pass\n0 fail".into()))
);
assert_eq!(result[1].name, "Read");
assert_eq!(result[1].ordinal, 1);
assert_eq!(result[1].args, Some(json!({"file_path": "/tmp/x.txt"})));
assert_eq!(result[1].result, None);
}
#[test]
fn returns_empty_when_no_tool_use_blocks() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("no-tools.jsonl");
write_jsonl(
&path,
&[
json!({"type": "user", "message": {"role": "user", "content": "hi"}}),
json!({"type": "assistant", "message": {"role": "assistant", "content": [{"type": "text", "text": "hello"}]}}),
],
);
assert_eq!(invocations(&path), vec![]);
}
#[test]
fn skips_malformed_jsonl_lines() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("malformed.jsonl");
let good_a = json!({"type": "assistant", "message": {"role": "assistant", "content": [
{"type": "tool_use", "id": "toolu_a", "name": "Bash", "input": {"command": "ls"}}
]}});
let good_b = json!({"type": "assistant", "message": {"role": "assistant", "content": [
{"type": "tool_use", "id": "toolu_b", "name": "Read", "input": {"file_path": "/tmp"}}
]}});
let body = format!("{good_a}\nnot valid json\n{good_b}\n");
fs::write(&path, body).unwrap();
let result = invocations(&path);
assert_eq!(result.len(), 2);
assert_eq!(
result.iter().map(|r| r.name.as_str()).collect::<Vec<_>>(),
vec!["Bash", "Read"]
);
}
#[test]
fn handles_tool_result_with_array_content() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("array-result.jsonl");
write_jsonl(
&path,
&[
json!({"type": "assistant", "message": {"role": "assistant", "content": [
{"type": "tool_use", "id": "toolu_x", "name": "Bash", "input": {"command": "echo hi"}}
]}}),
json!({"type": "user", "message": {"role": "user", "content": [
{"type": "tool_result", "tool_use_id": "toolu_x", "content": [{"type": "text", "text": "hi"}]}
]}}),
],
);
let result = invocations(&path);
assert_eq!(result.len(), 1);
assert_eq!(result[0].result, Some(Value::String("hi".into())));
}
#[test]
fn last_assistant_text_concatenates_text_of_last_assistant_message() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("final-text.jsonl");
write_jsonl(
&path,
&[
json!({"type": "assistant", "message": {"id": "msg_1", "role": "assistant", "content": [{"type": "text", "text": "intermediate"}]}}),
json!({"type": "assistant", "message": {"id": "msg_2", "role": "assistant", "content": [
{"type": "text", "text": "All tests pass."},
{"type": "tool_use", "id": "toolu_z", "name": "Bash", "input": {"command": "true"}},
{"type": "text", "text": "Wrapping up."}
]}}),
json!({"type": "user", "message": {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_z", "content": "ok"}]}}),
],
);
assert_eq!(
last_assistant_text(&read_jsonl::<TranscriptRecord>(&path).unwrap()),
Some("All tests pass.\nWrapping up.".into())
);
}
#[test]
fn last_assistant_text_is_null_when_no_assistant_text() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("no-text.jsonl");
write_jsonl(
&path,
&[json!({"type": "user", "message": {"role": "user", "content": "hi"}})],
);
assert_eq!(
last_assistant_text(&read_jsonl::<TranscriptRecord>(&path).unwrap()),
None
);
}
}