main.rs

use anyhow::{Context, Result};
use chrono::{SecondsFormat, Utc};
use native_host_rust::context_packet::{
    context_packet_paths, generate_source_context, prepare_agent_context_handoff,
    AgentContextHandoff, AgentContextHandoffInput,
};
use native_host_rust::protocol::{SourceContextResult, SummaryProviderSettings};
use serde::{Deserialize, Serialize};
use std::fs;
use std::io::{BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
use std::path::{Path, PathBuf};
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};
use std::thread;
use std::time::Duration;

const DEFAULT_AGENT_SERVER_PORT: u16 = 49731;
const AGENT_SERVER_PORT_ENV: &str = "MIRROR_NOTE_AGENT_SERVER_PORT";
const JOB_DIRECTORY_NAME: &str = "agent-jobs";
const GENERATE_CONTEXT_JOB_FILE_NAME: &str = "generate-context.json";

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct AgentRequest {
    id: String,
    command: AgentCommand,
    target_directory: Option<String>,
    meetings_root: Option<String>,
    task: Option<String>,
    project_hint: Option<String>,
    working_directory: Option<String>,
    summary_settings: Option<SummaryProviderSettings>,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
enum AgentCommand {
    Capabilities,
    GenerateContext,
    PrepareAgentContext,
    JobStatus,
    Shutdown,
}

#[derive(Debug, Serialize)]
struct AgentResponse {
    #[serde(rename = "type")]
    message_type: &'static str,
    id: String,
    ok: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    payload: Option<AgentPayload>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<String>,
}

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct AgentPayload {
    #[serde(skip_serializing_if = "Option::is_none")]
    agent_version: Option<&'static str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    job: Option<JobStatus>,
    #[serde(skip_serializing_if = "Option::is_none")]
    source_context: Option<SourceContextResult>,
    #[serde(skip_serializing_if = "Option::is_none")]
    handoff: Option<AgentContextHandoff>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct JobStatus {
    job_id: String,
    command: String,
    state: JobState,
    target_directory: String,
    queued_at: String,
    updated_at: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    started_at: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    completed_at: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    packet_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    object_count: Option<usize>,
    #[serde(skip_serializing_if = "Option::is_none")]
    relation_count: Option<usize>,
    #[serde(skip_serializing_if = "Option::is_none")]
    warning_count: Option<usize>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
enum JobState {
    Queued,
    Running,
    Succeeded,
    Failed,
}

fn main() -> Result<()> {
    let port = std::env::var(AGENT_SERVER_PORT_ENV)
        .ok()
        .and_then(|value| value.parse::<u16>().ok())
        .unwrap_or(DEFAULT_AGENT_SERVER_PORT);
    run_server(port)
}

fn run_server(port: u16) -> Result<()> {
    let listener = TcpListener::bind(("127.0.0.1", port))
        .with_context(|| format!("MirrorNoteAgentServer could not bind 127.0.0.1:{port}"))?;
    listener.set_nonblocking(true)?;
    let running = Arc::new(AtomicBool::new(true));

    while running.load(Ordering::SeqCst) {
        match listener.accept() {
            Ok((stream, _address)) => {
                stream.set_nonblocking(false)?;
                let running_for_client = Arc::clone(&running);
                thread::spawn(move || {
                    if let Err(error) = handle_client(stream, running_for_client) {
                        eprintln!("{error:#}");
                    }
                });
            }
            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                thread::sleep(Duration::from_millis(50));
            }
            Err(error) => return Err(error.into()),
        }
    }
    Ok(())
}

fn handle_client(mut stream: TcpStream, running: Arc<AtomicBool>) -> Result<()> {
    let peer_reader = stream.try_clone()?;
    let reader = BufReader::new(peer_reader);
    for line in reader.lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }
        let response = match serde_json::from_str::<AgentRequest>(&line) {
            Ok(request) => handle_request(request, &running),
            Err(error) => failure_response(
                "unknown".to_string(),
                format!("MirrorNoteAgentServer could not decode the request: {error}"),
            ),
        };
        write_response(&mut stream, response)?;
    }
    Ok(())
}

