poimen-memory/crates/mem-ingest/src/claude_transcript.rs

use anyhow::anyhow;
use futures::stream::{Stream, StreamExt};
use mem_chunk::RecordSource;
use mem_core::{Provenance, Record, Role};
use serde_json::Value;
use std::path::{Path, PathBuf};
use tokio::fs::File;
use tokio::io::{AsyncBufReadExt, BufReader};
use crate::pi_session::flatten_content;

/// A source that reads Claude transcript JSONL files.
pub struct ClaudeTranscriptSource {
    file_path: PathBuf,
}

impl ClaudeTranscriptSource {
    /// Create a new Claude transcript source from a file path.
    pub fn new(file_path: PathBuf) -> Self {
        ClaudeTranscriptSource { file_path }
    }

    /// Extract the project key from the transcript file's cwd field.
    pub async fn read_project_key(&self) -> anyhow::Result<String> {
        let file = File::open(&self.file_path).await?;
        let reader = BufReader::new(file);
        let mut lines = reader.lines();

        while let Some(line_result) = lines.next_line().await? {
            let value: Value = match serde_json::from_str(&line_result) {
                Ok(v) => v,
                Err(_) => continue,
            };

            if let Some(cwd) = value.get("cwd").and_then(|v| v.as_str()) {
                return Ok(cwd.to_string());
            }
        }

        Err(anyhow!("No line with 'cwd' field found in transcript"))
    }
}

impl RecordSource for ClaudeTranscriptSource {
    fn records(self) -> Box<dyn Stream<Item = Result<Record, String>> + Unpin> {
        let file_path = self.file_path.clone();
        let session_id = file_path
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or("unknown")
            .to_string();

        // Create a stream that reads the file asynchronously
        let stream = futures::stream::once(async move {
            match read_claude_transcript_lines(&file_path, &session_id).await {
                Ok(records) => records,
                Err(e) => vec![Err(e.to_string())],
            }
        })
        .flat_map(futures::stream::iter);

        Box::new(stream.boxed())
    }
}

/// Read lines from a claude transcript file and convert them to Records.
pub async fn read_claude_transcript_lines(
    path: &Path,
    session_id: &str,
) -> anyhow::Result<Vec<Result<Record, String>>> {
    let file = File::open(path).await?;
    let reader = BufReader::new(file);
    let mut lines = reader.lines();
    let mut records = Vec::new();
    let mut line_num = 0u64;

    while let Some(line_result) = lines.next_line().await? {
        line_num += 1;

        match serde_json::from_str::<Value>(&line_result) {
            Ok(value) => {
                let record_type = value
                    .get("type")
                    .and_then(|v| v.as_str())
                    .unwrap_or("unknown");

                match record_type {
                    "user" => {
                        if let Ok(Some(record)) = parse_claude_message(&value, Role::User, session_id, line_num) {
                            records.push(Ok(record));
                        }
                    }
                    "assistant" => {
                        if let Ok(Some(record)) = parse_claude_message(&value, Role::Assistant, session_id, line_num) {
                            records.push(Ok(record));
                        }
                    }
                    "system" => {
                        // Only emit system messages with api_error subtype
                        if value.get("subtype").and_then(|v| v.as_str()) == Some("api_error") {
                            if let Ok(Some(record)) = parse_claude_message(&value, Role::System, session_id, line_num) {
                                records.push(Ok(record));
                            }
                        }
                    }
                    "attachment" | "queue-operation" | "file-history-snapshot" | "summary" | "mode" | "ai-title" | "last-prompt" => {
                        // Skip these record types
                    }
                    _ => {}
                }
            }
            Err(_) => {
                // Malformed line - skip with warning
                tracing::warn!(
                    "Skipping malformed JSON at line {}: {}",
                    line_num,
                    line_result.trim()
                );
            }
        }
    }

    Ok(records)
}

fn parse_claude_message(
    value: &Value,
    role: Role,
    session_id: &str,
    line_num: u64,
) -> Result<Option<Record>, String> {
    let message = value.get("message");
    let timestamp = parse_claude_timestamp(value);

    let text = if let Some(msg) = message {
        if let Some(content) = msg.get("content") {
            flatten_content(content)?
        } else {
            "".to_string()
        }
    } else {
        "".to_string()
    };

    if text.is_empty() && role != Role::System {
        return Ok(None);
    }

    Ok(Some(Record {
        role,
        text,
        timestamp,
        provenance: Provenance {
            source_id: format!("claude:{}", session_id),
            offset: line_num,
        },
    }))
}

fn parse_claude_timestamp(value: &Value) -> time::OffsetDateTime {
    value
        .get("timestamp")
        .and_then(|v| v.as_str())
        .and_then(|s| {
            time::OffsetDateTime::parse(s, &time::format_description::well_known::Rfc3339).ok()
        })
        .unwrap_or_else(time::OffsetDateTime::now_utc)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_claude_transcript_source_creation() {
        let source = ClaudeTranscriptSource::new(PathBuf::from("test.jsonl"));
        assert_eq!(source.file_path.to_string_lossy(), "test.jsonl");
    }
}