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

use cm_db::repo::{agents, messages, run_events, runs, sessions, steps, threads, 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(_)));
}

async fn mk_agent(pool: &sqlx::PgPool, ws: &Workspace, owner: &User, name: &str) -> Agent {
    let agent = Agent {
        id: AgentId::new(),
        workspace_id: ws.id,
        name: name.into(),
        job_title: "Worker".into(),
        system_prompt: "You help.".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();
    agent
}

#[tokio::test]
async fn room_is_visible_to_all_active_members() {
    let pool = cm_testkit::test_pool().await;
    let (ws, owner, a) = seeded(&pool).await;
    let b = mk_agent(&pool, &ws, &owner, "Bee").await;
    let c = mk_agent(&pool, &ws, &owner, "Cee").await;

    let room = threads::create_room(&pool, ws.id, "standup", Some(a.id), &[a.id, b.id, c.id])
        .await
        .unwrap();
    threads::add_message(&pool, room, a.id, json!({"text": "hi team"}), &[])
        .await
        .unwrap();

    for member in [a.id, b.id, c.id] {
        let listed = threads::list_for_agent(&pool, member).await.unwrap();
        assert!(
            listed.iter().any(|t| t.id == room && t.kind == "room"),
            "member should see the room",
        );
    }
    let mut members = threads::participants(&pool, room).await.unwrap();
    members.sort();
    let mut expected = vec![a.id.as_uuid(), b.id.as_uuid(), c.id.as_uuid()];
    expected.sort();
    assert_eq!(members, expected);
}

#[tokio::test]
async fn find_or_create_dm_never_matches_a_room() {
    let pool = cm_testkit::test_pool().await;
    let (ws, owner, a) = seeded(&pool).await;
    let b = mk_agent(&pool, &ws, &owner, "Bee").await;
    let c = mk_agent(&pool, &ws, &owner, "Cee").await;

    // A room containing both a and b must NOT satisfy the 1:1 lookup.
    let room = threads::create_room(&pool, ws.id, "group", Some(a.id), &[a.id, b.id, c.id])
        .await
        .unwrap();
    let dm = threads::find_or_create(&pool, ws.id, a.id, b.id, "dm")
        .await
        .unwrap();
    assert_ne!(dm, room, "DM must be a fresh 2-person thread, not the room");
    // Calling again returns the same DM (idempotent), still not the room.
    let dm2 = threads::find_or_create(&pool, ws.id, a.id, b.id, "dm")
        .await
        .unwrap();
    assert_eq!(dm, dm2);
}

#[tokio::test]
async fn leaving_a_room_hides_it_but_keeps_history() {
    let pool = cm_testkit::test_pool().await;
    let (ws, owner, a) = seeded(&pool).await;
    let b = mk_agent(&pool, &ws, &owner, "Bee").await;

    let room = threads::create_room(&pool, ws.id, "pair", Some(a.id), &[a.id, b.id])
        .await
        .unwrap();
    threads::add_message(&pool, room, b.id, json!({"text": "present"}), &[])
        .await
        .unwrap();

    threads::remove_participant(&pool, room, b.id).await.unwrap();
    assert!(
        !threads::list_for_agent(&pool, b.id)
            .await
            .unwrap()
            .iter()
            .any(|t| t.id == room),
        "left room should not appear in the inbox",
    );
    assert!(
        !threads::is_participant(&pool, room, b.id).await.unwrap(),
        "left member is no longer an active participant",
    );
    // History stays attributable.
    let msgs = threads::messages(&pool, room).await.unwrap();
    assert_eq!(msgs.len(), 1);
    assert_eq!(msgs[0].from_agent, b.id.as_uuid());
    // And a is still in.
    assert_eq!(threads::participants(&pool, room).await.unwrap(), vec![a.id.as_uuid()]);
}

#[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));
}