Skip to main content

agent_sessions/
import.rs

1//! Snapshot imports for native JSON, patch-based JSONL, and read-only databases.
2//! JSON pointers / database row keys are retained instead of fabricated byte offsets.
3mod cursor_cli;
4mod database;
5mod grok;
6mod replay;
7mod warp;
8use crate::parser::{Parsed, State, tally};
9use crate::*;
10use chrono::{DateTime, Utc};
11use serde::Serialize;
12use serde_json::Value;
13use std::{
14    fs::File,
15    io::{BufRead, BufReader},
16    path::Path,
17};
18
19// ProtoJSON cannot carry unknown wire fields. Preserve a diagnostic while
20// returning known events, so a new field cannot look like a fully supported import.
21pub(super) fn note_protobuf(message: &prost_reflect::DynamicMessage, summary: &mut ReadSummary) {
22    fn visit(value: &prost_reflect::Value, summary: &mut ReadSummary) {
23        match value {
24            prost_reflect::Value::Message(m) => note_protobuf(m, summary),
25            prost_reflect::Value::List(v) => {
26                for item in v {
27                    visit(item, summary);
28                }
29            }
30            prost_reflect::Value::Map(v) => {
31                for item in v.values() {
32                    visit(item, summary);
33                }
34            }
35            _ => {}
36        }
37    }
38    for field in message.unknown_fields() {
39        tally(
40            &mut summary.unknown_types,
41            &format!(
42                "protobuf:{}:{}",
43                prost_reflect::ReflectMessage::descriptor(message).full_name(),
44                field.number()
45            ),
46        );
47    }
48    for (_, value) in message.fields() {
49        visit(value, summary);
50    }
51}
52
53#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
54pub enum ImportSource {
55    JsonPointer(String),
56    Record(Location),
57    DatabaseRow { table: String, key: String },
58}
59#[derive(Debug, Clone, PartialEq, Serialize)]
60pub struct ImportedEvent {
61    /// All native records contributing to this event, including replay patches.
62    pub sources: Vec<ImportSource>,
63    pub at: Option<DateTime<Utc>>,
64    pub timestamp_text: Option<String>,
65    pub session_id: Option<String>,
66    pub message_id: Option<String>,
67    pub record_id: Option<String>,
68    pub value: Event,
69}
70#[derive(Debug, Default, Serialize)]
71pub struct SessionImport {
72    pub events: Vec<ImportedEvent>,
73    pub summary: ReadSummary,
74}
75#[derive(Debug)]
76pub enum ImportError {
77    Read(ReadError),
78    Stream(StreamError),
79}
80impl std::fmt::Display for ImportError {
81    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82        match self {
83            Self::Read(e) => e.fmt(f),
84            Self::Stream(e) => e.fmt(f),
85        }
86    }
87}
88impl std::error::Error for ImportError {
89    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
90        Some(match self {
91            Self::Read(e) => e,
92            Self::Stream(e) => e,
93        })
94    }
95}
96impl From<ReadError> for ImportError {
97    fn from(e: ReadError) -> Self {
98        Self::Read(e)
99    }
100}
101impl From<StreamError> for ImportError {
102    fn from(e: StreamError) -> Self {
103        Self::Stream(e)
104    }
105}
106impl From<LineErrorKind> for ImportError {
107    fn from(e: LineErrorKind) -> Self {
108        Self::Read(ReadError::Format(e))
109    }
110}
111
112/// Import an immutable file prefix. Database sources use import_database instead.
113/// The result must pass summary.is_supported() before claiming complete support.
114pub fn import_session(
115    agent: Agent,
116    path: &Path,
117    opts: &ReadOptions,
118) -> Result<SessionImport, ImportError> {
119    let file = File::open(path).map_err(ReadError::Io)?;
120    let mut opts = opts.clone();
121    if opts.stop_at_byte.is_none() {
122        opts.stop_at_byte = Some(file.metadata().map_err(ReadError::Io)?.len());
123    }
124    if let (Some(end), Some(limit)) = (opts.stop_at_byte, opts.max_file_bytes)
125        && end > limit
126    {
127        return Err(ReadError::TooLarge { limit }.into());
128    }
129    import_session_from(agent, BufReader::new(file), &opts)
130}
131/// Reads one native session/export. Patch-based sources are replayed before events
132/// are returned, so old text and rewound messages cannot masquerade as final text.
133pub fn import_session_from<R: BufRead>(
134    agent: Agent,
135    source: R,
136    opts: &ReadOptions,
137) -> Result<SessionImport, ImportError> {
138    opts.validate()?;
139    if agent.streaming() && agent != Agent::KimiCli {
140        let mut reader = read_from(agent, source, opts)?;
141        let mut events = Vec::new();
142        for e in reader.by_ref() {
143            let e = e?;
144            events.push(ImportedEvent {
145                sources: vec![ImportSource::Record(e.location)],
146                at: e.at,
147                timestamp_text: e.timestamp_text,
148                session_id: e.session_id,
149                message_id: e.message_id,
150                record_id: e.record_id,
151                value: e.value,
152            });
153        }
154        return Ok(SessionImport {
155            events,
156            summary: reader.finish(),
157        });
158    }
159    if matches!(
160        agent,
161        Agent::Cursor | Agent::ZCode | Agent::CursorCli | Agent::Warp
162    ) {
163        return Err(ReadError::InvalidOptions(
164            "this agent requires import_database with a native session ID",
165        )
166        .into());
167    }
168    // Reuse the existing bounded physical framer; no second IO/size contract.
169    let mut raw = read_raw_from(
170        source,
171        &RawReadOptions {
172            start_offset: 0,
173            stop_at_byte: opts.stop_at_byte,
174            max_read_bytes: opts.max_file_bytes,
175            max_line_bytes: opts.max_line_bytes,
176        },
177    )?;
178    let mut bytes = Vec::new();
179    let mut records = Vec::new();
180    for record in raw.by_ref() {
181        let record = record?;
182        bytes.extend_from_slice(&record.bytes);
183        records.push(record);
184    }
185    let raw_summary = raw.finish();
186    let summary = ReadSummary {
187        status: raw_summary.status,
188        lines: raw_summary.records,
189        bytes_read: raw_summary.bytes_read,
190        last_complete_byte: raw_summary.delivered_through,
191        ..Default::default()
192    };
193    let mut result = SessionImport {
194        summary,
195        ..Default::default()
196    };
197    if agent == Agent::GeminiCli {
198        // A legacy JSON document parses as one value; otherwise use native JSONL replay.
199        if let Ok(v) = serde_json::from_slice::<Value>(&bytes)
200            && v.get("messages").is_some()
201        {
202            document(agent, &v, opts, &mut result)?;
203        } else {
204            replay::gemini(&records, opts, &mut result)?;
205        }
206    } else if agent == Agent::KimiCli {
207        replay::kimi(&records, opts, &mut result)?;
208    } else if agent == Agent::Grok {
209        grok::import(&records, opts, &mut result)?;
210    } else {
211        let v = serde_json::from_slice::<Value>(&bytes).map_err(|_| LineErrorKind::InvalidJson)?;
212        document(agent, &v, opts, &mut result)?;
213    }
214    if result.summary.truncated_tail {
215        result.summary.status = ReadStatus::IncompleteTail;
216    }
217    Ok(result)
218}
219fn array<'a>(v: &'a Value, key: &'static str) -> Result<&'a Vec<Value>, LineErrorKind> {
220    v.get(key)
221        .and_then(Value::as_array)
222        .ok_or(LineErrorKind::MissingField(key))
223}
224pub(super) fn document(
225    agent: Agent,
226    v: &Value,
227    opts: &ReadOptions,
228    result: &mut SessionImport,
229) -> Result<(), ImportError> {
230    let mut state = State {
231        codex_content: opts.codex_content,
232        ..Default::default()
233    };
234    let (rows, prefix) = match agent {
235        Agent::GeminiCli => {
236            state.session_id = crate::parser::string(v, "sessionId");
237            if v.get("kind").and_then(Value::as_str) == Some("subagent") {
238                state.sidechain = true;
239            }
240            (array(v, "messages")?, "/messages")
241        }
242        Agent::Cline | Agent::RooCode => (
243            v.as_array().ok_or_else(|| {
244                LineErrorKind::InvalidField("API history must be an array".into())
245            })?,
246            "",
247        ),
248        Agent::ClineCli => {
249            state.session_id = crate::parser::string(v, "sessionId");
250            (array(v, "messages")?, "/messages")
251        }
252        Agent::Hermes => {
253            state.session_id = crate::parser::string(v, "session_id");
254            emit_meta(
255                MetaUpdate {
256                    session_id: state.session_id.clone(),
257                    model: crate::parser::string(v, "model"),
258                    ..Default::default()
259                },
260                vec![ImportSource::JsonPointer("".into())],
261                opts,
262                result,
263            );
264            (array(v, "messages")?, "/messages")
265        }
266        Agent::GrokBot => {
267            let value = v.get("value").ok_or(LineErrorKind::MissingField("value"))?;
268            (array(value, "entries")?, "/value/entries")
269        }
270        Agent::Zed => (array(v, "messages")?, "/messages"),
271        Agent::Antigravity => {
272            let trajectory = v.get("trajectory").unwrap_or(v);
273            state.session_id = crate::parser::string(trajectory, "cascadeId")
274                .or_else(|| crate::parser::string(trajectory, "trajectoryId"));
275            let rows = array(trajectory, "steps")?;
276            if let Some(total) = v.get("numTotalSteps") {
277                let total = total
278                    .as_u64()
279                    .or_else(|| total.as_str().and_then(|v| v.parse().ok()))
280                    .ok_or_else(|| LineErrorKind::InvalidField("numTotalSteps".into()))?;
281                if total != rows.len() as u64 {
282                    return Err(LineErrorKind::InvalidField(
283                        "incomplete Antigravity trajectory; fetch all steps".into(),
284                    )
285                    .into());
286                }
287            }
288            (
289                rows,
290                if v.get("trajectory").is_some() {
291                    "/trajectory/steps"
292                } else {
293                    "/steps"
294                },
295            )
296        }
297        Agent::OpenCode => {
298            let info = v.get("info").ok_or(LineErrorKind::MissingField("info"))?;
299            state.session_id = crate::parser::string(info, "id");
300            state.sidechain = info.get("parentID").is_some_and(|v| !v.is_null());
301            emit_meta(
302                MetaUpdate {
303                    session_id: state.session_id.clone(),
304                    cwd: crate::parser::string(info, "directory"),
305                    ..Default::default()
306                },
307                vec![ImportSource::JsonPointer("/info".into())],
308                opts,
309                result,
310            );
311            (array(v, "messages")?, "/messages")
312        }
313        Agent::Goose => {
314            state.session_id = crate::parser::string(v, "id");
315            emit_meta(
316                MetaUpdate {
317                    session_id: state.session_id.clone(),
318                    cwd: crate::parser::string(v, "working_dir"),
319                    ..Default::default()
320                },
321                vec![ImportSource::JsonPointer("".into())],
322                opts,
323                result,
324            );
325            let c = v
326                .get("conversation")
327                .ok_or(LineErrorKind::MissingField("conversation"))?;
328            if let Some(rows) = c.as_array() {
329                (rows, "/conversation")
330            } else {
331                (array(c, "messages")?, "/conversation/messages")
332            }
333        }
334        Agent::Continue => {
335            state.session_id = crate::parser::string(v, "sessionId");
336            emit_meta(
337                MetaUpdate {
338                    session_id: state.session_id.clone(),
339                    cwd: crate::parser::string(v, "workspaceDirectory"),
340                    model: crate::parser::string(v, "chatModelTitle"),
341                    ..Default::default()
342                },
343                vec![ImportSource::JsonPointer("".into())],
344                opts,
345                result,
346            );
347            (array(v, "history")?, "/history")
348        }
349        _ => {
350            return Err(
351                ReadError::InvalidOptions("this agent does not use a JSON document").into(),
352            );
353        }
354    };
355    let mut tool_ids = std::collections::HashSet::new();
356    let mut usage_ids = std::collections::HashSet::new();
357    for (i, row) in rows.iter().enumerate() {
358        let start = result.events.len();
359        parse_record(
360            agent,
361            row,
362            &mut state,
363            vec![ImportSource::JsonPointer(format!("{prefix}/{i}"))],
364            opts,
365            result,
366        )?;
367        if agent == Agent::Antigravity {
368            let mut tail = result.events.split_off(start);
369            tail.retain(|e| match &e.value {
370                Event::ToolCall(c) => c.id.as_ref().is_none_or(|id| tool_ids.insert(id.clone())),
371                Event::Usage(u) => u
372                    .dedup_key
373                    .as_ref()
374                    .is_none_or(|id| usage_ids.insert(id.clone())),
375                _ => true,
376            });
377            result.events.extend(tail);
378        }
379    }
380    if agent == Agent::Zed
381        && opts.include.contains(EventKinds::USAGE)
382        && let Some(usage) = v.get("request_token_usage").and_then(Value::as_object)
383    {
384        for (id, counts) in usage {
385            let mut p = Parsed {
386                include: opts.include,
387                ..Default::default()
388            };
389            p.emit(
390                0,
391                Event::Usage(Usage {
392                    counts: TokenCounts {
393                        input: crate::adapters::counter(counts, "input_tokens")?,
394                        output: crate::adapters::counter(counts, "output_tokens")?,
395                        cache_read: crate::adapters::counter(counts, "cache_read_input_tokens")?,
396                        cache_write: crate::adapters::counter(
397                            counts,
398                            "cache_creation_input_tokens",
399                        )?,
400                        ..Default::default()
401                    },
402                    semantics: TokenSemantics::Unknown,
403                    model: None,
404                    dedup_key: Some(id.clone()),
405                    cumulative: None,
406                    basis: UsageBasis::Response,
407                    stop_reason: None,
408                    endpoint: Endpoint::Unknown,
409                    inference_geo: None,
410                    adjustments: Vec::new(),
411                }),
412            );
413            append(
414                p,
415                &state,
416                vec![ImportSource::JsonPointer(format!(
417                    "/request_token_usage/{}",
418                    id.replace('~', "~0").replace('/', "~1")
419                ))],
420                result,
421            );
422        }
423    }
424    Ok(())
425}
426pub(crate) fn parse_record(
427    agent: Agent,
428    v: &Value,
429    state: &mut State,
430    sources: Vec<ImportSource>,
431    opts: &ReadOptions,
432    result: &mut SessionImport,
433) -> Result<(), ImportError> {
434    let mut p = Parsed {
435        include: opts.include,
436        accounting: opts.accounting,
437        ..Default::default()
438    };
439    p.at = if matches!(agent, Agent::GeminiCli | Agent::QwenCode) {
440        crate::parser::timestamp(v)?
441    } else {
442        None
443    };
444    p.timestamp_text = v
445        .get("timestamp")
446        .and_then(Value::as_str)
447        .map(str::to_owned);
448    if matches!(agent, Agent::Cline | Agent::RooCode | Agent::ClineCli) {
449        p.at = crate::adapters::millis(v.get("ts"))?;
450    }
451    if agent == Agent::Hermes {
452        p.at = crate::adapters::seconds(v.get("timestamp"))?;
453    }
454    crate::adapters::parse(agent, v, state, &mut p)?;
455    append(p, state, sources, result);
456    Ok(())
457}
458pub(crate) fn append(
459    p: Parsed,
460    state: &State,
461    sources: Vec<ImportSource>,
462    result: &mut SessionImport,
463) {
464    for unknown in p.unknown {
465        tally(&mut result.summary.unknown_types, &unknown);
466    }
467    for ignored in p.ignored {
468        tally(&mut result.summary.ignored_types, &ignored);
469    }
470    for (_, mut value) in p.events {
471        if let Event::Message(m) = &mut value {
472            m.is_sidechain |= state.sidechain;
473        }
474        result.events.push(ImportedEvent {
475            sources: sources.clone(),
476            at: p.at,
477            timestamp_text: p.timestamp_text.clone(),
478            session_id: state.session_id.clone(),
479            message_id: p.message_id.clone(),
480            record_id: p.record_id.clone(),
481            value,
482        });
483    }
484}
485pub(crate) fn emit_meta(
486    meta: MetaUpdate,
487    sources: Vec<ImportSource>,
488    opts: &ReadOptions,
489    result: &mut SessionImport,
490) {
491    if opts.include.contains(EventKinds::META) {
492        result.events.push(ImportedEvent {
493            session_id: meta.session_id.clone(),
494            sources,
495            at: None,
496            timestamp_text: None,
497            message_id: None,
498            record_id: None,
499            value: Event::Meta(meta),
500        });
501    }
502}
503pub(crate) fn decode_record(
504    record: &RawRecord,
505    opts: &ReadOptions,
506    result: &mut SessionImport,
507) -> Result<Option<Value>, ImportError> {
508    if record.bytes.iter().all(u8::is_ascii_whitespace) {
509        return Ok(None);
510    }
511    if std::str::from_utf8(&record.bytes).is_err() {
512        return Err(StreamError::Line {
513            line_no: record.line_no,
514            byte_start: record.byte_start,
515            kind: LineErrorKind::InvalidUtf8,
516        }
517        .into());
518    }
519    match serde_json::from_slice(&record.bytes) {
520        Ok(v) => Ok(Some(v)),
521        Err(e) if !record.terminated && e.is_eof() && opts.tail == TailMode::AllowIncomplete => {
522            result.summary.truncated_tail = true;
523            result.summary.last_complete_byte = record.byte_start;
524            Ok(None)
525        }
526        Err(_) => Err(StreamError::Line {
527            line_no: record.line_no,
528            byte_start: record.byte_start,
529            kind: LineErrorKind::InvalidJson,
530        }
531        .into()),
532    }
533}
534pub(crate) fn record_source(r: &RawRecord) -> ImportSource {
535    ImportSource::Record(Location {
536        record_index: r.line_no - 1,
537        line_no: r.line_no,
538        byte_start: r.byte_start,
539        byte_end: r.byte_end,
540        event_index: 0,
541    })
542}
543pub use database::{DatabaseSession, import_database, list_database_sessions};