clawmates/crates/cm-db/tests/chat_repos.rs

use cm_db::repo::{agents, messages, run_events, runs, sessions, steps, users, workspaces};
use cm_db::DbError;
use cm_domain::{
    AccessPolicy, Agent, AgentId, AgentStatus, MessageRole, Role, RunState, Session, SessionId,
    Step, StepStatus, User, UserId, Workspace, WorkspaceId,
};
use serde_json::json;

async fn seeded(pool: &sqlx::PgPool) -> (Workspace, User, Agent) {
    let ws = Workspace {
        id: WorkspaceId::new(),
        name: "Acme".into(),
        plan: "team".into(),
    };
    workspaces::insert(pool, &ws).await.unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: ws.id,
        email: format!("{}@acme.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    users::insert(pool, &owner).await.unwrap();
    let agent = Agent {
        id: AgentId::new(),
        workspace_id: ws.id,
        name: "Scout".into(),
        job_title: "Analyst".into(),
        system_prompt: "You research.".into(),
        avatar: String::new(),
        accent: String::new(),
        wallpaper: String::new(),
        managed_by: owner.id,
        status: AgentStatus::Online,
    };
    agents::insert(pool, &agent, &AccessPolicy::default())
        .await
        .unwrap();
    (ws, owner, agent)
}

async fn new_session(pool: &sqlx::PgPool, agent: &Agent, title: &str) -> Session {
    let session = sessions::create(pool, agent.id, agent.workspace_id, title)
        .await
        .unwrap();
    assert_eq!(session.title, title);
    session
}

#[tokio::test]
async fn session_create_get_and_shard_are_consistent() {
    let pool = cm_testkit::test_pool().await;
    let (_, _, agent) = seeded(&pool).await;
    let session = new_session(&pool, &agent, "Quarterly research").await;

    let fetched = sessions::get(&pool, session.id).await.unwrap();
    assert_eq!(fetched, session);
    assert_eq!(fetched.shard, cm_domain::shard_of(agent.id));
}

#[tokio::test]
async fn sessions_list_most_recently_active_first() {
    let pool = cm_testkit::test_pool().await;
    let (_, _, agent) = seeded(&pool).await;
    let older = new_session(&pool, &agent, "First").await;
    let newer = new_session(&pool, &agent, "Second").await;

    let listed = sessions::list_by_agent(&pool, agent.id).await.unwrap();
    assert_eq!(
        listed.iter().map(|s| s.id).collect::<Vec<_>>(),
        vec![newer.id, older.id]
    );

    // Touching the older session moves it back to the top.
    sessions::touch(&pool, older.id).await.unwrap();
    let relisted = sessions::list_by_agent(&pool, agent.id).await.unwrap();
    assert_eq!(relisted[0].id, older.id);
}

#[tokio::test]
async fn messages_get_increasing_seq_starting_at_one() {
    let pool = cm_testkit::test_pool().await;
    let (_, _, agent) = seeded(&pool).await;
    let session = new_session(&pool, &agent, "Chat").await;

    let first = messages::append(&pool, session.id, MessageRole::User, json!({"text": "hi"}))
        .await
        .unwrap();
    let second = messages::append(
        &pool,
        session.id,
        MessageRole::Agent,
        json!({"text": "hello!"}),
    )
    .await
    .unwrap();
    assert_eq!(first.seq, 1);
    assert_eq!(second.seq, 2);
    assert_eq!(second.role, MessageRole::Agent);
}

#[tokio::test]
async fn history_returns_messages_with_ordered_steps() {
    let pool = cm_testkit::test_pool().await;
    let (_, _, agent) = seeded(&pool).await;
    let session = new_session(&pool, &agent, "Chat").await;

    messages::append(
        &pool,
        session.id,
        MessageRole::User,
        json!({"text": "time?"}),
    )
    .await
    .unwrap();
    let reply = messages::append(
        &pool,
        session.id,
        MessageRole::Agent,
        json!({"text": "It is noon."}),
    )
    .await
    .unwrap();
    for (seq, tool) in [(1, "clock.now"), (2, "clock.now")] {
        steps::append(
            &pool,
            &Step {
                id: uuid::Uuid::now_v7(),
                message_id: reply.id,
                seq,
                kind: "tool_call".into(),
                tool_name: Some(tool.into()),
                input: Some(json!({})),
                output: Some(json!({"now": "12:00"})),
                taint: vec![],
                status: StepStatus::Ok,
            },
        )
        .await
        .unwrap();
    }

    let history = messages::history(&pool, session.id).await.unwrap();
    assert_eq!(history.len(), 2);
    assert!(history[0].steps.is_empty());
    let trace = &history[1].steps;
    assert_eq!(trace.len(), 2);
    assert_eq!(trace[0].seq, 1);
    assert_eq!(trace[1].seq, 2);
    assert_eq!(trace[0].tool_name.as_deref(), Some("clock.now"));
    assert_eq!(trace[0].status, StepStatus::Ok);
}

#[tokio::test]
async fn runs_track_state_transitions() {
    let pool = cm_testkit::test_pool().await;
    let (_, _, agent) = seeded(&pool).await;
    let session = new_session(&pool, &agent, "Chat").await;

    let run_id = runs::create(&pool, session.id).await.unwrap();
    let run = runs::get(&pool, run_id).await.unwrap();
    assert_eq!(run.state, RunState::Running);

    runs::set_state(&pool, run_id, RunState::Completed, None)
        .await
        .unwrap();
    let done = runs::get(&pool, run_id).await.unwrap();
    assert_eq!(done.state, RunState::Completed);

    runs::set_state(&pool, run_id, RunState::Failed, Some("llm unreachable"))
        .await
        .unwrap();
    let failed = runs::get(&pool, run_id).await.unwrap();
    assert_eq!(failed.error.as_deref(), Some("llm unreachable"));
}

#[tokio::test]
async fn run_events_replay_from_an_offset() {
    let pool = cm_testkit::test_pool().await;
    let (_, _, agent) = seeded(&pool).await;
    let session = new_session(&pool, &agent, "Chat").await;
    let run_id = runs::create(&pool, session.id).await.unwrap();

    for (seq, kind) in [(1, "run_started"), (2, "text_delta"), (3, "run_completed")] {
        run_events::append(&pool, run_id, seq, kind, json!({"seq": seq}))
            .await
            .unwrap();
    }

    let all = run_events::list_after(&pool, run_id, 0).await.unwrap();
    assert_eq!(all.len(), 3);
    assert_eq!(all[0].event_type, "run_started");

    let tail = run_events::list_after(&pool, run_id, 2).await.unwrap();
    assert_eq!(tail.len(), 1);
    assert_eq!(tail[0].event_type, "run_completed");
    assert_eq!(tail[0].seq, 3);
}

#[tokio::test]
async fn duplicate_event_seq_is_a_conflict() {
    let pool = cm_testkit::test_pool().await;
    let (_, _, agent) = seeded(&pool).await;
    let session = new_session(&pool, &agent, "Chat").await;
    let run_id = runs::create(&pool, session.id).await.unwrap();

    run_events::append(&pool, run_id, 1, "run_started", json!({}))
        .await
        .unwrap();
    let err = run_events::append(&pool, run_id, 1, "run_started", json!({}))
        .await
        .unwrap_err();
    assert!(matches!(err, DbError::Conflict(_)));
}

#[tokio::test]
async fn missing_session_is_not_found() {
    let pool = cm_testkit::test_pool().await;
    let err = sessions::get(&pool, SessionId::new()).await.unwrap_err();
    assert!(matches!(err, DbError::NotFound));
}