use crate::engine::context::ir::{build_tool_call, extract_text};
use crate::engine::context::{ContextError, ContextReader, for_each_jsonl_record};
use crate::engine::message::{
BashOutput, Context, ContextListing, ConversationMessage, Entry, Image, MessageKind, Part,
ProviderId, Reasoning, ToolResultData,
};
use anyhow::Context as _;
use chrono::{DateTime, NaiveDateTime, Utc};
use serde_json::Value;
use std::fs;
use std::io::{BufRead, BufReader};
use std::path::{Path, PathBuf};
const BRANCH_SUMMARY_PREFIX: &str =
"The following is a summary of a branch that this conversation came back from:\n\n<summary>\n";
const BRANCH_SUMMARY_SUFFIX: &str = "</summary>";
fn pi_start_from_filename(file_path: &Path) -> Option<DateTime<Utc>> {
let stem = file_path.file_stem()?.to_str()?;
let ts_part = stem.split('_').next()?;
NaiveDateTime::parse_from_str(ts_part, "%Y-%m-%dT%H-%M-%S-%3fZ")
.ok()
.map(|ndt| DateTime::<Utc>::from_naive_utc_and_offset(ndt, Utc))
}
pub struct JsonlReader {
file_path: PathBuf,
}
impl JsonlReader {
pub fn new(file_path: PathBuf) -> Self {
Self { file_path }
}
}
impl ContextReader for JsonlReader {
fn list_contexts(&self) -> anyhow::Result<Vec<ContextListing>> {
let file_path = &self.file_path;
let file =
fs::File::open(file_path).with_context(|| format!("open {}", file_path.display()))?;
let mut reader = BufReader::new(file);
let mut first_line = String::new();
reader
.read_line(&mut first_line)
.with_context(|| format!("read {}", file_path.display()))?;
let first_line = first_line.trim_end().to_string();
if first_line.is_empty() {
return Err(ContextError::EmptyFile.into());
}
let session: Value = serde_json::from_str(&first_line)?;
if session["type"].as_str() != Some("session") {
return Err(ContextError::InvalidSessionHeader("session").into());
}
if session["version"].as_u64() != Some(3) {
return Err(ContextError::UnsupportedPiSessionVersion.into());
}
let id = session["id"]
.as_str()
.ok_or(ContextError::MissingSessionId)?
.to_string();
let cwd = session["cwd"].as_str().unwrap_or("").to_string();
Ok(vec![ContextListing {
id,
provider_id: ProviderId {
label: None,
cwd: PathBuf::from(&cwd),
count: 0,
from: pi_start_from_filename(file_path),
until: None,
},
path: file_path.clone(),
parent_id: None,
}])
}
fn backing_file(&self) -> Option<&Path> {
Some(self.file_path.as_path())
}
fn read_context(&self, _context_id: &str) -> anyhow::Result<Context> {
self.list_contexts()?;
let mut entries = Vec::new();
let mut messages = Vec::new();
let mut cwd = String::new();
for_each_jsonl_record(&self.file_path, |_line_num, entry_json| {
let entry_type = entry_json["type"].as_str().unwrap_or("").to_string();
if entry_type == "session" {
if cwd.is_empty()
&& let Some(value) = entry_json["cwd"].as_str()
{
cwd = value.to_string();
}
return;
}
let id = entry_json["id"].as_str().unwrap_or("").to_string();
let parent_id = entry_json["parentId"].as_str().unwrap_or("").to_string();
let native_data = if entry_type == "compaction" {
let summary = entry_json["summary"].as_str().unwrap_or("").to_string();
let mut message =
ConversationMessage::new(id.clone(), "user", vec![Part::Text(summary.clone())]);
message.kind = MessageKind::PiCompaction {
summary,
first_kept_entry_id: entry_json["firstKeptEntryId"]
.as_str()
.unwrap_or("")
.to_string(),
tokens_before: entry_json["tokensBefore"].as_u64().unwrap_or(0),
details: entry_json
.get("details")
.filter(|value| !value.is_null())
.cloned(),
};
messages.push(message);
None
} else if entry_type == "custom_message" {
let content = entry_json.get("content").cloned().unwrap_or(Value::Null);
let mut message = ConversationMessage::new(
id.clone(),
"user",
jsonl_text_and_image_parts(Some(&content)),
);
message.kind = MessageKind::PiCustomMessage {
custom_type: entry_json["customType"].as_str().unwrap_or("").to_string(),
content,
display: entry_json["display"].as_bool().unwrap_or(false),
details: entry_json
.get("details")
.filter(|value| !value.is_null())
.cloned(),
};
messages.push(message);
None
} else if entry_type == "branch_summary" {
let summary = entry_json["summary"].as_str().unwrap_or("").to_string();
let text = format!("{BRANCH_SUMMARY_PREFIX}{summary}{BRANCH_SUMMARY_SUFFIX}");
let mut message =
ConversationMessage::new(id.clone(), "user", vec![Part::Text(text)]);
message.kind = MessageKind::PiBranchSummary {
summary,
from_id: entry_json["fromId"].as_str().unwrap_or("").to_string(),
details: entry_json
.get("details")
.filter(|value| !value.is_null())
.cloned(),
from_hook: entry_json.get("fromHook").and_then(Value::as_bool),
};
messages.push(message);
None
} else if let Some(raw_msg) = entry_json.get("message") {
let (role, parts) = jsonl_build_message(&entry_type, raw_msg);
messages.push(ConversationMessage::new(id.clone(), role, parts));
None
} else {
Some(entry_json.clone())
};
entries.push(Entry {
id,
parent_id,
native_data,
});
})?;
Ok(Context {
entries,
messages,
cwd: (!cwd.is_empty()).then_some(cwd),
})
}
}
fn jsonl_build_message(msg_type: &str, msg: &Value) -> (String, Vec<Part>) {
let role = msg["role"].as_str().unwrap_or(msg_type).to_string();
let content = msg.get("content");
match role.as_str() {
"assistant" => {
let mut parts = Vec::new();
if let Some(Value::Array(blocks)) = content {
for block in blocks {
match block["type"].as_str() {
Some("text") => {
parts
.push(Part::Text(block["text"].as_str().unwrap_or("").to_string()));
}
Some("thinking") => {
if let Some(text) = block["thinking"].as_str() {
parts.push(Part::Reasoning(Reasoning::new(text)));
}
}
Some("image") => {
parts.push(Part::Image(jsonl_image(block)));
}
Some("toolCall") => {
let name = block["name"].as_str().unwrap_or("");
let id = block["id"].as_str().unwrap_or("");
parts.push(Part::ToolCall(build_tool_call(
id,
name,
block.get("arguments"),
)));
}
_ => {}
}
}
}
("assistant".to_string(), parts)
}
"toolResult" => {
let tool_name = msg["toolName"].as_str().unwrap_or("").to_string();
let is_error = msg["isError"].as_bool().unwrap_or(false);
let call_id = msg["toolCallId"].as_str().unwrap_or("").to_string();
(
"user".to_string(),
vec![Part::ToolResult(ToolResultData {
call_id,
tool_name,
content: extract_text(content),
is_error,
})],
)
}
"bashExecution" => (
"user".to_string(),
vec![Part::Bash(BashOutput {
command: msg["command"].as_str().unwrap_or("").to_string(),
output: msg["output"].as_str().unwrap_or("").to_string(),
})],
),
_ => (role, jsonl_text_and_image_parts(content)),
}
}
fn jsonl_text_and_image_parts(content: Option<&Value>) -> Vec<Part> {
let mut parts = Vec::new();
if let Some(Value::Array(blocks)) = content {
for block in blocks {
if block["type"].as_str() == Some("image") {
parts.push(Part::Image(jsonl_image(block)));
}
}
}
let text = extract_text(content);
if !text.is_empty() || parts.is_empty() {
parts.insert(0, Part::Text(text));
}
parts
}
fn jsonl_image(block: &Value) -> Image {
Image {
mime_type: block["mimeType"].as_str().unwrap_or("").to_string(),
data: block["data"].as_str().unwrap_or("").to_string(),
}
}