johnandersen777.bsky.social/asyncai

Switch repository

Loading…

pub mod harnesses;

use std::io;
use std::path::{Path, PathBuf};
use std::time::Duration;

use serde_json::{json, Value};
use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
use tokio::task::JoinHandle;

use crate::tools::ToolObserver;

pub use harnesses::claude_code_cli::ClaudeCodeCliHarness;

#[derive(Clone, Debug, PartialEq)]
pub struct MonitorEvent {
    pub task_id: String,
    pub description: String,
    pub event: String,
    pub source_uuid: String,
    pub delivery_id: String,
    pub human_turn: bool,
    /// Instruction the model saw with this event; replaces the default
    /// PushNotification sentence in Claude transcripts.
    pub note: Option<String>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum TranscriptEvent {
    User {
        text: String,
    },
    Assistant {
        text: String,
    },
    ToolCall {
        id: String,
        name: String,
        arguments: Value,
    },
    ToolResult {
        id: String,
        content: String,
        is_error: bool,
        /// Harness metadata (code mode nested calls); not model input.
        audit: Option<Value>,
    },
    MonitorEvent(MonitorEvent),
    MonitorConsumed {
        event: MonitorEvent,
        reason: String,
    },
}

pub trait SessionHarness: Send {
    fn format_name(&self) -> &'static str;
    fn session_id(&self) -> &str;
    fn session_file(&self) -> &Path;
    fn write_event(&mut self, event: &TranscriptEvent) -> io::Result<()>;
    fn finish(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionInfo {
    pub format: &'static str,
    pub session_id: String,
    pub session_file: PathBuf,
}

pub struct SessionLogger {
    pub sessions: Vec<SessionInfo>,
    pub agent_name: String,
    tx: UnboundedSender<TranscriptEvent>,
    join: JoinHandle<()>,
}

impl SessionLogger {
    pub fn start(
        harnesses: Vec<Box<dyn SessionHarness>>,
        agent_name: impl Into<String>,
    ) -> io::Result<Self> {
        if harnesses.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "at least one session harness is required",
            ));
        }

        let sessions = harnesses
            .iter()
            .map(|harness| SessionInfo {
                format: harness.format_name(),
                session_id: harness.session_id().to_string(),
                session_file: harness.session_file().to_path_buf(),
            })
            .collect();
        let (tx, rx) = mpsc::unbounded_channel();
        let join = tokio::spawn(async move { run_harnesses(harnesses, rx).await });

        Ok(Self {
            sessions,
            agent_name: agent_name.into(),
            tx,
            join,
        })
    }

    /// Convenience accessor for the first configured harness.
    pub fn session_id(&self) -> &str {
        &self.sessions[0].session_id
    }

    /// Convenience accessor for the first configured harness.
    pub fn session_file(&self) -> &Path {
        &self.sessions[0].session_file
    }

    pub fn first_line(&self) -> io::Result<String> {
        let first = self.sessions.first().ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidInput, "no session harness configured")
        })?;
        serde_json::to_string(&json!({
            "work_name": "session",
            "result": {
                "session_id": first.session_id,
                "session_file": first.session_file.to_string_lossy(),
                "harness": first.format,
                "agent_name": self.agent_name,
                "harnesses": self.sessions.iter().map(|session| json!({
                    "format": session.format,
                    "session_id": session.session_id,
                    "session_file": session.session_file.to_string_lossy(),
                })).collect::<Vec<_>>(),
            },
        }))
        .map_err(io::Error::other)
    }

    pub fn user_message(&self, text: impl Into<String>) {
        let _ = self.tx.send(TranscriptEvent::User { text: text.into() });
    }

    pub fn assistant_message(&self, text: impl Into<String>) {
        let _ = self
            .tx
            .send(TranscriptEvent::Assistant { text: text.into() });
    }

    pub fn monitor_event(&self, event: MonitorEvent) {
        let _ = self.tx.send(TranscriptEvent::MonitorEvent(event));
    }

    pub fn monitor_consumed(&self, event: MonitorEvent, reason: impl Into<String>) {
        let _ = self.tx.send(TranscriptEvent::MonitorConsumed {
            event,
            reason: reason.into(),
        });
    }

    pub fn transcript_sender(&self) -> UnboundedSender<TranscriptEvent> {
        self.tx.clone()
    }

    pub fn tool_observer(&self) -> std::sync::Arc<dyn ToolObserver> {
        std::sync::Arc::new(SessionToolObserver {
            tx: self.tx.clone(),
        })
    }

    pub async fn close(self) {
        let Self { tx, join, .. } = self;
        drop(tx);
        match tokio::time::timeout(Duration::from_secs(2), join).await {
            Ok(Ok(())) => {}
            Ok(Err(error)) => crate::log::emit(
                "error",
                "session.writer_failed",
                serde_json::json!({"error": error.to_string()}),
            ),
            Err(_) => crate::log::emit(
                "error",
                "session.flush_timeout",
                serde_json::json!({"error": "writer did not flush before shutdown"}),
            ),
        }
    }
}