fn handle_request(request: AgentRequest, running: &Arc<AtomicBool>) -> AgentResponse {
    match request.command {
        AgentCommand::Capabilities => success_response(
            request.id,
            Some(AgentPayload {
                agent_version: Some(env!("CARGO_PKG_VERSION")),
                job: None,
                source_context: None,
                handoff: None,
            }),
        ),
        AgentCommand::GenerateContext => match generate_context_job(&request) {
            Ok((job, result)) => success_response(
                request.id,
                Some(AgentPayload {
                    agent_version: None,
                    job: Some(job),
                    source_context: Some(result),
                    handoff: None,
                }),
            ),
            Err(error) => failure_response(request.id, format!("{error:#}")),
        },
        AgentCommand::PrepareAgentContext => match prepare_agent_context(&request) {
            Ok(handoff) => success_response(
                request.id,
                Some(AgentPayload {
                    agent_version: None,
                    job: None,
                    source_context: None,
                    handoff: Some(handoff),
                }),
            ),
            Err(error) => failure_response(request.id, format!("{error:#}")),
        },
        AgentCommand::JobStatus => match read_job_status_for_request(&request) {
            Ok(job) => success_response(
                request.id,
                Some(AgentPayload {
                    agent_version: None,
                    job,
                    source_context: None,
                    handoff: None,
                }),
            ),
            Err(error) => failure_response(request.id, format!("{error:#}")),
        },
        AgentCommand::Shutdown => {
            running.store(false, Ordering::SeqCst);
            success_response(
                request.id,
                Some(AgentPayload {
                    agent_version: Some(env!("CARGO_PKG_VERSION")),
                    job: None,
                    source_context: None,
                    handoff: None,
                }),
            )
        }
    }
}

fn prepare_agent_context(request: &AgentRequest) -> Result<AgentContextHandoff> {
    let meetings_root = request
        .meetings_root
        .as_deref()
        .context("A meetings root directory is required.")?;
    let task = request
        .task
        .as_deref()
        .filter(|task| !task.trim().is_empty())
        .context("A task is required to prepare agent context.")?;
    let result = prepare_agent_context_handoff(
        Path::new(meetings_root),
        AgentContextHandoffInput {
            task: task.to_string(),
            project_hint: request.project_hint.clone(),
            working_directory: request.working_directory.clone(),
        },
    )?;
    Ok(result.handoff)
}

fn generate_context_job(request: &AgentRequest) -> Result<(JobStatus, SourceContextResult)> {
    let target_directory = request
        .target_directory
        .as_deref()
        .context("A target bundle directory is required.")?;
    let summary_settings = request
        .summary_settings
        .as_ref()
        .context("Summary provider settings are required.")?;
    let bundle_directory = Path::new(target_directory);
    let now = current_iso8601();
    let mut job = JobStatus {
        job_id: format!("generate-context-{}", Utc::now().timestamp_millis()),
        command: "generateContext".to_string(),
        state: JobState::Queued,
        target_directory: target_directory.to_string(),
        queued_at: now.clone(),
        updated_at: now,
        started_at: None,
        completed_at: None,
        packet_id: None,
        object_count: None,
        relation_count: None,
        warning_count: None,
        error: None,
    };
    write_job_status(bundle_directory, &job)?;

    job.state = JobState::Running;
    job.started_at = Some(current_iso8601());
    job.updated_at = current_iso8601();
    write_job_status(bundle_directory, &job)?;

    match generate_source_context(bundle_directory, summary_settings) {
        Ok(result) => {
            job.state = JobState::Succeeded;
            job.updated_at = current_iso8601();
            job.completed_at = Some(job.updated_at.clone());
            job.packet_id = Some(result.packet_id.clone());
            job.object_count = Some(result.object_count);
            job.relation_count = Some(result.relation_count);
            job.warning_count = Some(result.warning_count);
            write_job_status(bundle_directory, &job)?;
            Ok((job, result))
        }
        Err(error) => {
            job.state = JobState::Failed;
            job.updated_at = current_iso8601();
            job.completed_at = Some(job.updated_at.clone());
            job.error = Some(format!("{error:#}"));
            write_job_status(bundle_directory, &job)?;
            Err(error)
        }
    }
}

fn read_job_status_for_request(request: &AgentRequest) -> Result<Option<JobStatus>> {
    let target_directory = request
        .target_directory
        .as_deref()
        .context("A target bundle directory is required.")?;
    read_job_status(Path::new(target_directory))
}

fn job_status_path(bundle_directory: &Path) -> PathBuf {
    context_packet_paths(bundle_directory)
        .context_directory
        .join(JOB_DIRECTORY_NAME)
        .join(GENERATE_CONTEXT_JOB_FILE_NAME)
}

fn write_job_status(bundle_directory: &Path, job: &JobStatus) -> Result<()> {
    let path = job_status_path(bundle_directory);
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    let bytes = serde_json::to_vec_pretty(job)?;
    fs::write(path, bytes)?;
    Ok(())
}

fn read_job_status(bundle_directory: &Path) -> Result<Option<JobStatus>> {
    let path = job_status_path(bundle_directory);
    if !path.exists() {
        return Ok(None);
    }
    let bytes = fs::read(path)?;
    let status = serde_json::from_slice(&bytes)?;
    Ok(Some(status))
}

fn write_response(out: &mut impl Write, response: AgentResponse) -> Result<()> {
    serde_json::to_writer(&mut *out, &response)?;
    out.write_all(b"\n")?;
    out.flush()?;
    Ok(())
}

fn success_response(id: String, payload: Option<AgentPayload>) -> AgentResponse {
    AgentResponse {
        message_type: "response",
        id,
        ok: true,
        payload,
        error: None,
    }
}

fn failure_response(id: String, error: String) -> AgentResponse {
    AgentResponse {
        message_type: "response",
        id,
        ok: false,
        payload: None,
        error: Some(error),
    }
}

fn current_iso8601() -> String {
    Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true)
}

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

    #[test]
    fn persists_generate_context_job_status_under_bundle_context_directory() {
        let temp = tempfile::tempdir().unwrap();
        let bundle = temp.path();
        let job = JobStatus {
            job_id: "job-1".to_string(),
            command: "generateContext".to_string(),
            state: JobState::Running,
            target_directory: bundle.to_string_lossy().into_owned(),
            queued_at: "2026-05-04T00:00:00.000Z".to_string(),
            updated_at: "2026-05-04T00:00:01.000Z".to_string(),
            started_at: Some("2026-05-04T00:00:01.000Z".to_string()),
            completed_at: None,
            packet_id: None,
            object_count: None,
            relation_count: None,
            warning_count: None,
            error: None,
        };

        write_job_status(bundle, &job).unwrap();

        let path = job_status_path(bundle);
        assert_eq!(
            path,
            bundle
                .join("context")
                .join("agent-jobs")
                .join("generate-context.json")
        );
        let restored = read_job_status(bundle).unwrap().unwrap();
        assert_eq!(restored.job_id, "job-1");
        assert!(matches!(restored.state, JobState::Running));
        assert_eq!(
            restored.started_at.as_deref(),
            Some("2026-05-04T00:00:01.000Z")
        );
    }

    #[test]
    fn missing_job_status_restores_to_none() {
        let temp = tempfile::tempdir().unwrap();
        assert!(read_job_status(temp.path()).unwrap().is_none());
    }

    #[test]
    fn prepare_agent_context_request_writes_handoff_files() {
        let temp = tempfile::tempdir().unwrap();
        let meetings_root = temp.path();
        let project = meetings_root
            .join("projects")
            .join("mirrornote-agent-runtime");
        fs::create_dir_all(&project).unwrap();
        fs::write(
            project.join("source-index.jsonl"),
            "{\"id\":\"meeting:note-1\",\"type\":\"meeting\",\"title\":\"Agent Runtime Sync\",\"observedAt\":\"2026-05-05T01:00:00Z\",\"projectHint\":\"MirrorNote Agent Runtime\"}\n",
        )
        .unwrap();
        fs::write(
            project.join("memory-objects.jsonl"),
            "{\"id\":\"meeting:note-1::decision-1\",\"type\":\"decision\",\"title\":\"Keep context automatic\",\"body\":\"Keep context automatic\",\"status\":\"active\",\"confidence\":0.95,\"evidenceCoverage\":\"direct\",\"sourceRefs\":[{\"sourceId\":\"meeting:note-1\",\"sourceType\":\"meeting\",\"location\":{\"kind\":\"metadata\"}}]}\n",
        )
        .unwrap();

        let response = handle_request(
            AgentRequest {
                id: "request-1".to_string(),
                command: AgentCommand::PrepareAgentContext,
                target_directory: None,
                meetings_root: Some(meetings_root.to_string_lossy().into_owned()),
                task: Some("Run Codex with prior context".to_string()),
                project_hint: Some("MirrorNote Agent Runtime".to_string()),
                working_directory: Some("/tmp/MirrorNote".to_string()),
                summary_settings: None,
            },
            &Arc::new(AtomicBool::new(true)),
        );

        assert!(response.ok);
        let payload = response.payload.unwrap();
        let handoff = payload.handoff.unwrap();
        assert_eq!(handoff.project_id, "mirrornote-agent-runtime");
        assert_eq!(handoff.decisions.len(), 1);
        assert!(project.join("handoffs").join("handoff.json").exists());
        assert!(project.join("handoffs").join("handoff.md").exists());
    }
}