async fn run_harnesses(
    mut harnesses: Vec<Box<dyn SessionHarness>>,
    mut rx: UnboundedReceiver<TranscriptEvent>,
) {
    while let Some(event) = rx.recv().await {
        for harness in &mut harnesses {
            if let Err(error) = harness.write_event(&event) {
                crate::log::emit(
                    "error",
                    "session.write_failed",
                    serde_json::json!({"harness": harness.format_name(), "error": error.to_string()}),
                );
            }
        }
    }
    for harness in &mut harnesses {
        if let Err(error) = harness.finish() {
            crate::log::emit(
                "error",
                "session.finish_failed",
                serde_json::json!({"harness": harness.format_name(), "error": error.to_string()}),
            );
        }
    }
}

struct SessionToolObserver {
    tx: UnboundedSender<TranscriptEvent>,
}

impl ToolObserver for SessionToolObserver {
    fn tool_call(&self, id: &str, name: &str, arguments: &Value) {
        let _ = self.tx.send(TranscriptEvent::ToolCall {
            id: id.to_string(),
            name: name.to_string(),
            arguments: arguments.clone(),
        });
    }

    fn tool_result(&self, id: &str, content: &str, is_error: bool, audit: Option<&Value>) {
        let _ = self.tx.send(TranscriptEvent::ToolResult {
            id: id.to_string(),
            content: content.to_string(),
            is_error,
            audit: audit.cloned(),
        });
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use super::*;

    struct MemoryHarness {
        format: &'static str,
        session_id: String,
        path: PathBuf,
        events: Arc<Mutex<Vec<TranscriptEvent>>>,
    }

    impl SessionHarness for MemoryHarness {
        fn format_name(&self) -> &'static str {
            self.format
        }

        fn session_id(&self) -> &str {
            &self.session_id
        }

        fn session_file(&self) -> &Path {
            &self.path
        }

        fn write_event(&mut self, event: &TranscriptEvent) -> io::Result<()> {
            self.events.lock().unwrap().push(event.clone());
            Ok(())
        }
    }

    fn memory_harness(
        format: &'static str,
        id: &str,
    ) -> (Box<dyn SessionHarness>, Arc<Mutex<Vec<TranscriptEvent>>>) {
        let events = Arc::new(Mutex::new(Vec::new()));
        (
            Box::new(MemoryHarness {
                format,
                session_id: id.to_string(),
                path: PathBuf::from(format!("/tmp/{id}.jsonl")),
                events: Arc::clone(&events),
            }),
            events,
        )
    }

    #[tokio::test]
    async fn fans_out_events_to_multiple_harnesses() {
        let (first, first_events) = memory_harness("first", "session-one");
        let (second, second_events) = memory_harness("second", "session-two");
        let logger = SessionLogger::start(vec![first, second], "alice").unwrap();

        let first_line: Value = serde_json::from_str(&logger.first_line().unwrap()).unwrap();
        assert_eq!(first_line["result"]["session_id"], "session-one");
        assert_eq!(
            first_line["result"]["harnesses"].as_array().unwrap().len(),
            2
        );

        logger.user_message("hello");
        logger.assistant_message("hi");
        logger.close().await;

        let expected = vec![
            TranscriptEvent::User {
                text: "hello".to_string(),
            },
            TranscriptEvent::Assistant {
                text: "hi".to_string(),
            },
        ];
        assert_eq!(*first_events.lock().unwrap(), expected);
        assert_eq!(*second_events.lock().unwrap(), expected);
    }
}